|
- package user_migration
-
- import (
- "bytes"
- "fmt"
- "io"
- "net/http"
- "net/url"
- "strings"
- "time"
-
- "github.com/QuantumNous/new-api/common"
- )
-
- type migrationClient interface {
- ListUsers(page, pageSize int, keyword string) (*ListRemoteUsersResponse, error)
- QueryUsers(page, pageSize int, sourceUserIDs []int) (*QueryRemoteUsersResponse, error)
- GetUser(userId int) (*RemoteUserSnapshot, error)
- ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error
- HasSyncedCopy(cnUserId int) (bool, error)
- }
-
- type Client struct {
- endpoint string
- apiKey string
- http *http.Client
- }
-
- func NewClient(endpoint, apiKey string) *Client {
- return &Client{
- endpoint: strings.TrimRight(endpoint, "/"),
- apiKey: apiKey,
- http: &http.Client{Timeout: 30 * time.Second},
- }
- }
-
- func (c *Client) doRequest(method, path string, body any) ([]byte, error) {
- var bodyReader io.Reader
- if body != nil {
- data, err := common.Marshal(body)
- if err != nil {
- return nil, fmt.Errorf("marshal request: %w", err)
- }
- bodyReader = bytes.NewReader(data)
- }
-
- req, err := http.NewRequest(method, c.endpoint+path, bodyReader)
- if err != nil {
- return nil, err
- }
- req.Header.Set("X-Sync-API-Key", c.apiKey)
- req.Header.Set("X-Sync-Node", "cn-migration")
- if body != nil {
- req.Header.Set("Content-Type", "application/json")
- }
-
- resp, err := c.http.Do(req)
- if err != nil {
- return nil, err
- }
- defer resp.Body.Close()
-
- data, err := io.ReadAll(resp.Body)
- if err != nil {
- return nil, err
- }
- if resp.StatusCode != http.StatusOK {
- return nil, fmt.Errorf("status %d: %s", resp.StatusCode, string(data))
- }
- return data, nil
- }
-
- func (c *Client) ListUsers(page, pageSize int, keyword string) (*ListRemoteUsersResponse, error) {
- path := fmt.Sprintf("/api/internal/migration/users?page=%d&page_size=%d", page, pageSize)
- if keyword != "" {
- path += "&keyword=" + url.QueryEscape(keyword)
- }
- data, err := c.doRequest(http.MethodGet, path, nil)
- if err != nil {
- return nil, err
- }
- var resp ListRemoteUsersResponse
- if err := common.Unmarshal(data, &resp); err != nil {
- return nil, err
- }
- return &resp, nil
- }
-
- func (c *Client) QueryUsers(page, pageSize int, sourceUserIDs []int) (*QueryRemoteUsersResponse, error) {
- data, err := c.doRequest(http.MethodPost, "/api/internal/migration/users/query", &QueryRemoteUsersRequest{
- Page: page,
- PageSize: pageSize,
- SelectionMode: SelectionModeExplicitIDs,
- SourceUserIDs: sourceUserIDs,
- })
- if err != nil {
- return nil, err
- }
- var resp QueryRemoteUsersResponse
- if err := common.Unmarshal(data, &resp); err != nil {
- return nil, err
- }
- return &resp, nil
- }
-
- func (c *Client) GetUser(userId int) (*RemoteUserSnapshot, error) {
- data, err := c.doRequest(http.MethodGet,
- fmt.Sprintf("/api/internal/migration/users/%d", userId), nil)
- if err != nil {
- return nil, err
- }
- var resp struct {
- Success bool `json:"success"`
- Data RemoteUserSnapshot `json:"data"`
- Error string `json:"error,omitempty"`
- }
- if err := common.Unmarshal(data, &resp); err != nil {
- return nil, err
- }
- if !resp.Success {
- return nil, fmt.Errorf("get user failed: %s", resp.Error)
- }
- return &resp.Data, nil
- }
-
- func (c *Client) ConvertToSynced(userId int, req *ConvertRemoteUserRequest) error {
- data, err := c.doRequest(http.MethodPost,
- fmt.Sprintf("/api/internal/migration/users/%d/convert-to-synced", userId), req)
- if err != nil {
- return err
- }
- var resp struct {
- Success bool `json:"success"`
- Error string `json:"error,omitempty"`
- }
- if err := common.Unmarshal(data, &resp); err != nil {
- return err
- }
- if !resp.Success {
- return fmt.Errorf("convert-to-synced failed: %s", resp.Error)
- }
- return nil
- }
-
- func (c *Client) HasSyncedCopy(cnUserId int) (bool, error) {
- data, err := c.doRequest(http.MethodGet,
- fmt.Sprintf("/api/internal/migration/users/check-synced-copy?cn_user_id=%d", cnUserId), nil)
- if err != nil {
- return false, err
- }
- var resp struct {
- Success bool `json:"success"`
- HasCopy bool `json:"has_copy"`
- Error string `json:"error,omitempty"`
- }
- if err := common.Unmarshal(data, &resp); err != nil {
- return false, err
- }
- if !resp.Success {
- return false, fmt.Errorf("check-synced-copy failed: %s", resp.Error)
- }
- return resp.HasCopy, nil
- }
|