25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

164 lines
4.2 KiB

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