diff --git a/go/bigquery_database.go b/go/bigquery_database.go index 62814c7..aacea0b 100644 --- a/go/bigquery_database.go +++ b/go/bigquery_database.go @@ -25,6 +25,8 @@ package bigquery import ( "context" "fmt" + "net/url" + "strconv" "strings" "time" @@ -53,6 +55,7 @@ type databaseImpl struct { tableID string location string quotaProject string + endpoint string } func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) { @@ -71,6 +74,7 @@ func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) { catalog: d.projectID, dbSchema: d.datasetID, location: d.location, + endpoint: d.endpoint, resultRecordBufferSize: defaultQueryResultBufferSize, prefetchConcurrency: defaultQueryPrefetchConcurrency, quotaProject: d.quotaProject, @@ -113,6 +117,8 @@ func (d *databaseImpl) GetOption(key string) (string, error) { return d.datasetID, nil case OptionStringTableID: return d.tableID, nil + case OptionStringEndpoint: + return d.endpoint, nil case OptionStringImpersonateLifetime: if d.impersonateLifetime == 0 { // If no lifetime is set but impersonation is enabled, return the default @@ -145,6 +151,18 @@ func (d *databaseImpl) hasImpersonationOptions() bool { func (d *databaseImpl) SetOption(key string, value string) error { switch key { + case "uri": + params, err := ParseBigQueryURIToParams(value) + if err != nil { + return err + } + + for paramKey, paramValue := range params { + if err := d.SetOption(paramKey, paramValue); err != nil { + return err + } + } + return nil case OptionStringAuthType: switch value { case OptionValueAuthTypeDefault, @@ -190,6 +208,8 @@ func (d *databaseImpl) SetOption(key string, value string) error { d.datasetID = value case OptionStringTableID: d.tableID = value + case OptionStringEndpoint: + d.endpoint = value case OptionStringLocation: d.location = value default: @@ -197,3 +217,176 @@ func (d *databaseImpl) SetOption(key string, value string) error { } return nil } + +// ParseBigQueryURIToParams parses a BigQuery URI and returns the extracted parameters +func ParseBigQueryURIToParams(uri string) (map[string]string, error) { + if uri == "" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: "[bq] URI cannot be empty", + } + } + + parsedURI, err := url.Parse(uri) + if err != nil { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid BigQuery URI format: %v", err), + } + } + + if parsedURI.Scheme != "bigquery" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid BigQuery URI scheme: expected 'bigquery', got '%s'", parsedURI.Scheme), + } + } + + projectID := strings.TrimPrefix(parsedURI.Path, "/") + if projectID == "" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: "[bq] project ID is required in URI path", + } + } + + params := make(map[string]string) + params[OptionStringProjectID] = projectID + + // Handle host and port + var endpoint string + if parsedURI.Host != "" && parsedURI.Hostname() != "" { + // Custom endpoint specified with valid hostname + if parsedURI.Port() != "" { + endpoint = parsedURI.Host + } else { + endpoint = fmt.Sprintf("%s:443", parsedURI.Hostname()) + } + } else if parsedURI.Host != "" && parsedURI.Hostname() == "" && parsedURI.Port() != "" { + // Port without hostname. use default host with custom port + endpoint = fmt.Sprintf("bigquery.googleapis.com:%s", parsedURI.Port()) + } else { + // No host specified, use default BigQuery endpoint + endpoint = "bigquery.googleapis.com:443" + } + + // Store endpoint as hostname:port (Google client library handles https:// internally) + params[OptionStringEndpoint] = endpoint + + queryParams := parsedURI.Query() + + oauthTypeStr := queryParams.Get("OAuthType") + if oauthTypeStr != "" { + oauthType, err := strconv.Atoi(oauthTypeStr) + if err != nil { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid OAuthType value: %s", oauthTypeStr), + } + } + + switch oauthType { + case 0: + params[OptionStringAuthType] = OptionValueAuthTypeAppDefaultCredentials + case 1: + params[OptionStringAuthType] = OptionValueAuthTypeJSONCredentialFile + + authCredentials := queryParams.Get("AuthCredentials") + if authCredentials == "" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: "[bq] AuthCredentials required for service account authentication", + } + } + + decodedCreds, err := url.QueryUnescape(authCredentials) + if err != nil { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid AuthCredentials format: %v", err), + } + } + params[OptionStringAuthCredentials] = decodedCreds + queryParams.Del("AuthCredentials") + case 2: + params[OptionStringAuthType] = OptionValueAuthTypeJSONCredentialString + + authCredentials := queryParams.Get("AuthCredentials") + if authCredentials == "" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: "[bq] AuthCredentials required for service account authentication", + } + } + + decodedCreds, err := url.QueryUnescape(authCredentials) + if err != nil { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid AuthCredentials format: %v", err), + } + } + params[OptionStringAuthCredentials] = decodedCreds + queryParams.Del("AuthCredentials") + case 3: + params[OptionStringAuthType] = OptionValueAuthTypeUserAuthentication + + clientID := queryParams.Get("AuthClientId") + clientSecret := queryParams.Get("AuthClientSecret") + refreshToken := queryParams.Get("AuthRefreshToken") + if clientID == "" || clientSecret == "" || refreshToken == "" { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: "[bq] AuthClientId, AuthClientSecret and AuthRefreshToken required for OAuth authentication", + } + } + params[OptionStringAuthClientID] = clientID + params[OptionStringAuthClientSecret] = clientSecret + params[OptionStringAuthRefreshToken] = refreshToken + queryParams.Del("AuthClientId") + queryParams.Del("AuthClientSecret") + queryParams.Del("AuthRefreshToken") + default: + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] invalid OAuthType value: %d", oauthType), + } + } + queryParams.Del("OAuthType") + } else { + // if not provided default to ADC + params[OptionStringAuthType] = OptionValueAuthTypeAppDefaultCredentials + } + + parameterMap := map[string]string{ + "DatasetId": OptionStringDatasetID, + "Location": OptionStringLocation, + "TableId": OptionStringTableID, + "QuotaProject": OptionStringAuthQuotaProject, + + // Auth parameters - processed in OAuthType switch above, here for consistency + "AuthCredentials": OptionStringAuthCredentials, + "AuthClientId": OptionStringAuthClientID, + "AuthClientSecret": OptionStringAuthClientSecret, + "AuthRefreshToken": OptionStringAuthRefreshToken, + + "ImpersonateTargetPrincipal": OptionStringImpersonateTargetPrincipal, + "ImpersonateDelegates": OptionStringImpersonateDelegates, + "ImpersonateScopes": OptionStringImpersonateScopes, + "ImpersonateLifetime": OptionStringImpersonateLifetime, + } + + // Process all query parameters to convert URI params to option constants + for paramName, paramValues := range queryParams { + if optionName, exists := parameterMap[paramName]; exists { + params[optionName] = paramValues[0] + } else { + return nil, adbc.Error{ + Code: adbc.StatusInvalidArgument, + Msg: fmt.Sprintf("[bq] unknown parameter '%s' in URI", paramName), + } + } + } + + return params, nil +} diff --git a/go/connection.go b/go/connection.go index eb2f7ad..9db937c 100644 --- a/go/connection.go +++ b/go/connection.go @@ -69,6 +69,8 @@ type connectionImpl struct { dbSchema string // tableID is the default table for statement tableID string + // endpoint is the custom BigQuery API endpoint + endpoint string sessionID *string @@ -730,7 +732,13 @@ func (c *connectionImpl) newClient(ctx context.Context) error { authOptions = []option.ClientOption{option.WithTokenSource(tokenSource)} } - client, err := bigquery.NewClient(ctx, c.catalog, authOptions...) + // Add custom endpoint if specified for BigQuery API client + bigQueryAuthOptions := authOptions + if c.endpoint != "" { + bigQueryAuthOptions = append(bigQueryAuthOptions, option.WithEndpoint(c.endpoint)) + } + + client, err := bigquery.NewClient(ctx, c.catalog, bigQueryAuthOptions...) if err != nil { return errToAdbcErr(adbc.StatusIO, err, "create client") } @@ -739,6 +747,7 @@ func (c *connectionImpl) newClient(ctx context.Context) error { client.Location = c.location } + // Use original authOptions without custom endpoint for Storage Read API err = client.EnableStorageReadClient(ctx, authOptions...) if err != nil { return errToAdbcErr(adbc.StatusIO, err, "enable storage read client") diff --git a/go/driver.go b/go/driver.go index 3be5709..12a4e4c 100644 --- a/go/driver.go +++ b/go/driver.go @@ -40,6 +40,7 @@ const ( OptionStringProjectID = "adbc.bigquery.sql.project_id" OptionStringDatasetID = "adbc.bigquery.sql.dataset_id" OptionStringTableID = "adbc.bigquery.sql.table_id" + OptionStringEndpoint = "adbc.bigquery.sql.endpoint" OptionValueAuthTypeDefault = "adbc.bigquery.sql.auth_type.auth_bigquery" diff --git a/go/driver_test.go b/go/driver_test.go index 432df13..dff7aa3 100644 --- a/go/driver_test.go +++ b/go/driver_test.go @@ -45,6 +45,8 @@ import ( "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/pqarrow" "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" ) @@ -1658,3 +1660,315 @@ func (s *BigQueryTestSuite) TearDownSuite() { } s.mem.AssertSize(s.T(), 0) } + +func TestBigQueryURIParsing(t *testing.T) { + tests := []struct { + name string + uri string + expectedProjectID string + expectedDatasetID string + expectedAuthType string + expectedCredentials string + expectedClientID string + expectedClientSecret string + expectedRefreshToken string + expectedLocation string + expectedEndpoint string + expectedQuotaProject string + expectedTableID string + expectedImpersonateTarget string + expectedImpersonateLifetime string + expectedImpersonateDelegates string + expectedImpersonateScopes string + shouldError bool + errorContains string + }{ + { + name: "application default credentials basic", + uri: "bigquery:///my-project-123?OAuthType=0", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + }, + { + name: "adc with dataset and location", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test_dataset&Location=EU", + expectedProjectID: "my-project-123", + expectedDatasetID: "test_dataset", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedLocation: "EU", + }, + { + name: "service account json file", + uri: "bigquery://bigquery.googleapis.com/my-project-123?OAuthType=1&AuthCredentials=/path/to/key.json", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeJSONCredentialFile, + expectedCredentials: "/path/to/key.json", + }, + { + name: "service account with dataset", + uri: "bigquery:///my-project-123?OAuthType=1&AuthCredentials=/path/to/key.json&DatasetId=prod", + expectedProjectID: "my-project-123", + expectedDatasetID: "prod", + expectedAuthType: driver.OptionValueAuthTypeJSONCredentialFile, + expectedCredentials: "/path/to/key.json", + }, + { + name: "service account json string", + uri: "bigquery:///my-project-123?OAuthType=2&AuthCredentials=%7B%22type%22%3A%22service_account%22%7D", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeJSONCredentialString, + expectedCredentials: `{"type":"service_account"}`, + }, + { + name: "user oauth authentication", + uri: "bigquery:///my-project-123?OAuthType=3&AuthClientId=client_id&AuthClientSecret=secret&AuthRefreshToken=token", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeUserAuthentication, + expectedClientID: "client_id", + expectedClientSecret: "secret", + expectedRefreshToken: "token", + }, + { + name: "minimal uri with default auth", + uri: "bigquery:///my-project-123", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + }, + { + name: "default auth with dataset", + uri: "bigquery:///my-project-123?DatasetId=analytics", + expectedProjectID: "my-project-123", + expectedDatasetID: "analytics", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + }, + { + name: "default auth with multiple params", + uri: "bigquery:///my-project-123?DatasetId=test&Location=EU", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedLocation: "EU", + }, + { + name: "custom endpoint with port", + uri: "bigquery://custom-endpoint.googleapis.com:443/my-project-123?OAuthType=0", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedEndpoint: "custom-endpoint.googleapis.com:443", + }, + { + name: "custom endpoint without port", + uri: "bigquery://localhost/my-project-123?OAuthType=0", + expectedProjectID: "my-project-123", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedEndpoint: "localhost:443", + }, + { + name: "bigquery googleapis endpoint", + uri: "bigquery://bigquery.googleapis.com/my-project-123?OAuthType=0&DatasetId=test", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedEndpoint: "bigquery.googleapis.com:443", + }, + { + name: "default endpoint no host specified", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedEndpoint: "bigquery.googleapis.com:443", + }, + { + name: "multiple parameters", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=analytics&Location=US", + expectedProjectID: "my-project-123", + expectedDatasetID: "analytics", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedLocation: "US", + }, + { + name: "quota project parameter", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&QuotaProject=billing-project", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedQuotaProject: "billing-project", + }, + { + name: "table id parameter", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&TableId=my_table", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedTableID: "my_table", + }, + { + name: "impersonation parameters", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&ImpersonateTargetPrincipal=service@example.com&ImpersonateLifetime=3600s", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedImpersonateTarget: "service@example.com", + expectedImpersonateLifetime: "3600s", + }, + { + name: "impersonation delegates parameter", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&ImpersonateDelegates=delegate1@example.com,delegate2@example.com", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedImpersonateDelegates: "delegate1@example.com,delegate2@example.com", + }, + { + name: "impersonation scopes parameter", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&ImpersonateScopes=https://www.googleapis.com/auth/bigquery,https://www.googleapis.com/auth/cloud-platform", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedImpersonateScopes: "https://www.googleapis.com/auth/bigquery,https://www.googleapis.com/auth/cloud-platform", + }, + { + name: "all impersonation parameters", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=test&ImpersonateTargetPrincipal=service@example.com&ImpersonateDelegates=delegate@example.com&ImpersonateScopes=https://www.googleapis.com/auth/bigquery&ImpersonateLifetime=1800s", + expectedProjectID: "my-project-123", + expectedDatasetID: "test", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedImpersonateTarget: "service@example.com", + expectedImpersonateDelegates: "delegate@example.com", + expectedImpersonateScopes: "https://www.googleapis.com/auth/bigquery", + expectedImpersonateLifetime: "1800s", + }, + { + name: "all optional parameters", + uri: "bigquery:///my-project-123?OAuthType=0&DatasetId=analytics&Location=EU&QuotaProject=billing&TableId=orders&ImpersonateTargetPrincipal=svc@example.com", + expectedProjectID: "my-project-123", + expectedDatasetID: "analytics", + expectedAuthType: driver.OptionValueAuthTypeAppDefaultCredentials, + expectedLocation: "EU", + expectedQuotaProject: "billing", + expectedTableID: "orders", + expectedImpersonateTarget: "svc@example.com", + }, + { + name: "missing project id", + uri: "bigquery:///?OAuthType=0", + shouldError: true, + errorContains: "project ID is required in URI path", + }, + { + name: "invalid oauth type", + uri: "bigquery:///my-project-123?OAuthType=999", + shouldError: true, + errorContains: "invalid OAuthType value", + }, + { + name: "missing auth credentials for service account file", + uri: "bigquery:///my-project-123?OAuthType=1", + shouldError: true, + errorContains: "AuthCredentials required for service account authentication", + }, + { + name: "missing auth credentials for service account string", + uri: "bigquery:///my-project-123?OAuthType=2", + shouldError: true, + errorContains: "AuthCredentials required for service account authentication", + }, + { + name: "missing oauth credentials", + uri: "bigquery:///my-project-123?OAuthType=3&AuthClientId=client", + shouldError: true, + errorContains: "AuthClientId, AuthClientSecret and AuthRefreshToken required for OAuth authentication", + }, + { + name: "invalid scheme", + uri: "mysql://localhost/db", + shouldError: true, + errorContains: "invalid BigQuery URI scheme", + }, + { + name: "malformed uri", + uri: "bigquery://[invalid", + shouldError: true, + errorContains: "invalid BigQuery URI format", + }, + { + name: "empty project id", + uri: "bigquery:///?OAuthType=0", + shouldError: true, + errorContains: "project ID is required in URI path", + }, + { + name: "unknown parameter", + uri: "bigquery:///my-project-123?OAuthType=0&UnknownParam=value", + shouldError: true, + errorContains: "unknown parameter 'UnknownParam' in URI", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + params, err := driver.ParseBigQueryURIToParams(tt.uri) + + if tt.shouldError { + require.ErrorContains(t, err, tt.errorContains) + return + } + + require.NoError(t, err, "unexpected error during URI parsing") + + if tt.expectedProjectID != "" { + assert.Equal(t, tt.expectedProjectID, params[driver.OptionStringProjectID], "project ID mismatch") + } + + if tt.expectedDatasetID != "" { + assert.Equal(t, tt.expectedDatasetID, params[driver.OptionStringDatasetID], "dataset ID mismatch") + } + + if tt.expectedAuthType != "" { + assert.Equal(t, tt.expectedAuthType, params[driver.OptionStringAuthType], "auth type mismatch") + } + + if tt.expectedCredentials != "" { + assert.Equal(t, tt.expectedCredentials, params[driver.OptionStringAuthCredentials], "credentials mismatch") + } + + if tt.expectedClientID != "" { + assert.Equal(t, tt.expectedClientID, params[driver.OptionStringAuthClientID], "client ID mismatch") + } + + if tt.expectedClientSecret != "" { + assert.Equal(t, tt.expectedClientSecret, params[driver.OptionStringAuthClientSecret], "client secret mismatch") + } + + if tt.expectedRefreshToken != "" { + assert.Equal(t, tt.expectedRefreshToken, params[driver.OptionStringAuthRefreshToken], "refresh token mismatch") + } + + if tt.expectedLocation != "" { + assert.Equal(t, tt.expectedLocation, params[driver.OptionStringLocation], "location mismatch") + } + if tt.expectedEndpoint != "" { + assert.Equal(t, tt.expectedEndpoint, params[driver.OptionStringEndpoint], "endpoint mismatch") + } + if tt.expectedQuotaProject != "" { + assert.Equal(t, tt.expectedQuotaProject, params[driver.OptionStringAuthQuotaProject], "quota project mismatch") + } + if tt.expectedTableID != "" { + assert.Equal(t, tt.expectedTableID, params[driver.OptionStringTableID], "table ID mismatch") + } + if tt.expectedImpersonateTarget != "" { + assert.Equal(t, tt.expectedImpersonateTarget, params[driver.OptionStringImpersonateTargetPrincipal], "impersonate target mismatch") + } + if tt.expectedImpersonateLifetime != "" { + assert.Equal(t, tt.expectedImpersonateLifetime, params[driver.OptionStringImpersonateLifetime], "impersonate lifetime mismatch") + } + if tt.expectedImpersonateDelegates != "" { + assert.Equal(t, tt.expectedImpersonateDelegates, params[driver.OptionStringImpersonateDelegates], "impersonate delegates mismatch") + } + if tt.expectedImpersonateScopes != "" { + assert.Equal(t, tt.expectedImpersonateScopes, params[driver.OptionStringImpersonateScopes], "impersonate scopes mismatch") + } + }) + } +} diff --git a/go/validation/tests/bigquery/conftest.py b/go/validation/tests/bigquery/conftest.py new file mode 100644 index 0000000..93787df --- /dev/null +++ b/go/validation/tests/bigquery/conftest.py @@ -0,0 +1,44 @@ +# Copyright (c) 2025 ADBC Drivers Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import pytest + + +def pytest_generate_tests(metafunc) -> None: + metafunc.parametrize( + "driver", + [pytest.param("bigquery:", id="bigquery")], + scope="module", + indirect=["driver"], + ) + + +@pytest.fixture(scope="session") +def bigquery_project() -> str: + """BigQuery project ID. Example: GOOGLE_CLOUD_PROJECT=my-project-123""" + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + if not project: + pytest.skip("Must set GOOGLE_CLOUD_PROJECT environment variable") + return project + + +@pytest.fixture(scope="session") +def bigquery_dataset() -> str: + """BigQuery dataset ID. Example: BIGQUERY_DATASET_ID=test_dataset""" + dataset = os.environ.get("BIGQUERY_DATASET_ID") + if not dataset: + pytest.skip("Must set BIGQUERY_DATASET_ID environment variable") + return dataset diff --git a/go/validation/tests/bigquery/test_uri.py b/go/validation/tests/bigquery/test_uri.py new file mode 100644 index 0000000..f9533da --- /dev/null +++ b/go/validation/tests/bigquery/test_uri.py @@ -0,0 +1,342 @@ +# Copyright (c) 2025 ADBC Drivers Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import urllib.parse + +import adbc_driver_manager.dbapi +import pytest + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_application_default_credentials_uri_parsing( + driver_path: str, +) -> None: + """Test that Application Default Credentials URI is parsed correctly.""" + uri = "bigquery:///dummyproject?OAuthType=0&DatasetId=dummydataset&Location=dummylocation" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.app_default_credentials" + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + location = db.get_option("adbc.bigquery.sql.location") + assert location == "dummylocation" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_service_account_file_uri_parsing( + driver_path: str, +) -> None: + """Test that service account JSON file URI is parsed correctly.""" + credentials_path = "/path/to/service-account.json" + uri = f"bigquery:///dummyproject?OAuthType=1&AuthCredentials={credentials_path}&DatasetId=dummydataset&TableId=mytable" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.json_credential_file" + + computed_credentials = db.get_option("adbc.bigquery.sql.auth_credentials") + assert computed_credentials == credentials_path + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + table_id = db.get_option("adbc.bigquery.sql.table_id") + assert table_id == "mytable" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_service_account_string_uri_parsing( + driver_path: str, +) -> None: + """Test that service account JSON string URI is parsed correctly.""" + json_credentials = '{"type":"service_account","project_id":"test"}' + encoded_credentials = urllib.parse.quote(json_credentials) + uri = f"bigquery:///dummyproject?OAuthType=2&AuthCredentials={encoded_credentials}&DatasetId=dummydataset&Location=dummylocation" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.json_credential_string" + + computed_credentials = db.get_option("adbc.bigquery.sql.auth_credentials") + assert computed_credentials == json_credentials + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + location = db.get_option("adbc.bigquery.sql.location") + assert location == "dummylocation" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_user_oauth_uri_parsing( + driver_path: str, +) -> None: + """Test that User OAuth authentication URI is parsed correctly.""" + client_id = "test_client_id" + client_secret = "test_client_secret" + refresh_token = "test_refresh_token" + uri = f"bigquery:///dummyproject?OAuthType=3&AuthClientId={client_id}&AuthClientSecret={client_secret}&AuthRefreshToken={refresh_token}&DatasetId=dummydataset&Location=dummylocation" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.user_authentication" + + endpoint = db.get_option("adbc.bigquery.sql.endpoint") + assert endpoint == "bigquery.googleapis.com:443" + + computed_client_id = db.get_option("adbc.bigquery.sql.auth.client_id") + assert computed_client_id == client_id + + computed_client_secret = db.get_option("adbc.bigquery.sql.auth.client_secret") + assert computed_client_secret == client_secret + + computed_refresh_token = db.get_option("adbc.bigquery.sql.auth.refresh_token") + assert computed_refresh_token == refresh_token + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + location = db.get_option("adbc.bigquery.sql.location") + assert location == "dummylocation" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_custom_endpoint_uri_parsing( + driver_path: str, +) -> None: + """Test that custom endpoint in URI is parsed correctly.""" + uri = "bigquery://bigquery.dummyapis.com:445/dummyproject?DatasetId=dummydataset&QuotaProject=billing-project" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.app_default_credentials" + + endpoint = db.get_option("adbc.bigquery.sql.endpoint") + assert endpoint == "bigquery.dummyapis.com:445" + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + quota_project = db.get_option("adbc.bigquery.sql.auth.quota_project") + assert quota_project == "billing-project" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_custom_endpoint_without_port_uri_parsing( + driver_path: str, +) -> None: + """Test that custom endpoint without port defaults to 443.""" + uri = "bigquery://bigquery.dummyapis.com/dummyproject?OAuthType=0&DatasetId=dummydataset" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.app_default_credentials" + + endpoint = db.get_option("adbc.bigquery.sql.endpoint") + assert endpoint == "bigquery.dummyapis.com:443" + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + dataset_id = db.get_option("adbc.bigquery.sql.dataset_id") + assert dataset_id == "dummydataset" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_port_without_host_uri_parsing( + driver_path: str, +) -> None: + """Test that port without host in URI defaults to custom port on the default host.""" + uri = "bigquery://:448/dummyproject?OAuthType=0" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.app_default_credentials" + + endpoint = db.get_option("adbc.bigquery.sql.endpoint") + assert endpoint == "bigquery.googleapis.com:448" + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_missing_project_id_uri_error( + driver_path: str, +) -> None: + """Test that missing project ID in URI raises error during parsing.""" + uri = "bigquery:///?OAuthType=0" + + with pytest.raises( + adbc_driver_manager.dbapi.ProgrammingError, + match="project ID is required in URI path", + ): + with adbc_driver_manager.dbapi.connect( + driver=driver_path, + db_kwargs={"uri": uri}, + ): + pass + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_minimal_uri_parsing( + driver_path: str, +) -> None: + """Test that minimal URI (project only) works with default authentication.""" + uri = "bigquery:///dummyproject" + + params = { + "uri": uri, + } + + with adbc_driver_manager.AdbcDatabase(driver=driver_path, **params) as db: + auth_type = db.get_option("adbc.bigquery.sql.auth_type") + assert auth_type == "adbc.bigquery.sql.auth_type.app_default_credentials" + + project_id = db.get_option("adbc.bigquery.sql.project_id") + assert project_id == "dummyproject" + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_invalid_oauth_type_uri_error( + driver_path: str, +) -> None: + """Test that invalid OAuthType in URI raises error during parsing.""" + uri = "bigquery:///my-project-123?OAuthType=999" + + with pytest.raises( + adbc_driver_manager.dbapi.ProgrammingError, + match="invalid OAuthType value", + ): + with adbc_driver_manager.dbapi.connect( + driver=driver_path, + db_kwargs={"uri": uri}, + ): + pass + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_invalid_uri_scheme_error( + driver_path: str, +) -> None: + """Test that invalid URI scheme raises error during parsing.""" + uri = "mysql://localhost/database" + + with pytest.raises( + adbc_driver_manager.dbapi.ProgrammingError, + match="invalid BigQuery URI scheme", + ): + with adbc_driver_manager.dbapi.connect( + driver=driver_path, + db_kwargs={"uri": uri}, + ): + pass + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_missing_auth_credentials_service_account_uri_error( + driver_path: str, +) -> None: + """Test that missing AuthCredentials for service account raises error during parsing.""" + uri = "bigquery:///my-project-123?OAuthType=1" + + with pytest.raises( + adbc_driver_manager.dbapi.ProgrammingError, + match="AuthCredentials required for service account authentication", + ): + with adbc_driver_manager.dbapi.connect( + driver=driver_path, + db_kwargs={"uri": uri}, + ): + pass + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_missing_oauth_credentials_uri_error( + driver_path: str, +) -> None: + """Test that incomplete OAuth credentials in URI raise error during parsing.""" + uri = "bigquery:///my-project-123?OAuthType=3&AuthClientId=client_id" + + with pytest.raises( + adbc_driver_manager.dbapi.ProgrammingError, + match="AuthClientId, AuthClientSecret and AuthRefreshToken required for OAuth authentication", + ): + with adbc_driver_manager.dbapi.connect( + driver=driver_path, + db_kwargs={"uri": uri}, + ): + pass + + +@pytest.mark.feature(group="Configuration", name="Connect with URI") +def test_end_to_end_real_connection( + driver_path: str, + bigquery_project: str, + bigquery_dataset: str, +) -> None: + """Test end-to-end connection with real BigQuery project.""" + uri = f"bigquery://bigquery.googleapis.com/{bigquery_project}?DatasetId={bigquery_dataset}" + + with adbc_driver_manager.dbapi.connect( + driver=driver_path, db_kwargs={"uri": uri}, autocommit=True + ) as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT 1")