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.
 
 
 

116 lines
3.0 KiB

  1. package helpers
  2. import (
  3. "fmt"
  4. "net/http"
  5. "time"
  6. "github.com/QuantumNous/new-api/common"
  7. )
  8. type optionResponse struct {
  9. Success bool `json:"success"`
  10. Message string `json:"message,omitempty"`
  11. }
  12. type internalListUsersResponse struct {
  13. Success bool `json:"success"`
  14. Total int `json:"total"`
  15. }
  16. func SetOption(client *APIClient, key string, value any) error {
  17. var resp optionResponse
  18. if err := client.PutJSON("/api/option/", map[string]any{
  19. "key": key,
  20. "value": value,
  21. }, http.StatusOK, &resp); err != nil {
  22. return err
  23. }
  24. if !resp.Success {
  25. return fmt.Errorf("set option %s failed: %s", key, resp.Message)
  26. }
  27. return nil
  28. }
  29. func ConfigureCNRegionSync(client *APIClient, cnBaseURL, ovBaseURL, apiKey string) error {
  30. slaveEndpoints, err := common.Marshal([]string{ovBaseURL})
  31. if err != nil {
  32. return fmt.Errorf("marshal slave endpoints: %w", err)
  33. }
  34. options := []struct {
  35. Key string
  36. Value any
  37. }{
  38. {Key: "region_sync.enabled", Value: true},
  39. {Key: "region_sync.region_id", Value: "cn"},
  40. {Key: "region_sync.is_master", Value: true},
  41. {Key: "region_sync.master_endpoint", Value: cnBaseURL},
  42. {Key: "region_sync.slave_endpoints", Value: string(slaveEndpoints)},
  43. {Key: "region_sync.sync_api_key", Value: apiKey},
  44. }
  45. for _, option := range options {
  46. if err := SetOption(client, option.Key, option.Value); err != nil {
  47. return err
  48. }
  49. }
  50. return nil
  51. }
  52. func ConfigureOVRegionSync(client *APIClient, cnBaseURL, apiKey string) error {
  53. emptyEndpoints, err := common.Marshal([]string{})
  54. if err != nil {
  55. return fmt.Errorf("marshal empty slave endpoints: %w", err)
  56. }
  57. options := []struct {
  58. Key string
  59. Value any
  60. }{
  61. {Key: "region_sync.enabled", Value: true},
  62. {Key: "region_sync.region_id", Value: "ov"},
  63. {Key: "region_sync.is_master", Value: false},
  64. {Key: "region_sync.master_endpoint", Value: cnBaseURL},
  65. {Key: "region_sync.slave_endpoints", Value: string(emptyEndpoints)},
  66. {Key: "region_sync.sync_api_key", Value: apiKey},
  67. }
  68. for _, option := range options {
  69. if err := SetOption(client, option.Key, option.Value); err != nil {
  70. return err
  71. }
  72. }
  73. return nil
  74. }
  75. func WaitOptionApplied(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error {
  76. return WaitInternalMigrationReady(baseURL, apiKey, expectedTotal, timeout)
  77. }
  78. func WaitInternalMigrationReady(baseURL, apiKey string, expectedTotal int, timeout time.Duration) error {
  79. client, err := NewAPIClient(baseURL)
  80. if err != nil {
  81. return err
  82. }
  83. deadline := time.Now().Add(timeout)
  84. for time.Now().Before(deadline) {
  85. var resp internalListUsersResponse
  86. err := client.DoJSONWithHeaders(
  87. http.MethodGet,
  88. "/api/internal/migration/users?page=0&page_size=1",
  89. nil,
  90. http.StatusOK,
  91. map[string]string{
  92. "X-Sync-API-Key": apiKey,
  93. "X-Sync-Node": "e2e-ready-check",
  94. },
  95. &resp,
  96. )
  97. if err == nil && resp.Success && resp.Total == expectedTotal {
  98. return nil
  99. }
  100. time.Sleep(250 * time.Millisecond)
  101. }
  102. return fmt.Errorf("internal migration API at %s not ready within %s", baseURL, timeout)
  103. }