package lake import ( "context" "encoding/json" "errors" "fmt" "io" "net/http" "net/http/httptest" "reflect" "strconv" "strings" "sync" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) var testNow = time.Date(2026, time.August, 2, 5, 6, 7, 0, time.UTC) var testInvocation = Invocation{ TiDBOrgID: "1000000000000000001", LakeOrgID: "7", OrgSlug: "test-org", ContextType: "data_pipeline", ContextID: "dp-test-001", ContextVersion: 3, InitiatedBy: "test-user-001", } type contractServer struct { t *testing.T secret string invocation Invocation resolvedOrgSlug string mu sync.Mutex exchanges int businessCalls int latestToken string lastExchangeBody []byte nonces map[string]struct{} business func(http.ResponseWriter, *http.Request, []byte) } func newContractServer( t *testing.T, business func(http.ResponseWriter, *http.Request, []byte), ) (*httptest.Server, *contractServer) { t.Helper() state := &contractServer{ t: t, secret: "test-hmac-secret", invocation: testInvocation, resolvedOrgSlug: testInvocation.OrgSlug, nonces: make(map[string]struct{}), business: business, } server := httptest.NewServer(http.HandlerFunc(state.serveHTTP)) t.Cleanup(server.Close) return server, state } func (s *contractServer) serveHTTP(response http.ResponseWriter, request *http.Request) { body, err := io.ReadAll(request.Body) if err != nil { s.t.Errorf("read request body: %v", err) response.WriteHeader(http.StatusInternalServerError) return } response.Header().Set("Content-Type", "application/json") if request.URL.Path == tokenExchangePath { s.serveTokenExchange(response, request, body) return } s.serveBusiness(response, request, body) } func (s *contractServer) serveTokenExchange(response http.ResponseWriter, request *http.Request, body []byte) { s.mu.Lock() defer s.mu.Unlock() assert.Equal(s.t, http.MethodPost, request.Method) assert.Equal(s.t, strconv.FormatInt(testNow.Unix(), 10), request.Header.Get("X-Trusted-Timestamp")) nonce := request.Header.Get("X-Trusted-Nonce") assert.NotEmpty(s.t, nonce) if _, exists := s.nonces[nonce]; exists { s.t.Errorf("nonce %q was reused", nonce) } s.nonces[nonce] = struct{}{} expectedSignature := sign(s.secret, http.MethodPost, tokenExchangePath, testNow.Unix(), nonce, body) assert.Equal(s.t, expectedSignature, request.Header.Get("X-Trusted-Signature")) var exchange tokenExchangeRequest if err := json.Unmarshal(body, &exchange); err != nil { s.t.Errorf("decode exchange request: %v", err) } assert.Equal(s.t, s.invocation.TiDBOrgID, exchange.TiDBOrgID) assert.Equal(s.t, s.invocation.OrgSlug, exchange.OrgSlug) assert.Equal(s.t, s.invocation.ContextType, exchange.ContextType) assert.Equal(s.t, s.invocation.ContextID, exchange.ContextID) assert.Equal(s.t, s.invocation.InitiatedBy, exchange.InitiatedBy) s.lastExchangeBody = append([]byte(nil), body...) s.exchanges++ s.latestToken = fmt.Sprintf("test-token-%d", s.exchanges) _, _ = fmt.Fprintf(response, `{"success":true,"data":{"accessToken":%q,"tokenType":"Bearer","expiresAt":"2026-08-02T05:16:07Z","orgSlug":%q,"tenantId":"test-tenant"}}`, s.latestToken, s.resolvedOrgSlug) } func (s *contractServer) serveBusiness(response http.ResponseWriter, request *http.Request, body []byte) { s.mu.Lock() expectedToken := s.latestToken s.latestToken = "" s.businessCalls++ s.mu.Unlock() assert.True(s.t, strings.HasPrefix(request.URL.Path, businessPathPrefix+"/"), request.URL.Path) assert.NotEmpty(s.t, expectedToken, "business call did not exchange a fresh token") assert.Equal(s.t, "Bearer "+expectedToken, request.Header.Get("Authorization")) assert.NotEmpty(s.t, request.Header.Get("X-Request-ID")) assert.Equal(s.t, s.invocation.ContextType, request.Header.Get("X-TiDB-Context-Type")) assert.Equal(s.t, s.invocation.ContextID, request.Header.Get("X-TiDB-Context-Id")) var payload map[string]any if err := json.Unmarshal(body, &payload); err != nil { s.t.Errorf("decode business request: %v", err) } assert.Equal(s.t, s.resolvedOrgSlug, payload["orgSlug"]) s.business(response, request, body) } func (s *contractServer) counts() (int, int) { s.mu.Lock() defer s.mu.Unlock() return s.exchanges, s.businessCalls } func (s *contractServer) exchangeBody() []byte { s.mu.Lock() defer s.mu.Unlock() return append([]byte(nil), s.lastExchangeBody...) } func newContractClient(t *testing.T, server *httptest.Server, state *contractServer) *HTTPClient { t.Helper() var idMu sync.Mutex idSequence := 0 ids := func(prefix string) IDGenerator { return func() (string, error) { idMu.Lock() defer idMu.Unlock() idSequence++ return fmt.Sprintf("%s-%d", prefix, idSequence), nil } } client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret(state.secret), HTTPClient: server.Client(), Clock: func() time.Time { return testNow }, Nonce: ids("nonce"), RequestID: ids("request"), }) require.NoError(t, err) return client } func testDataSourceSpec(name string) DataSourceSpec { return DataSourceSpec{ Name: name, S3: S3Credentials{ Bucket: "test-bucket", Region: "us-test-1", Endpoint: "https://s3.test.invalid", AccessKey: NewSecret("test-access-key"), SecretKey: NewSecret("test-secret-key"), SQSQueueURL: "https://sqs.test.invalid/test-queue.fifo", SQSRegion: "us-test-1", ExternalID: NewSecret("test-external-id"), }, } } func testTaskSpec(name, warehouse string) TaskSpec { pollInterval := int32(30) return TaskSpec{ Name: name, DataSourceName: "test-datasource", Warehouse: warehouse, ContinuousIngestion: true, TiDB: TiDBTaskSpec{ SyncMode: SyncModeSnapshotAndCDC, TableRules: TableRules{"test_database.orders", "test_database.users", "!test_database.audit"}, PollIntervalSeconds: &pollInterval, AllowDelete: true, DumplingS3Prefix: "pipelines/test/snapshot/", DumplingCSVSeparator: ",", DumplingCSVNull: `\N`, DumplingAutoCreateTable: true, DumplingOnError: "abort", ChangefeedS3Prefix: "pipelines/test/incremental/", TargetNaming: TargetNaming{ DatabasePrefix: "target_", TableSuffix: "_copy", }, }, } } func TestHTTPClientDataSourceContract(t *testing.T) { var requestedPaths []string business := func(response http.ResponseWriter, request *http.Request, body []byte) { requestedPaths = append(requestedPaths, request.URL.Path) switch request.URL.Path { case businessPathPrefix + "/datasource/test", businessPathPrefix + "/datasource/create": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, "test-datasource", payload["name"]) assert.Equal(t, ServiceIDTiDBCDC, payload["serviceId"]) assert.Equal(t, testInvocation.OrgSlug, payload["orgSlug"]) credentials := payload["tidbCredentials"].(map[string]any) assert.Equal(t, "test-bucket", credentials["s3Bucket"]) assert.Equal(t, "us-test-1", credentials["s3Region"]) assert.Equal(t, "test-access-key", credentials["s3AccessKey"]) assert.Equal(t, "test-secret-key", credentials["s3SecretKey"]) // A rules task carries no source database on the datasource. assert.NotContains(t, credentials, "database") if strings.HasSuffix(request.URL.Path, "/create") { _, _ = response.Write([]byte(`{"success":true,"id":11}`)) } else { _, _ = response.Write([]byte(`{"success":true}`)) } case businessPathPrefix + "/datasource/11/info": assert.JSONEq(t, `{"datasourceId":11,"orgSlug":"test-org"}`, string(body)) serviceArgs, _ := json.Marshal(map[string]any{ "s3Bucket": "test-bucket", "s3Region": "us-test-1", "s3AccessKey": "test-access-key", "s3SecretKey": "test-secret-key", "externalID": "test-external-id", "s3Endpoint": "https://s3.test.invalid", "roleARN": "arn:aws:iam::000000000000:role/test-lake-role", "sqsQueueURL": "https://sqs.test.invalid/test-queue.fifo", "sqsRegion": "us-test-1", }) encodedArgs, _ := json.Marshal(string(serviceArgs)) _, _ = fmt.Fprintf(response, `{"success":true,"data":{"id":11,"name":"test-datasource","serviceId":"TiDBCDC","serviceArgs":%s,"createdAt":"2026-08-02T04:00:00Z"}}`, encodedArgs) case businessPathPrefix + "/datasource/11/update": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, float64(11), payload["datasourceId"]) assert.Equal(t, "test-datasource-updated", payload["name"]) _, _ = response.Write([]byte(`{"success":true}`)) case businessPathPrefix + "/datasource/delete": assert.JSONEq(t, `{"datasourceIds":[11],"orgSlug":"test-org"}`, string(body)) _, _ = response.Write([]byte(`{"success":true}`)) case businessPathPrefix + "/datasources": assert.JSONEq(t, `{"page":1,"perPage":100,"serviceId":"TiDBCDC","orgSlug":"test-org"}`, string(body)) _, _ = response.Write([]byte(`{"success":true,"data":{"list":[{"id":11,"name":"test-datasource","serviceId":"TiDBCDC","serviceArgs":"{\"s3Bucket\":\"test-bucket\",\"s3Region\":\"us-test-1\",\"s3AccessKey\":\"test-access-key\",\"s3SecretKey\":\"test-secret-key\",\"externalID\":\"test-external-id\",\"database\":\"test_database\"}"}],"total":1}}`)) default: t.Errorf("unexpected path %s", request.URL.Path) response.WriteHeader(http.StatusNotFound) } } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testDataSourceSpec("test-datasource") require.NoError(t, client.TestDataSource(context.Background(), testInvocation, spec)) require.NoError(t, client.CreateDataSource(context.Background(), testInvocation, spec)) dataSource, err := client.GetDataSource(context.Background(), testInvocation, "test-datasource") require.NoError(t, err) assert.Equal(t, "test-bucket", dataSource.S3.Bucket) assert.Equal(t, "us-test-1", dataSource.S3.Region) assert.Equal(t, "https://s3.test.invalid", dataSource.S3.Endpoint) assert.Equal(t, "arn:aws:iam::000000000000:role/test-lake-role", dataSource.S3.RoleARN) assert.True(t, dataSource.S3.AccessKey.Empty()) assert.True(t, dataSource.S3.SecretKey.Empty()) assert.True(t, dataSource.S3.ExternalID.Empty()) // The queue round-trips: the matcher that decides whether a datasource // still matches its spec compares it, so dropping it here would make // attaching a queue never converge and detaching one never happen. assert.Equal(t, "https://sqs.test.invalid/test-queue.fifo", dataSource.S3.SQSQueueURL) assert.Equal(t, "us-test-1", dataSource.S3.SQSRegion) serialized, err := json.Marshal(dataSource) require.NoError(t, err) assert.NotContains(t, string(serialized), "test-access-key") assert.NotContains(t, string(serialized), "test-secret-key") assert.NotContains(t, string(serialized), "test-external-id") updated := testDataSourceSpec("test-datasource-updated") require.NoError(t, client.UpdateDataSource(context.Background(), testInvocation, "test-datasource", updated)) assert.True(t, dataSource.S3.AccessKey.Empty()) require.NoError(t, client.DeleteDataSource(context.Background(), testInvocation, "test-datasource")) exchanges, calls := state.counts() assert.Equal(t, calls, exchanges, "every business call must exchange a fresh token") // Each name-keyed call resolves the name through the list endpoint first. assert.Len(t, requestedPaths, 8) } func TestHTTPClientTaskAndObservationContract(t *testing.T) { resolvedTS := uint64(testNow.Add(-45*time.Second).UnixMilli()) << 18 createdAt := testNow.Add(-time.Minute).Format(time.RFC3339) var operations []string business := func(response http.ResponseWriter, request *http.Request, body []byte) { switch request.URL.Path { case businessPathPrefix + "/task/create": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, "warehouse / exact", payload["warehouse"], "warehouse must be passed unchanged") assert.Equal(t, float64(1), payload["continuousIngestion"]) tidbTask := payload["tidbTask"].(map[string]any) assert.Equal(t, string(SyncModeSnapshotAndCDC), tidbTask["syncMode"]) assert.Equal(t, []any{"test_database.orders", "test_database.users", "!test_database.audit"}, tidbTask["tableRules"]) assert.NotContains(t, tidbTask, "includeTables") assert.NotContains(t, tidbTask, "dumplingCsvSkipHeader", "omitted and zero have different semantics") _, _ = response.Write([]byte(`{"success":true,"id":22}`)) case businessPathPrefix + "/task/22/info": assert.JSONEq(t, `{"taskId":22,"orgSlug":"test-org"}`, string(body)) _, _ = fmt.Fprintf(response, `{"success":true,"data":{"id":22,"name":"test-task","datasourceId":11,"warehouse":"warehouse / exact","continuousIngestion":1,"taskStatus":"TaskRunning","database":"target_test_database","tidbTask":{"syncMode":"TiDBSyncModeSnapshotAndCdc","tableRules":["test_database.orders","test_database.users","!test_database.audit"],"pollIntervalSeconds":30},"lastRunData":{"kind":"tidb","tidbLastRun":{"lastResolvedTS":"%d","rowsSynced":"9","tablesActive":2,"tablesError":0}},"createdAt":%q}}`, resolvedTS, createdAt) case businessPathPrefix + "/task/22/update": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, float64(22), payload["taskId"]) assert.Equal(t, "warehouse / exact", payload["warehouse"]) assert.Equal(t, float64(0), payload["tidbTask"].(map[string]any)["dumplingCsvSkipHeader"]) _, _ = response.Write([]byte(`{"success":true}`)) case businessPathPrefix + "/task/operate": var payload struct { Operate string `json:"operate"` } assert.NoError(t, json.Unmarshal(body, &payload)) operations = append(operations, payload.Operate) _, _ = response.Write([]byte(`{"success":true}`)) case businessPathPrefix + "/datasources": // The client addresses Lake by name and resolves the wire's // numeric id itself. _, _ = response.Write([]byte(`{"success":true,"data":{"list":[{"id":11,"name":"test-datasource","serviceId":"TiDBCDC","serviceArgs":{"s3Bucket":"b","s3Region":"r","database":"test_database"}}],"total":1}}`)) case businessPathPrefix + "/tasks": assert.JSONEq(t, `{"page":1,"perPage":100,"serviceId":"TiDBCDC","orgSlug":"test-org"}`, string(body)) _, _ = response.Write([]byte(`{"success":true,"data":{"list":[{"id":22,"name":"test-task","datasourceId":11,"warehouse":"warehouse / exact","taskStatus":1,"tidbTask":{"syncMode":3}}],"total":1}}`)) case businessPathPrefix + "/task/runs": // Two callers reach this: ListTaskRuns with the caller's own // options, and ObserveTask, which reads every page. payload := map[string]any{} require.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, float64(22), payload["taskId"]) assert.Equal(t, "test-org", payload["orgSlug"]) assert.Contains(t, []any{float64(20), float64(listPageSize)}, payload["perPage"]) _, _ = fmt.Fprintf(response, `{"success":true,"data":{"list":[{"executeId":"tidb-22-test","taskId":22,"status":1,"responseData":{"kind":"tidb","tidbLastRun":{"lastResolvedTS":"%d","rowsSynced":"10","insertCount":"4","updateCount":"3","deleteCount":"3","tablesActive":2,"tablesError":0}},"createdAt":%q}],"total":1}}`, resolvedTS+(1<<18), createdAt) case businessPathPrefix + "/task/runs/tidb-22-test": assert.JSONEq(t, `{"executeId":"tidb-22-test","orgSlug":"test-org"}`, string(body)) _, _ = fmt.Fprintf(response, `{"success":true,"data":{"executeId":"tidb-22-test","taskId":22,"status":"TaskRunRunning","requestData":{"kind":"tidb"},"responseData":{"tidbLastRun":{"lastResolvedTS":"%d","rowsSynced":"10"}}}}`, resolvedTS+(1<<18)) default: t.Errorf("unexpected path %s", request.URL.Path) response.WriteHeader(http.StatusNotFound) } } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testTaskSpec("test-task", "warehouse / exact") require.NoError(t, client.CreateTask(context.Background(), testInvocation, spec)) task, err := client.GetTask(context.Background(), testInvocation, "test-task") require.NoError(t, err) assert.Equal(t, TaskStatusRunning, task.Status) assert.Equal(t, resolvedTS, task.LastProgress.LastResolvedTS) assert.Equal(t, uint64(9), task.LastProgress.RowsSynced) skipHeader := int32(0) spec.TiDB.DumplingCSVSkipHeader = &skipHeader require.NoError(t, client.UpdateTask(context.Background(), testInvocation, "test-task", spec)) require.NoError(t, client.StartTask(context.Background(), testInvocation, "test-task")) require.NoError(t, client.StopTask(context.Background(), testInvocation, "test-task")) require.NoError(t, client.DropTask(context.Background(), testInvocation, "test-task")) assert.Equal(t, []string{"Start", "Stop", "Drop"}, operations) runPage, err := client.ListTaskRuns(context.Background(), testInvocation, "test-task", ListTaskRunsOptions{}) require.NoError(t, err) require.Len(t, runPage.Runs, 1) assert.Equal(t, uint64(10), runPage.Runs[0].Progress.RowsSynced) run, err := client.GetTaskRun(context.Background(), testInvocation, "tidb-22-test") require.NoError(t, err) assert.Equal(t, TaskRunStatusRunning, run.Status) observation, err := client.ObserveTask(context.Background(), testInvocation, "test-task") require.NoError(t, err) assert.Equal(t, resolvedTS+(1<<18), observation.LastResolvedTS) assert.Equal(t, testNow.Add(-45*time.Second).Add(time.Millisecond), observation.LastResolvedAt) assert.Equal(t, testNow, observation.ObservedAt) assert.Equal(t, "tidb-22-test", observation.LatestRunID) exchanges, calls := state.counts() assert.Equal(t, calls, exchanges) } func TestHTTPClientFailsClosedForMultipleDatabasesBeforeHTTP(t *testing.T) { requestCount := 0 server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { requestCount++ response.WriteHeader(http.StatusInternalServerError) })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) datasource := testDataSourceSpec("test-datasource") datasource.Name = "" require.Error(t, client.TestDataSource(context.Background(), testInvocation, datasource)) task := testTaskSpec("test-task", "test-warehouse") task.TiDB.TableRules = nil require.Error(t, client.CreateTask(context.Background(), testInvocation, task)) assert.Zero(t, requestCount, "spec rejection must happen before token exchange") } func TestHTTPClientRoleARNDataSourceContract(t *testing.T) { business := func(response http.ResponseWriter, _ *http.Request, body []byte) { var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) credentials := payload["tidbCredentials"].(map[string]any) assert.Equal(t, "arn:aws:iam::000000000000:role/test-lake-role", credentials["roleARN"]) assert.Equal(t, "test-external-id", credentials["externalID"]) assert.Equal(t, "us-test-1", credentials["s3Region"]) assert.NotContains(t, credentials, "s3AccessKey") assert.NotContains(t, credentials, "s3SecretKey") _, _ = response.Write([]byte(`{"success":true}`)) } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testDataSourceSpec("role-datasource") spec.S3.AccessKey = Secret{} spec.S3.SecretKey = Secret{} spec.S3.RoleARN = "arn:aws:iam::000000000000:role/test-lake-role" require.NoError(t, client.TestDataSource(context.Background(), testInvocation, spec)) } func TestHTTPClientRedactsSecretsFromErrorsAndValues(t *testing.T) { business := func(response http.ResponseWriter, _ *http.Request, _ []byte) { response.WriteHeader(http.StatusBadRequest) _, _ = response.Write([]byte(`{"error":"BadRequest","message":"credentials test-access-key test-secret-key test-external-id test-hmac-secret"}`)) } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testDataSourceSpec("test-datasource") err := client.TestDataSource(context.Background(), testInvocation, spec) var apiErr *APIError require.ErrorAs(t, err, &apiErr) assert.Equal(t, "BadRequest", apiErr.Code) for _, secret := range []string{"test-access-key", "test-secret-key", "test-external-id", "test-hmac-secret"} { assert.NotContains(t, err.Error(), secret) assert.NotContains(t, apiErr.Message, secret) } serialized, marshalErr := json.Marshal(spec) require.NoError(t, marshalErr) assert.NotContains(t, string(serialized), "test-access-key") assert.NotContains(t, string(serialized), "test-secret-key") assert.NotContains(t, fmt.Sprintf("%#v", NewSecret("test-secret-key")), "test-secret-key") } func TestHTTPClientUnknownCreateResultCanBeAdoptedByDeterministicName(t *testing.T) { names, err := DeterministicResourceNames("pipeline/test/001") require.NoError(t, err) created := false business := func(response http.ResponseWriter, request *http.Request, body []byte) { switch request.URL.Path { case businessPathPrefix + "/datasource/create": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, names.DataSource, payload["name"]) created = true _, _ = response.Write([]byte(`{"success":true,"id":`)) case businessPathPrefix + "/datasource/71/info": _, _ = fmt.Fprintf(response, `{"success":true,"data":{"id":71,"name":%q,"serviceId":"TiDBCDC","serviceArgs":{"s3Bucket":"test-bucket","s3Region":"us-test-1","database":"test_database"}}}`, names.DataSource) case businessPathPrefix + "/datasources": if !created { t.Error("the name lookup ran before the simulated create") } _, _ = fmt.Fprintf(response, `{"success":true,"data":{"list":[{"id":71,"name":%q,"serviceId":"TiDBCDC","serviceArgs":{"s3Bucket":"test-bucket","s3Region":"us-test-1","database":"test_database"}}],"total":1}}`, names.DataSource) default: t.Errorf("unexpected path %s", request.URL.Path) } } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testDataSourceSpec(names.DataSource) err = client.CreateDataSource(context.Background(), testInvocation, spec) var unknown *UnknownWriteResultError assert.True(t, errors.As(err, &unknown)) assert.NotContains(t, err.Error(), "test-secret-key") found, err := client.GetDataSource(context.Background(), testInvocation, names.DataSource) require.NoError(t, err) assert.Equal(t, int32(71), found.ID) } func TestHTTPClientServerErrorMakesWriteResultUnknown(t *testing.T) { business := func(response http.ResponseWriter, _ *http.Request, _ []byte) { response.WriteHeader(http.StatusServiceUnavailable) _, _ = response.Write([]byte(`{"error":"InternalError","message":"temporary failure"}`)) } server, state := newContractServer(t, business) client := newContractClient(t, server, state) err := client.CreateDataSource(context.Background(), testInvocation, testDataSourceSpec("test-datasource")) var unknown *UnknownWriteResultError assert.True(t, errors.As(err, &unknown)) var apiErr *APIError require.ErrorAs(t, err, &apiErr) assert.Equal(t, http.StatusServiceUnavailable, apiErr.StatusCode) } func TestHTTPClientResolvesOrganizationForOrgIDOnlyInvocation(t *testing.T) { invocation := Invocation{TiDBOrgID: "1000000000000000001", LakeOrgID: "7"} business := func(response http.ResponseWriter, _ *http.Request, body []byte) { var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, "resolved-test-org", payload["orgSlug"]) _, _ = response.Write([]byte(`{"success":true}`)) } server, state := newContractServer(t, business) state.invocation = invocation state.resolvedOrgSlug = "resolved-test-org" client := newContractClient(t, server, state) require.NoError(t, client.TestDataSource(context.Background(), invocation, testDataSourceSpec("test-datasource"))) assert.NotContains(t, string(state.exchangeBody()), "orgSlug", "orgSlug must be omitted when the caller supplies only an org ID") exchanges, businessCalls := state.counts() assert.Equal(t, 1, exchanges) assert.Equal(t, 1, businessCalls) } func TestHTTPClientRejectsResolvedOrganizationMismatch(t *testing.T) { invocation := Invocation{TiDBOrgID: "1000000000000000001", LakeOrgID: "7", OrgSlug: "caller-test-org"} business := func(http.ResponseWriter, *http.Request, []byte) { t.Error("business request must not be sent after an organization mismatch") } server, state := newContractServer(t, business) state.invocation = invocation state.resolvedOrgSlug = "different-test-org" client := newContractClient(t, server, state) err := client.TestDataSource(context.Background(), invocation, testDataSourceSpec("test-datasource")) assert.EqualError(t, err, "Lake token exchange returned a different organization") assert.NotContains(t, err.Error(), invocation.OrgSlug) assert.NotContains(t, err.Error(), state.resolvedOrgSlug) exchanges, businessCalls := state.counts() assert.Equal(t, 1, exchanges) assert.Zero(t, businessCalls) } func TestHTTPClientRequiresResolvedOrganization(t *testing.T) { invocation := Invocation{TiDBOrgID: "1000000000000000001", LakeOrgID: "7"} server, state := newContractServer(t, func(http.ResponseWriter, *http.Request, []byte) { t.Error("business request must not be sent without a resolved organization") }) state.invocation = invocation state.resolvedOrgSlug = "" client := newContractClient(t, server, state) err := client.TestDataSource(context.Background(), invocation, testDataSourceSpec("test-datasource")) assert.EqualError(t, err, "Lake token exchange response has no organization slug") _, businessCalls := state.counts() assert.Zero(t, businessCalls) } func TestHTTPClientRedactsTokenExchangeErrors(t *testing.T) { const ( hmacSecret = "test-token-exchange-hmac" orgSlug = "test-org-identifier" ) server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { assert.Equal(t, tokenExchangePath, request.URL.Path) response.WriteHeader(http.StatusUnauthorized) _, _ = fmt.Fprintf(response, `{"error":"AuthFailed","message":"organization %s rejected secret %s"}`, orgSlug, hmacSecret) })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret(hmacSecret), HTTPClient: server.Client(), }) require.NoError(t, err) err = client.TestDataSource(context.Background(), Invocation{ TiDBOrgID: "1000000000000000001", LakeOrgID: "7", OrgSlug: orgSlug, }, testDataSourceSpec("test-datasource")) var apiErr *APIError require.ErrorAs(t, err, &apiErr) assert.Equal(t, "AuthFailed", apiErr.Code) assert.NotContains(t, apiErr.Message, hmacSecret) assert.Contains(t, apiErr.Message, "[REDACTED]") // The organization slug survives. It identifies the tenant rather than // authenticating it - callers are handed the same value back as the // external ID - and an error that cannot say which organization was // rejected is markedly harder to act on. assert.Contains(t, apiErr.Message, orgSlug) } func TestHTTPClientContextHeadersAreAllOrNothing(t *testing.T) { withoutContext := Invocation{TiDBOrgID: "1000000000000000001", LakeOrgID: "7", OrgSlug: "test-org"} business := func(response http.ResponseWriter, request *http.Request, _ []byte) { assert.Empty(t, request.Header.Get("X-TiDB-Context-Type")) assert.Empty(t, request.Header.Get("X-TiDB-Context-Id")) _, _ = response.Write([]byte(`{"success":true}`)) } server, state := newContractServer(t, business) state.invocation = withoutContext client := newContractClient(t, server, state) require.NoError(t, client.TestDataSource(context.Background(), withoutContext, testDataSourceSpec("test-datasource"))) invalidContext := withoutContext invalidContext.ContextType = "data_pipeline" err := client.TestDataSource(context.Background(), invalidContext, testDataSourceSpec("test-datasource")) assert.ErrorContains(t, err, "must be provided together") } func TestSignMatchesDocumentedCanonicalForm(t *testing.T) { body := []byte(`{"tidbOrgId":"1000000000000000001"}`) actual := sign("test-hmac-secret", "post", tokenExchangePath, 1785647167, "test-nonce", body) assert.Equal(t, "7ed2d4c07f8dd4e2607414e1d10d837fd62e2183e2ce4ff3f31c7ca9898d9f3d", actual) } func TestHTTPClientConstructorValidation(t *testing.T) { _, err := NewHTTPClient(HTTPClientConfig{BaseURL: "https://user:password@example.invalid/path", HMACSecret: NewSecret("secret")}) assert.Error(t, err) _, err = NewHTTPClient(HTTPClientConfig{BaseURL: "https://example.invalid", HMACSecret: Secret{}}) assert.Error(t, err) } // TestTaskRunNeverCarriesTheRequestSnapshot pins the redaction rule for the // task-run payload: Lake echoes back the request that created the run, and for // a TiDB CDC task that request contains the S3 credentials. func TestTaskRunNeverCarriesTheRequestSnapshot(t *testing.T) { t.Parallel() runType := reflect.TypeOf(TaskRun{}) for index := 0; index < runType.NumField(); index++ { name := runType.Field(index).Name assert.NotContains(t, strings.ToLower(name), "request", "TaskRun must not expose the request snapshot Lake echoes back") } } // TestBusinessFailureInsideOKEnvelopePreservesTheCode covers the second shape // Lake uses to report a missing resource: HTTP 200 with success=false. Losing // the code here makes a dropped task indistinguishable from a transient error, // so the caller polls for a child that will never come back. func TestBusinessFailureInsideOKEnvelopePreservesTheCode(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) if request.URL.Path == tokenExchangePath { _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) return } _, _ = response.Write([]byte( `{"success":false,"error":"NotFound","message":"task 7 does not exist"}`)) })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) _, err = client.GetTask(context.Background(), testInvocation, "test-task") require.Error(t, err) var apiErr *APIError require.ErrorAs(t, err, &apiErr) require.Equal(t, "NotFound", apiErr.Code) require.Contains(t, apiErr.Message, "task 7 does not exist") } // TestHTTPClientUpdateTaskResolvesTheRequestedDatasource pins that an update // honours the datasource the spec names. Reusing the task's existing one would // let the update report success while the task kept ingesting from the old // bucket and database. func TestHTTPClientUpdateTaskResolvesTheRequestedDatasource(t *testing.T) { var sentDataSourceID float64 business := func(response http.ResponseWriter, request *http.Request, body []byte) { switch request.URL.Path { case businessPathPrefix + "/tasks": _, _ = response.Write([]byte(`{"success":true,"data":{"list":[{"id":22,"name":"test-task","datasourceId":11}],"total":1}}`)) case businessPathPrefix + "/task/22/info": _, _ = response.Write([]byte(`{"success":true,"data":{"id":22,"name":"test-task","datasourceId":11,"warehouse":"w","taskStatus":1,"tidbTask":{"syncMode":3}}}`)) case businessPathPrefix + "/datasources": // The task currently points at 11; the spec asks for "moved-to", 12. _, _ = response.Write([]byte(`{"success":true,"data":{"list":[{"id":12,"name":"moved-to","serviceId":"TiDBCDC","serviceArgs":{"s3Bucket":"b","s3Region":"r","database":"test_database"}}],"total":1}}`)) case businessPathPrefix + "/task/22/update": var payload map[string]any assert.NoError(t, json.Unmarshal(body, &payload)) sentDataSourceID = payload["datasourceId"].(float64) _, _ = response.Write([]byte(`{"success":true}`)) default: t.Errorf("unexpected path %s", request.URL.Path) } } server, state := newContractServer(t, business) client := newContractClient(t, server, state) spec := testTaskSpec("test-task", "test-warehouse") spec.DataSourceName = "moved-to" require.NoError(t, client.UpdateTask(context.Background(), testInvocation, "test-task", spec)) assert.Equal(t, float64(12), sentDataSourceID, "the update must carry the datasource the spec named, not the one the task had") } // TestHTTPClientDeleteDataSourceIsIdempotent pins that the HTTP client and the // fake agree: an absent datasource is the state Delete converges to. func TestHTTPClientDeleteDataSourceIsIdempotent(t *testing.T) { deleteCalled := false business := func(response http.ResponseWriter, request *http.Request, _ []byte) { switch request.URL.Path { case businessPathPrefix + "/datasources": _, _ = response.Write([]byte(`{"success":true,"data":{"list":[],"total":0}}`)) case businessPathPrefix + "/datasource/delete": deleteCalled = true _, _ = response.Write([]byte(`{"success":true}`)) default: t.Errorf("unexpected path %s", request.URL.Path) } } server, state := newContractServer(t, business) client := newContractClient(t, server, state) require.NoError(t, client.DeleteDataSource(context.Background(), testInvocation, "already-gone")) assert.False(t, deleteCalled, "nothing to delete means no write is issued") } // TestHTTPClientDecodesLakesOwnDataSourceProjection pins the response shape a // real Lake deployment returns, which differs from the request shape in two // ways that both broke decoding: serviceId comes back as Lake's internal // numeric service key rather than the "TiDBCDC" alias requests use, and // serviceArgs arrives as a JSON-encoded string rather than an object. Decoding // either one strictly fails the whole call, and every lifecycle that looks a // datasource up by name then reports a transport error instead of an answer. func TestHTTPClientDecodesLakesOwnDataSourceProjection(t *testing.T) { t.Parallel() const listBody = `{"success":true,"data":{"total":1,"page":1,"perPage":100,"list":[` + `{"id":15,"orgId":"1","name":"test-datasource","serviceId":6,` + `"serviceArgs":"{\"s3Bucket\":\"test-bucket\",\"s3Region\":\"us-west-2\",\"database\":\"app\"}",` + `"createdAt":"2026-08-10T20:03:48Z","updatedAt":"2026-08-10T20:03:48Z"}]}}` const infoBody = `{"success":true,"data":{"id":15,"orgId":"1","name":"test-datasource","serviceId":6,` + `"serviceArgs":"{\"s3Bucket\":\"test-bucket\",\"s3Region\":\"us-west-2\",\"s3Endpoint\":\"https://s3.example\",\"database\":\"app\"}",` + `"createdAt":"2026-08-10T20:03:48Z","updatedAt":"2026-08-10T20:03:48Z"}}` server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) switch { case request.URL.Path == tokenExchangePath: _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) case strings.HasSuffix(request.URL.Path, "/datasources"): _, _ = response.Write([]byte(listBody)) default: _, _ = response.Write([]byte(infoBody)) } })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) dataSource, err := client.GetDataSource(context.Background(), testInvocation, "test-datasource") require.NoError(t, err) require.Equal(t, int32(15), dataSource.ID) require.Equal(t, "test-datasource", dataSource.Name) require.Equal(t, "test-bucket", dataSource.S3.Bucket) require.Equal(t, "us-west-2", dataSource.S3.Region) require.Equal(t, "https://s3.example", dataSource.S3.Endpoint) } // TestHTTPClientSetupInfoContract pins the shape Lake answers with, using the // response a real deployment returned. This is the entry point for role-ARN // access: a caller that drops a principal hands the customer a trust policy // that validates and then fails during the load. func TestHTTPClientSetupInfoContract(t *testing.T) { t.Parallel() const body = `{"success":true,"data":{"externalID":"test-org","trustPrincipals":[` + `{"name":"Platform setup and validation role","arn":"arn:aws:iam::102915632338:role/databend-webapi","required":true},` + `{"name":"Platform data loading role","arn":"arn:aws:iam::102915632338:role/moon/tny42rdpm/moon-tny42rdpm","required":true}]}}` var sentOrgSlug any server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) if request.URL.Path == tokenExchangePath { _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) return } assert.True(t, strings.HasSuffix(request.URL.Path, "/tidb/setup-info")) payload := map[string]any{} require.NoError(t, json.NewDecoder(request.Body).Decode(&payload)) sentOrgSlug = payload["orgSlug"] _, _ = response.Write([]byte(body)) })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) info, err := client.SetupInfo(context.Background(), testInvocation) require.NoError(t, err) // The organization is resolved by the token exchange, not asserted by the // caller, and Lake keys the external ID off it. require.Equal(t, "test-org", sentOrgSlug) require.Equal(t, "test-org", info.ExternalID) require.Len(t, info.TrustPrincipals, 2) require.Equal(t, "arn:aws:iam::102915632338:role/databend-webapi", info.TrustPrincipals[0].ARN) require.Equal(t, "Platform data loading role", info.TrustPrincipals[1].Name) require.True(t, info.TrustPrincipals[1].Required) } // TestHTTPClientObserveTaskReadsTheNewestRuns pins the window an observation // reads, and what falls outside it. A continuous task accumulates a run per // resume, so its oldest run - the snapshot run carrying fullSyncCompleted - // eventually pages out. Reading the whole history on every poll would cost a // request per hundred runs forever; callers that need those facts across the // task's life record them instead. func TestHTTPClientObserveTaskReadsTheNewestRuns(t *testing.T) { t.Parallel() requestedPages := []float64{} server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) switch { case request.URL.Path == tokenExchangePath: _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) case strings.HasSuffix(request.URL.Path, "/tasks"): _, _ = response.Write([]byte( `{"success":true,"data":{"list":[{"id":9,"name":"test-task"}],"total":1}}`)) case strings.HasSuffix(request.URL.Path, "/runs"): payload := map[string]any{} require.NoError(t, json.NewDecoder(request.Body).Decode(&payload)) requestedPages = append(requestedPages, payload["page"].(float64)) assert.Equal(t, float64(listPageSize), payload["perPage"]) _, _ = fmt.Fprint(response, `{"success":true,"data":{"list":[`+ `{"executeId":"cdc-2","taskId":9,"status":"TaskRunRunning","responseData":{"kind":"tidb","tidbLastRun":{"rowsSynced":"7","lastResolvedTS":"99"}}},`+ `{"executeId":"dumpling-1","taskId":9,"status":"TaskRunSuccess","responseData":{"kind":"tidb","tidbLastRun":{"fullSyncCompleted":true,"backupTSO":"4321"}}}`+ `],"total":2}}`) default: _, _ = response.Write([]byte(`{"success":true,"data":{"id":9,"name":"test-task"}}`)) } })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) observed, err := client.ObserveTask(context.Background(), testInvocation, "test-task") require.NoError(t, err) require.Equal(t, []float64{1}, requestedPages, "one poll must cost one page") // Both runs are inside the window, so the disjoint halves still merge: the // snapshot run's sticky facts and the newest run's own counters. require.True(t, observed.Progress.FullSyncCompleted) require.Equal(t, uint64(4321), observed.Progress.BackupTSO) require.Equal(t, uint64(7), observed.Progress.RowsSynced) } // The failure reason travels only in the run payload, and everything above the // decode is exercised against the fake. A wrong JSON tag on LastMessage or // Error would therefore break nothing anywhere else, which is the drift this // whole change is about - so pin it here, on the real decode path. func TestHTTPClientObserveTaskCarriesTheRunFailureReason(t *testing.T) { t.Parallel() const wantMessage = "create target database e2e_crm: query error: code: 0, " + "message: warehouse tn8upxuaq dp_e2e_wh not found: Bad warehouse" for name, runJSON := range map[string]string{ // What staging actually returned: both fields set, to the same text. "lastMessage and error": `{"executeId":"tidb-1","taskId":9,"status":"TaskRunFail",` + `"responseData":{"kind":"tidb","error":"` + wantMessage + `",` + `"tidbLastRun":{"lastMessage":"` + wantMessage + `"}}}`, "lastMessage only": `{"executeId":"tidb-1","taskId":9,"status":"TaskRunFail",` + `"responseData":{"kind":"tidb","tidbLastRun":{"lastMessage":"` + wantMessage + `"}}}`, // The envelope-level fallback, for a run that reports the failure there. "error only": `{"executeId":"tidb-1","taskId":9,"status":"TaskRunFail",` + `"responseData":{"kind":"tidb","error":"` + wantMessage + `"}}`, } { t.Run(name, func(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc( func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) switch { case request.URL.Path == tokenExchangePath: _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) case strings.HasSuffix(request.URL.Path, "/tasks"): _, _ = response.Write([]byte( `{"success":true,"data":{"list":[{"id":9,"name":"test-task","status":"TaskStopped"}],"total":1}}`)) case strings.HasSuffix(request.URL.Path, "/runs"): _, _ = fmt.Fprint(response, `{"success":true,"data":{"list":[`+runJSON+`],"total":1}}`) default: _, _ = response.Write([]byte( `{"success":true,"data":{"id":9,"name":"test-task","status":"TaskStopped"}}`)) } })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) observed, err := client.ObserveTask(context.Background(), testInvocation, "test-task") require.NoError(t, err) require.Equal(t, TaskRunStatusFail, observed.LatestRunStatus) require.Equal(t, wantMessage, observed.LatestRunMessage) }) } } // A healthy run says nothing, and must not be given a message invented from an // empty field. func TestHTTPClientObserveTaskLeavesAHealthyRunUnexplained(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc( func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) switch { case request.URL.Path == tokenExchangePath: _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) case strings.HasSuffix(request.URL.Path, "/tasks"): _, _ = response.Write([]byte( `{"success":true,"data":{"list":[{"id":9,"name":"test-task"}],"total":1}}`)) case strings.HasSuffix(request.URL.Path, "/runs"): _, _ = response.Write([]byte( `{"success":true,"data":{"list":[{"executeId":"tidb-1","taskId":9,` + `"status":"TaskRunRunning","responseData":{"kind":"tidb",` + `"tidbLastRun":{"rowsSynced":"3"}}}],"total":1}}`)) default: _, _ = response.Write([]byte(`{"success":true,"data":{"id":9,"name":"test-task"}}`)) } })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) observed, err := client.ObserveTask(context.Background(), testInvocation, "test-task") require.NoError(t, err) require.Empty(t, observed.LatestRunMessage) } // A TiDB organization may own several Lakes, so the TiDB organization id alone // does not resolve to one. The token exchange rejects a request that omits the // Lake, and this pins that we send it - a client that dropped the field would // fail every call at the exchange, before any business request is attempted. func TestHTTPClientTokenExchangeCarriesTheLakeOrganization(t *testing.T) { t.Parallel() var sent map[string]any server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { response.WriteHeader(http.StatusOK) if request.URL.Path == tokenExchangePath { require.NoError(t, json.NewDecoder(request.Body).Decode(&sent)) _, _ = response.Write([]byte( `{"success":true,"data":{"accessToken":"test-token","orgSlug":"test-org"}}`)) return } _, _ = response.Write([]byte(`{"success":true,"data":{"externalID":"test-org"}}`)) })) t.Cleanup(server.Close) client, err := NewHTTPClient(HTTPClientConfig{ BaseURL: server.URL, HMACSecret: NewSecret("test-hmac-secret"), HTTPClient: server.Client(), }) require.NoError(t, err) _, err = client.SetupInfo(context.Background(), testInvocation) require.NoError(t, err) require.Equal(t, testInvocation.LakeOrgID, sent["lakeOrgId"]) require.Equal(t, testInvocation.TiDBOrgID, sent["tidbOrgId"]) } // Half an identity is not an identity: without the Lake the exchange cannot // pick one, so the client refuses before spending a request on it. func TestInvocationRequiresBothOrganizations(t *testing.T) { t.Parallel() for name, invocation := range map[string]Invocation{ "no Lake": {TiDBOrgID: "1000000000000000001"}, "zero Lake": {TiDBOrgID: "1000000000000000001", LakeOrgID: "0"}, "no TiDB org": {LakeOrgID: "7"}, "zero TiDB org": {TiDBOrgID: "0", LakeOrgID: "7"}, } { t.Run(name, func(t *testing.T) { t.Parallel() require.Error(t, invocation.validate()) }) } require.NoError(t, Invocation{TiDBOrgID: "1000000000000000001", LakeOrgID: "7"}.validate()) }