You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

142 lines
3.5 KiB

  1. package user_migration
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "net/http"
  7. "strings"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. )
  11. type migrationClient interface {
  12. ListUsers(page, pageSize int) (*ListRemoteUsersResponse, error)
  13. GetUser(userId int) (*RemoteUserSnapshot, error)
  14. ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error
  15. HasSyncedCopy(cnUserId int) (bool, error)
  16. }
  17. type Client struct {
  18. endpoint string
  19. apiKey string
  20. http *http.Client
  21. }
  22. func NewClient(endpoint, apiKey string) *Client {
  23. return &Client{
  24. endpoint: strings.TrimRight(endpoint, "/"),
  25. apiKey: apiKey,
  26. http: &http.Client{Timeout: 30 * time.Second},
  27. }
  28. }
  29. func (c *Client) doRequest(method, path string, body any) ([]byte, error) {
  30. var bodyReader io.Reader
  31. if body != nil {
  32. data, err := common.Marshal(body)
  33. if err != nil {
  34. return nil, fmt.Errorf("marshal request: %w", err)
  35. }
  36. bodyReader = bytes.NewReader(data)
  37. }
  38. req, err := http.NewRequest(method, c.endpoint+path, bodyReader)
  39. if err != nil {
  40. return nil, err
  41. }
  42. req.Header.Set("X-Sync-API-Key", c.apiKey)
  43. req.Header.Set("X-Sync-Node", "cn-migration")
  44. if body != nil {
  45. req.Header.Set("Content-Type", "application/json")
  46. }
  47. resp, err := c.http.Do(req)
  48. if err != nil {
  49. return nil, err
  50. }
  51. defer resp.Body.Close()
  52. data, err := io.ReadAll(resp.Body)
  53. if err != nil {
  54. return nil, err
  55. }
  56. if resp.StatusCode != http.StatusOK {
  57. return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(data))
  58. }
  59. return data, nil
  60. }
  61. func (c *Client) ListUsers(page, pageSize int) (*ListRemoteUsersResponse, error) {
  62. data, err := c.doRequest(http.MethodGet,
  63. fmt.Sprintf("/api/internal/migration/users?page=%d&page_size=%d", page, pageSize), nil)
  64. if err != nil {
  65. return nil, err
  66. }
  67. var resp ListRemoteUsersResponse
  68. if err := common.Unmarshal(data, &resp); err != nil {
  69. return nil, err
  70. }
  71. return &resp, nil
  72. }
  73. func (c *Client) GetUser(userId int) (*RemoteUserSnapshot, error) {
  74. data, err := c.doRequest(http.MethodGet,
  75. fmt.Sprintf("/api/internal/migration/users/%d", userId), nil)
  76. if err != nil {
  77. return nil, err
  78. }
  79. var resp struct {
  80. Success bool `json:"success"`
  81. Data RemoteUserSnapshot `json:"data"`
  82. Error string `json:"error,omitempty"`
  83. }
  84. if err := common.Unmarshal(data, &resp); err != nil {
  85. return nil, err
  86. }
  87. if !resp.Success {
  88. return nil, fmt.Errorf("get user failed: %s", resp.Error)
  89. }
  90. return &resp.Data, nil
  91. }
  92. func (c *Client) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error {
  93. data, err := c.doRequest(http.MethodPost,
  94. fmt.Sprintf("/api/internal/migration/users/%d/convert-to-synced", userId), req)
  95. if err != nil {
  96. return err
  97. }
  98. var resp struct {
  99. Success bool `json:"success"`
  100. Error string `json:"error,omitempty"`
  101. }
  102. if err := common.Unmarshal(data, &resp); err != nil {
  103. return err
  104. }
  105. if !resp.Success {
  106. return fmt.Errorf("convert-to-synced failed: %s", resp.Error)
  107. }
  108. return nil
  109. }
  110. func (c *Client) HasSyncedCopy(cnUserId int) (bool, error) {
  111. data, err := c.doRequest(http.MethodGet,
  112. fmt.Sprintf("/api/internal/migration/users/check-synced-copy?cn_user_id=%d", cnUserId), nil)
  113. if err != nil {
  114. return false, err
  115. }
  116. var resp struct {
  117. Success bool `json:"success"`
  118. HasCopy bool `json:"has_copy"`
  119. Error string `json:"error,omitempty"`
  120. }
  121. if err := common.Unmarshal(data, &resp); err != nil {
  122. return false, err
  123. }
  124. if !resp.Success {
  125. return false, fmt.Errorf("check-synced-copy failed: %s", resp.Error)
  126. }
  127. return resp.HasCopy, nil
  128. }