diff --git a/.env.example b/.env.example index b2ab905e2..29b1325c3 100644 --- a/.env.example +++ b/.env.example @@ -55,6 +55,15 @@ PEERDB_CATALOG_DATABASE=postgres PEERDB_CLICKHOUSE_AWS_S3_BUCKET_NAME=peerdb +# GCP workload identity deployment settings. The projected token must use +# .svc.id.goog as its audience. Project and cluster +# values can be omitted on GKE, where they are read from the metadata server. +# PEERDB_GCP_WORKLOAD_IDENTITY_TARGET_SERVICE_ACCOUNT= +# PEERDB_GCP_WORKLOAD_IDENTITY_TOKEN_FILE=/var/run/secrets/peerdb/gcp-token +# PEERDB_GCP_PROJECT_ID= +# PEERDB_GCP_CLUSTER_LOCATION= +# PEERDB_GCP_CLUSTER_NAME= + AWS_REGION=us-east-1 AWS_ACCESS_KEY_ID=_peerdb_minioadmin AWS_SECRET_ACCESS_KEY=_peerdb_minioadmin diff --git a/flow/connectors/bigquery/auth.go b/flow/connectors/bigquery/auth.go new file mode 100644 index 000000000..601abdedb --- /dev/null +++ b/flow/connectors/bigquery/auth.go @@ -0,0 +1,269 @@ +package connbigquery + +import ( + "context" + "encoding/json" + "fmt" + "net/url" + "os" + "strings" + + "cloud.google.com/go/auth/credentials" + "cloud.google.com/go/bigquery" + "cloud.google.com/go/compute/metadata" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" +) + +const ( + // BigQueryAuthTypeServiceAccount selects legacy service-account-key credentials. + BigQueryAuthTypeServiceAccount = "service_account" + // BigQueryAuthTypeServiceAccountWorkloadIdentity selects deployment-scoped workload identity credentials. + BigQueryAuthTypeServiceAccountWorkloadIdentity = "service_account_workload_identity" + + workloadIdentityServiceAccountEnv = "PEERDB_GCP_WORKLOAD_IDENTITY_TARGET_SERVICE_ACCOUNT" + //nolint:gosec // Environment variable name, not a credential. + workloadIdentityTokenFileEnv = "PEERDB_GCP_WORKLOAD_IDENTITY_TOKEN_FILE" + workloadIdentityProjectIDEnv = "PEERDB_GCP_PROJECT_ID" + workloadIdentityClusterLocationEnv = "PEERDB_GCP_CLUSTER_LOCATION" + workloadIdentityClusterNameEnv = "PEERDB_GCP_CLUSTER_NAME" + + //nolint:gosec // Fixed public Google authentication endpoints, not credentials. + googleSTSTokenURL = "https://sts.googleapis.com/v1/token" + //nolint:gosec // Fixed public Google authentication endpoints, not credentials. + googleIAMCredentialsURL = "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/" + //nolint:gosec // OAuth subject-token type identifier, not a credential. + jwtSubjectTokenType = "urn:ietf:params:oauth:token-type:jwt" +) + +type bigQueryCredentialConfig struct { + credentialType credentials.CredType + clientProjectID string + credentialsJSON []byte +} + +type workloadIdentityDeploymentConfig struct { + targetServiceAccount string + tokenFile string + projectID string + clusterLocation string + clusterName string +} + +type workloadIdentityConfigSource struct { + lookupEnv func(string) (string, bool) + projectID func(context.Context) (string, error) + instanceAttributeValue func(context.Context, string) (string, error) +} + +type externalAccountCredentials struct { + Type string `json:"type"` + Audience string `json:"audience"` + SubjectTokenType string `json:"subject_token_type"` + TokenURL string `json:"token_url"` + ServiceAccountImpersonationURL string `json:"service_account_impersonation_url"` + CredentialSource externalAccountCredentialSource `json:"credential_source"` +} + +type externalAccountCredentialSource struct { + File string `json:"file"` + Format externalAccountCredentialSourceFormat `json:"format"` +} + +type externalAccountCredentialSourceFormat struct { + Type string `json:"type"` +} + +func resolveBigQueryResource(config *protos.BigqueryConfig) (string, string, error) { + datasetID := config.GetDatasetId() + projectID := config.GetProjectId() + projectPart, datasetPart, found := strings.Cut(datasetID, ".") + if found && strings.Contains(datasetPart, ".") { + return "", "", fmt.Errorf( + "invalid dataset ID: %s. Ensure that it is just a single string or string1.string2", + datasetID, + ) + } + if projectPart != "" && datasetPart != "" { + datasetID = datasetPart + projectID = projectPart + } + return projectID, datasetID, nil +} + +func newBigQueryCredentialConfig( + ctx context.Context, + config *protos.BigqueryConfig, + resourceProjectID string, +) (*bigQueryCredentialConfig, error) { + authType := config.GetAuthType() + if authType != BigQueryAuthTypeServiceAccount && authType != BigQueryAuthTypeServiceAccountWorkloadIdentity { + return nil, fmt.Errorf( + "unsupported BigQuery auth_type %q: expected %q or %q", + authType, + BigQueryAuthTypeServiceAccount, + BigQueryAuthTypeServiceAccountWorkloadIdentity, + ) + } + + if authType == BigQueryAuthTypeServiceAccount { + serviceAccount, err := NewBigQueryServiceAccount(config) + if err != nil { + return nil, fmt.Errorf("failed to create BigQueryServiceAccount: %w", err) + } + serviceAccountJSON, err := json.Marshal(serviceAccount) //nolint:gosec // G117: credential struct marshaled for inline use + if err != nil { + return nil, fmt.Errorf("failed to marshal service account: %v", err) + } + return &bigQueryCredentialConfig{ + credentialsJSON: serviceAccountJSON, + credentialType: credentials.ServiceAccount, + clientProjectID: bigquery.DetectProjectID, + }, nil + } + + if resourceProjectID == "" { + return nil, fmt.Errorf("BigQuery project ID must be set in the peer when workload identity is selected") + } + + deploymentConfig, err := resolveWorkloadIdentityDeploymentConfig(ctx, workloadIdentityConfigSource{ + lookupEnv: os.LookupEnv, + projectID: metadata.ProjectIDWithContext, + instanceAttributeValue: metadata.InstanceAttributeValueWithContext, + }) + if err != nil { + return nil, err + } + credentialsJSON, err := deploymentConfig.credentialsJSON() + if err != nil { + return nil, err + } + return &bigQueryCredentialConfig{ + credentialsJSON: credentialsJSON, + credentialType: credentials.ExternalAccount, + clientProjectID: resourceProjectID, + }, nil +} + +func resolveWorkloadIdentityDeploymentConfig( + ctx context.Context, + source workloadIdentityConfigSource, +) (*workloadIdentityDeploymentConfig, error) { + targetServiceAccount := envValue(source.lookupEnv, workloadIdentityServiceAccountEnv) + if targetServiceAccount == "" { + return nil, fmt.Errorf( + "BigQuery workload identity requires deployment environment variable %s", + workloadIdentityServiceAccountEnv, + ) + } + tokenFile := envValue(source.lookupEnv, workloadIdentityTokenFileEnv) + if tokenFile == "" { + return nil, fmt.Errorf( + "BigQuery workload identity requires deployment environment variable %s", + workloadIdentityTokenFileEnv, + ) + } + + projectID, err := envOrMetadata( + ctx, + source.lookupEnv, + workloadIdentityProjectIDEnv, + source.projectID, + ) + if err != nil { + return nil, err + } + clusterLocation, err := envOrMetadata( + ctx, + source.lookupEnv, + workloadIdentityClusterLocationEnv, + func(ctx context.Context) (string, error) { + return source.instanceAttributeValue(ctx, "cluster-location") + }, + ) + if err != nil { + return nil, err + } + clusterName, err := envOrMetadata( + ctx, + source.lookupEnv, + workloadIdentityClusterNameEnv, + func(ctx context.Context) (string, error) { + return source.instanceAttributeValue(ctx, "cluster-name") + }, + ) + if err != nil { + return nil, err + } + + return &workloadIdentityDeploymentConfig{ + targetServiceAccount: targetServiceAccount, + tokenFile: tokenFile, + projectID: projectID, + clusterLocation: clusterLocation, + clusterName: clusterName, + }, nil +} + +func envValue(lookupEnv func(string) (string, bool), name string) string { + value, _ := lookupEnv(name) + return strings.TrimSpace(value) +} + +func envOrMetadata( + ctx context.Context, + lookupEnv func(string) (string, bool), + envName string, + metadataValue func(context.Context) (string, error), +) (string, error) { + if value := envValue(lookupEnv, envName); value != "" { + return value, nil + } + value, err := metadataValue(ctx) + if err != nil { + return "", fmt.Errorf( + "BigQuery workload identity requires %s or the corresponding GKE metadata: %w", + envName, + err, + ) + } + if strings.TrimSpace(value) == "" { + return "", fmt.Errorf( + "BigQuery workload identity requires %s or non-empty corresponding GKE metadata", + envName, + ) + } + return strings.TrimSpace(value), nil +} + +func (config *workloadIdentityDeploymentConfig) credentialsJSON() ([]byte, error) { + pool := config.projectID + ".svc.id.goog" + audience := fmt.Sprintf( + "identitynamespace:%s:https://container.googleapis.com/v1/projects/%s/locations/%s/clusters/%s", + pool, + config.projectID, + config.clusterLocation, + config.clusterName, + ) + credentialConfig := externalAccountCredentials{ + Type: "external_account", + Audience: audience, + SubjectTokenType: jwtSubjectTokenType, + TokenURL: googleSTSTokenURL, + ServiceAccountImpersonationURL: googleIAMCredentialsURL + + url.PathEscape(config.targetServiceAccount) + ":generateAccessToken", + CredentialSource: externalAccountCredentialSource{ + // The cloud auth file provider reopens this path for every token exchange, + // allowing Kubernetes to rotate the projected token in place. + File: config.tokenFile, + Format: externalAccountCredentialSourceFormat{ + Type: "text", + }, + }, + } + credentialsJSON, err := json.Marshal(credentialConfig) + if err != nil { + return nil, fmt.Errorf("failed to marshal BigQuery workload identity credentials: %w", err) + } + return credentialsJSON, nil +} diff --git a/flow/connectors/bigquery/auth_test.go b/flow/connectors/bigquery/auth_test.go new file mode 100644 index 000000000..46265de64 --- /dev/null +++ b/flow/connectors/bigquery/auth_test.go @@ -0,0 +1,415 @@ +package connbigquery + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "testing" + "time" + + "cloud.google.com/go/auth/credentials" + "cloud.google.com/go/bigquery" + "github.com/stretchr/testify/require" + "google.golang.org/api/option" + "google.golang.org/protobuf/encoding/protojson" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" +) + +func TestBigQueryServiceAccountAuthTypeRemainsLegacy(t *testing.T) { + var config protos.BigqueryConfig + require.NoError(t, protojson.Unmarshal( + []byte(`{"authType":"service_account","projectId":"resource-project","datasetId":"dataset"}`), + &config, + )) + require.Equal(t, BigQueryAuthTypeServiceAccount, config.GetAuthType()) + + config.PrivateKeyId = "key-id" + config.PrivateKey = "private-key" + config.ClientEmail = "legacy@example.com" + config.ClientId = "client-id" + config.AuthUri = "https://accounts.google.com/o/oauth2/auth" + config.TokenUri = "https://oauth2.googleapis.com/token" + config.AuthProviderX509CertUrl = "https://www.googleapis.com/oauth2/v1/certs" + config.ClientX509CertUrl = "https://www.googleapis.com/robot/v1/metadata/x509/legacy" + + credentialConfig, err := newBigQueryCredentialConfig(t.Context(), &config, config.ProjectId) + require.NoError(t, err) + require.Equal(t, credentials.ServiceAccount, credentialConfig.credentialType) + require.Equal(t, bigquery.DetectProjectID, credentialConfig.clientProjectID) + + var credentialsDocument map[string]any + require.NoError(t, json.Unmarshal(credentialConfig.credentialsJSON, &credentialsDocument)) + require.Equal(t, "service_account", credentialsDocument["type"]) + require.Equal(t, "legacy@example.com", credentialsDocument["client_email"]) +} + +func TestBigQueryRejectsUnknownAuthType(t *testing.T) { + tests := []struct { + name string + authType string + }{ + {name: "empty"}, + {name: "unknown", authType: "unknown"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := newBigQueryCredentialConfig(t.Context(), &protos.BigqueryConfig{AuthType: test.authType}, "project") + require.ErrorContains(t, err, "unsupported BigQuery auth_type") + require.ErrorContains(t, err, fmt.Sprintf("%q", test.authType)) + }) + } +} + +func TestResolveBigQueryResource(t *testing.T) { + tests := []struct { + name string + projectID string + datasetID string + wantProjectID string + wantDatasetID string + wantErr string + }{ + { + name: "separate project and dataset", + projectID: "resource-project", + datasetID: "dataset", + wantProjectID: "resource-project", + wantDatasetID: "dataset", + }, + { + name: "project without default dataset", + projectID: "resource-project", + wantProjectID: "resource-project", + }, + { + name: "qualified dataset overrides project", + projectID: "configured-project", + datasetID: "qualified-project.dataset", + wantProjectID: "qualified-project", + wantDatasetID: "dataset", + }, + { + name: "three part dataset is rejected", + datasetID: "project.dataset.table", + wantErr: "invalid dataset ID", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + projectID, datasetID, err := resolveBigQueryResource(&protos.BigqueryConfig{ + ProjectId: test.projectID, + DatasetId: test.datasetID, + }) + if test.wantErr != "" { + require.ErrorContains(t, err, test.wantErr) + return + } + require.NoError(t, err) + require.Equal(t, test.wantProjectID, projectID) + require.Equal(t, test.wantDatasetID, datasetID) + }) + } +} + +func TestWorkloadIdentityCredentialsJSON(t *testing.T) { + config := workloadIdentityDeploymentConfig{ //nolint:gosec // Test-only fake account and token path. + targetServiceAccount: "tenant@tenant-project.iam.gserviceaccount.com", + tokenFile: "/var/run/secrets/peerdb/gcp-token", + projectID: "platform-project", + clusterLocation: "us-central1", + clusterName: "clickpipes", + } + + credentialsJSON, err := config.credentialsJSON() + require.NoError(t, err) + + var document externalAccountCredentials + require.NoError(t, json.Unmarshal(credentialsJSON, &document)) + require.Equal(t, "external_account", document.Type) + require.Equal(t, + "identitynamespace:platform-project.svc.id.goog:"+ + "https://container.googleapis.com/v1/projects/platform-project/locations/us-central1/clusters/clickpipes", + document.Audience, + ) + require.Equal(t, jwtSubjectTokenType, document.SubjectTokenType) + require.Equal(t, googleSTSTokenURL, document.TokenURL) + require.Equal(t, + "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/"+ + "tenant@tenant-project.iam.gserviceaccount.com:generateAccessToken", + document.ServiceAccountImpersonationURL, + ) + require.Equal(t, "/var/run/secrets/peerdb/gcp-token", document.CredentialSource.File) + require.Equal(t, "text", document.CredentialSource.Format.Type) + + _, err = credentials.NewCredentialsFromJSON(credentials.ExternalAccount, credentialsJSON, &credentials.DetectOptions{ + Scopes: []string{bigquery.Scope}, + }) + require.NoError(t, err) +} + +func TestWorkloadIdentityCredentialsReloadProjectedToken(t *testing.T) { + tokenFile := filepath.Join(t.TempDir(), "projected-token") + require.NoError(t, os.WriteFile(tokenFile, []byte("projected-token-one"), 0o600)) + + config := workloadIdentityDeploymentConfig{ + targetServiceAccount: "tenant@tenant-project.iam.gserviceaccount.com", + tokenFile: tokenFile, + projectID: "platform-project", + clusterLocation: "us-central1", + clusterName: "clickpipes", + } + credentialsJSON, err := config.credentialsJSON() + require.NoError(t, err) + + var stsSubjectTokens []string + var iamAuthorizationHeaders []string + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + require.Equal(t, "https", request.URL.Scheme) + switch request.URL.Host { + case "sts.googleapis.com": + require.Equal(t, http.MethodPost, request.Method) + require.Equal(t, "/v1/token", request.URL.Path) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + form, err := url.ParseQuery(string(body)) + require.NoError(t, err) + stsSubjectTokens = append(stsSubjectTokens, form.Get("subject_token")) + //nolint:gosec // Test-only fake STS response values. + return jsonResponse(t, request, map[string]any{ + "access_token": fmt.Sprintf("sts-access-token-%d", len(stsSubjectTokens)), + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "token_type": "Bearer", + "expires_in": 1, + }), nil + case "iamcredentials.googleapis.com": + require.Equal(t, http.MethodPost, request.Method) + require.Equal(t, + "/v1/projects/-/serviceAccounts/tenant@tenant-project.iam.gserviceaccount.com:generateAccessToken", + request.URL.Path, + ) + iamAuthorizationHeaders = append(iamAuthorizationHeaders, request.Header.Get("Authorization")) + return jsonResponse(t, request, map[string]any{ + "accessToken": fmt.Sprintf("tenant-access-token-%d", len(iamAuthorizationHeaders)), + "expireTime": time.Now().Add(time.Second).UTC().Format(time.RFC3339), + }), nil + default: + return nil, fmt.Errorf("unexpected request URL %s", request.URL) + } + }) + + creds, err := credentials.NewCredentialsFromJSON( + credentials.ExternalAccount, + credentialsJSON, + &credentials.DetectOptions{ + Scopes: []string{bigquery.Scope}, + Client: &http.Client{Transport: transport}, + }, + ) + require.NoError(t, err) + + firstToken, err := creds.Token(t.Context()) + require.NoError(t, err) + require.Equal(t, "tenant-access-token-1", firstToken.Value) + require.Equal(t, []string{"projected-token-one"}, stsSubjectTokens) + require.Equal(t, []string{"Bearer sts-access-token-1"}, iamAuthorizationHeaders) + + require.NoError(t, os.WriteFile(tokenFile, []byte("projected-token-two"), 0o600)) + time.Sleep(time.Until(firstToken.Expiry) + time.Second) + + secondToken, err := creds.Token(t.Context()) + require.NoError(t, err) + require.Equal(t, "tenant-access-token-2", secondToken.Value) + require.Equal(t, []string{"projected-token-one", "projected-token-two"}, stsSubjectTokens) + require.Equal(t, + []string{"Bearer sts-access-token-1", "Bearer sts-access-token-2"}, + iamAuthorizationHeaders, + ) +} + +func TestResolveWorkloadIdentityDeploymentConfig(t *testing.T) { + environment := map[string]string{ + workloadIdentityServiceAccountEnv: "tenant@tenant-project.iam.gserviceaccount.com", + workloadIdentityTokenFileEnv: "/token", + } + config, err := resolveWorkloadIdentityDeploymentConfig(t.Context(), workloadIdentityConfigSource{ + lookupEnv: func(name string) (string, bool) { + value, ok := environment[name] + return value, ok + }, + projectID: func(context.Context) (string, error) { + return "metadata-project", nil + }, + instanceAttributeValue: func(_ context.Context, name string) (string, error) { + return map[string]string{ + "cluster-location": "metadata-location", + "cluster-name": "metadata-cluster", + }[name], nil + }, + }) + require.NoError(t, err) + require.Equal(t, "metadata-project", config.projectID) + require.Equal(t, "metadata-location", config.clusterLocation) + require.Equal(t, "metadata-cluster", config.clusterName) +} + +func TestResolveWorkloadIdentityDeploymentConfigMissing(t *testing.T) { + tests := []struct { + name string + environment map[string]string + wantError string + }{ + { + name: "target service account", + wantError: workloadIdentityServiceAccountEnv, + }, + { + name: "token file", + environment: map[string]string{ + workloadIdentityServiceAccountEnv: "tenant@tenant-project.iam.gserviceaccount.com", + }, + wantError: workloadIdentityTokenFileEnv, + }, + { + name: "project context", + environment: map[string]string{ + workloadIdentityServiceAccountEnv: "tenant@tenant-project.iam.gserviceaccount.com", + workloadIdentityTokenFileEnv: "/token", + }, + wantError: workloadIdentityProjectIDEnv, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := resolveWorkloadIdentityDeploymentConfig(t.Context(), workloadIdentityConfigSource{ + lookupEnv: func(name string) (string, bool) { + value, ok := test.environment[name] + return value, ok + }, + projectID: func(context.Context) (string, error) { + return "", errors.New("metadata unavailable") + }, + instanceAttributeValue: func(context.Context, string) (string, error) { + return "", errors.New("metadata unavailable") + }, + }) + require.ErrorContains(t, err, test.wantError) + }) + } +} + +func TestWorkloadIdentityUsesExplicitResourceProject(t *testing.T) { + t.Setenv(workloadIdentityServiceAccountEnv, "tenant@tenant-project.iam.gserviceaccount.com") + t.Setenv(workloadIdentityTokenFileEnv, "/token") + t.Setenv(workloadIdentityProjectIDEnv, "platform-project") + t.Setenv(workloadIdentityClusterLocationEnv, "us-central1") + t.Setenv(workloadIdentityClusterNameEnv, "clickpipes") + + config := &protos.BigqueryConfig{ + ProjectId: "resource-project", + AuthType: BigQueryAuthTypeServiceAccountWorkloadIdentity, + } + projectID, datasetID, err := resolveBigQueryResource(config) + require.NoError(t, err) + require.Equal(t, "resource-project", projectID) + require.Empty(t, datasetID) + + credentialConfig, err := newBigQueryCredentialConfig(t.Context(), config, projectID) + require.NoError(t, err) + require.Equal(t, credentials.ExternalAccount, credentialConfig.credentialType) + require.Equal(t, "resource-project", credentialConfig.clientProjectID) +} + +func TestValidateBigQueryConnectionWithoutDefaultDataset(t *testing.T) { + requestPaths := make(chan string, 1) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + requestPaths <- request.URL.Path + response.Header().Set("Content-Type", "application/json") + _, _ = response.Write([]byte(`{"kind":"bigquery#datasetList","datasets":[]}`)) + })) + defer server.Close() + + client, err := bigquery.NewClient( + t.Context(), + "resource-project", + option.WithEndpoint(server.URL), + option.WithoutAuthentication(), + ) + require.NoError(t, err) + defer client.Close() + + require.NoError(t, validateBigQueryConnection(t.Context(), client, "resource-project", "")) + require.Equal(t, "/projects/resource-project/datasets", <-requestPaths) +} + +func TestValidateBigQueryConnectionWithoutDefaultDatasetReportsListingFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + http.Error(response, `{"error":{"code":403,"message":"permission denied"}}`, http.StatusForbidden) + })) + defer server.Close() + + client, err := bigquery.NewClient( + t.Context(), + "resource-project", + option.WithEndpoint(server.URL), + option.WithoutAuthentication(), + ) + require.NoError(t, err) + defer client.Close() + + err = validateBigQueryConnection(t.Context(), client, "resource-project", "") + require.ErrorContains(t, err, "failed to list BigQuery datasets") +} + +func TestWorkloadIdentityRequiresPeerProject(t *testing.T) { + config := &protos.BigqueryConfig{ + AuthType: BigQueryAuthTypeServiceAccountWorkloadIdentity, + } + _, err := newBigQueryCredentialConfig(t.Context(), config, "") + require.ErrorContains(t, err, "project ID must be set in the peer") +} + +func TestWorkloadIdentityCredentialConfigReportsMissingDeploymentConfig(t *testing.T) { + t.Setenv(workloadIdentityServiceAccountEnv, "") + t.Setenv(workloadIdentityTokenFileEnv, "") + t.Setenv(workloadIdentityProjectIDEnv, "") + t.Setenv(workloadIdentityClusterLocationEnv, "") + t.Setenv(workloadIdentityClusterNameEnv, "") + + config := &protos.BigqueryConfig{ + ProjectId: "resource-project", + AuthType: BigQueryAuthTypeServiceAccountWorkloadIdentity, + } + _, err := newBigQueryCredentialConfig(t.Context(), config, config.ProjectId) + require.ErrorContains(t, err, workloadIdentityServiceAccountEnv) +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (fn roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return fn(request) +} + +func jsonResponse(t *testing.T, request *http.Request, value any) *http.Response { + t.Helper() + body, err := json.Marshal(value) + require.NoError(t, err) + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + }, + Body: io.NopCloser(bytes.NewReader(body)), + Request: request, + } +} diff --git a/flow/connectors/bigquery/bigquery.go b/flow/connectors/bigquery/bigquery.go index 0642b1ff3..2225fb86e 100644 --- a/flow/connectors/bigquery/bigquery.go +++ b/flow/connectors/bigquery/bigquery.go @@ -2,7 +2,6 @@ package connbigquery import ( "context" - "encoding/json" "errors" "fmt" "log/slog" @@ -67,51 +66,50 @@ type BigQueryConnector struct { func NewBigQueryConnector(ctx context.Context, config *protos.BigqueryConfig) (*BigQueryConnector, error) { logger := internal.LoggerFromCtx(ctx) - datasetID := config.GetDatasetId() - projectID := config.GetProjectId() - projectPart, datasetPart, found := strings.Cut(datasetID, ".") - if found && strings.Contains(datasetPart, ".") { - return nil, - fmt.Errorf("invalid dataset ID: %s. Ensure that it is just a single string or string1.string2", datasetID) - } - if projectPart != "" && datasetPart != "" { - datasetID = datasetPart - projectID = projectPart - } - - serviceAccount, err := NewBigQueryServiceAccount(config) + projectID, datasetID, err := resolveBigQueryResource(config) if err != nil { - return nil, fmt.Errorf("failed to create BigQueryServiceAccount: %w", err) + return nil, err } - saJSON, err := json.Marshal(serviceAccount) //nolint:gosec // G117: credential struct marshaled for inline use + credentialConfig, err := newBigQueryCredentialConfig(ctx, config, projectID) if err != nil { - return nil, fmt.Errorf("failed to marshal service account: %v", err) + return nil, err } - creds, err := credentials.DetectDefault(&credentials.DetectOptions{ - CredentialsJSON: saJSON, + detectOptions := &credentials.DetectOptions{ Scopes: []string{ bigquery.Scope, storage.ScopeFullControl, // we should split it into two clients later }, - }) + } + var creds *auth.Credentials + if credentialConfig.credentialType == credentials.ExternalAccount { + creds, err = credentials.NewCredentialsFromJSON( + credentials.ExternalAccount, + credentialConfig.credentialsJSON, + detectOptions, + ) + } else { + // Keep the legacy service-account-key path unchanged. + detectOptions.CredentialsJSON = credentialConfig.credentialsJSON //nolint:staticcheck // Preserve legacy service-account-key behavior. + creds, err = credentials.DetectDefault(detectOptions) + } if err != nil { return nil, fmt.Errorf("failed to create credentials: %v", err) } client, err := bigquery.NewClient( ctx, - bigquery.DetectProjectID, + credentialConfig.clientProjectID, option.WithAuthCredentials(creds), ) if err != nil { return nil, fmt.Errorf("failed to create BigQuery client: %v", err) } - if _, err := client.DatasetInProject(projectID, datasetID).Metadata(ctx); err != nil { - logger.Error("failed to get dataset metadata", slog.Any("error", err)) - return nil, fmt.Errorf("failed to get dataset metadata: %v", err) + if err := validateBigQueryConnection(ctx, client, projectID, datasetID); err != nil { + logger.Error("failed to validate BigQuery connection", slog.Any("error", err)) + return nil, err } storageClient, err := storage.NewClient(ctx, option.WithAuthCredentials(creds)) @@ -138,11 +136,7 @@ func NewBigQueryConnector(ctx context.Context, config *protos.BigqueryConfig) (* } func (c *BigQueryConnector) ValidateCheck(ctx context.Context) error { - if _, err := c.client.DatasetInProject(c.projectID, c.datasetID).Metadata(ctx); err != nil { - return fmt.Errorf("failed to get dataset metadata: %v", err) - } - - return nil + return validateBigQueryConnection(ctx, c.client, c.projectID, c.datasetID) } func (c *BigQueryConnector) ValidateMirrorDestination( @@ -200,10 +194,25 @@ func (c *BigQueryConnector) Close() error { // ConnectionActive returns nil if the connection is active. func (c *BigQueryConnector) ConnectionActive(ctx context.Context) error { - if _, err := c.client.DatasetInProject(c.projectID, c.datasetID).Metadata(ctx); err != nil { - return fmt.Errorf("failed to get dataset metadata: %v", err) - } + return validateBigQueryConnection(ctx, c.client, c.projectID, c.datasetID) +} +func validateBigQueryConnection( + ctx context.Context, + client *bigquery.Client, + projectID, datasetID string, +) error { + if datasetID == "" { + datasets := client.Datasets(ctx) + datasets.ProjectID = projectID + if _, err := datasets.Next(); err != nil && !errors.Is(err, iterator.Done) { + return fmt.Errorf("failed to list BigQuery datasets: %w", err) + } + return nil + } + if _, err := client.DatasetInProject(projectID, datasetID).Metadata(ctx); err != nil { + return fmt.Errorf("failed to get dataset metadata: %w", err) + } return nil } diff --git a/flow/go.mod b/flow/go.mod index a0e243da2..a0f334e52 100644 --- a/flow/go.mod +++ b/flow/go.mod @@ -6,6 +6,7 @@ require ( cloud.google.com/go v0.123.0 cloud.google.com/go/auth v0.20.0 cloud.google.com/go/bigquery v1.74.0 + cloud.google.com/go/compute/metadata v0.9.0 cloud.google.com/go/kms v1.31.0 cloud.google.com/go/pubsub/v2 v2.3.0 cloud.google.com/go/storage v1.62.1 @@ -101,7 +102,6 @@ require ( require ( cel.dev/expr v0.25.1 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect - cloud.google.com/go/compute/metadata v0.9.0 // indirect cloud.google.com/go/iam v1.7.0 // indirect cloud.google.com/go/longrunning v0.9.0 // indirect cloud.google.com/go/monitoring v1.24.3 // indirect