|
- package helpers
-
- import (
- "fmt"
- "net/http"
- "time"
-
- "github.com/QuantumNous/new-api/common"
- )
-
- type optionResponse struct {
- Success bool `json:"success"`
- Message string `json:"message,omitempty"`
- }
-
- type internalListUsersResponse struct {
- Success bool `json:"success"`
- Total int `json:"total"`
- }
-
- func SetOption(client *APIClient, key string, value any) error {
- var resp optionResponse
- if err := client.PutJSON("/api/option/", map[string]any{
- "key": key,
- "value": value,
- }, http.StatusOK, &resp); err != nil {
- return err
- }
- if !resp.Success {
- return fmt.Errorf("set option %s failed: %s", key, resp.Message)
- }
- return nil
- }
-
- func ConfigureCNRegionSync(client *APIClient, cnBaseURL, ovBaseURL, apiKey string) error {
- slaveEndpoints, err := common.Marshal([]string{ovBaseURL})
- if err != nil {
- return fmt.Errorf("marshal slave endpoints: %w", err)
- }
-
- options := []struct {
- Key string
- Value any
- }{
- {Key: "region_sync.enabled", Value: true},
- {Key: "region_sync.region_id", Value: "cn"},
- {Key: "region_sync.is_master", Value: true},
- {Key: "region_sync.master_endpoint", Value: cnBaseURL},
- {Key: "region_sync.slave_endpoints", Value: string(slaveEndpoints)},
- {Key: "region_sync.sync_api_key", Value: apiKey},
- }
- for _, option := range options {
- if err := SetOption(client, option.Key, option.Value); err != nil {
- return err
- }
- }
- return nil
- }
-
- func ConfigureOVRegionSync(client *APIClient, cnBaseURL, apiKey string) error {
- emptyEndpoints, err := common.Marshal([]string{})
- if err != nil {
- return fmt.Errorf("marshal empty slave endpoints: %w", err)
- }
-
- options := []struct {
- Key string
- Value any
- }{
- {Key: "region_sync.enabled", Value: true},
- {Key: "region_sync.region_id", Value: "ov"},
- {Key: "region_sync.is_master", Value: false},
- {Key: "region_sync.master_endpoint", Value: cnBaseURL},
- {Key: "region_sync.slave_endpoints", Value: string(emptyEndpoints)},
- {Key: "region_sync.sync_api_key", Value: apiKey},
- }
- for _, option := range options {
- if err := SetOption(client, option.Key, option.Value); err != nil {
- return err
- }
- }
- return nil
- }
-
- func WaitOptionApplied(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error {
- return WaitInternalMigrationReady(baseURL, apiKey, expectedTotal, timeout)
- }
-
- func WaitInternalMigrationReady(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error {
- client, err := NewAPIClient(baseURL)
- if err != nil {
- return err
- }
-
- deadline := time.Now().Add(timeout)
- for time.Now().Before(deadline) {
- var resp internalListUsersResponse
- err := client.DoJSONWithHeaders(
- http.MethodGet,
- "/api/internal/migration/users?page=0&page_size=1",
- nil,
- http.StatusOK,
- map[string]string{
- "X-Sync-API-Key": apiKey,
- "X-Sync-Node": "e2e-ready-check",
- },
- &resp,
- )
- if err == nil && resp.Success && resp.Total == expectedTotal {
- return nil
- }
- time.Sleep(250 * time.Millisecond)
- }
- return fmt.Errorf("internal migration API at %s not ready within %s", baseURL, timeout)
- }
|