Port over the commits from the legacy bigquery go driver, rebase the latest upstream main - #1
Conversation
strings.Split("", ",") yields [""], so an empty delegates or scopes option
was stored as a single empty entry rather than none. That reads as
"impersonation requested" downstream and, for scopes, passes an empty
scope to option.WithScopes.
Two deliberate deviations from the upstream commit:
- The guard is applied on databaseImpl as well as connectionImpl. Open()
copies these fields into the connection, so a connection-only guard is
defeated when the option is set at the database level — which is the
path dbt-auth uses.
- The upstream commit also drops impersonateScopes from
hasImpersonationOptions(). That is already covered here by the port of
adbc-drivers#75; upstream had reverted it in adbc-drivers#87 ("Merge with ADBC 21") and adbc-drivers#91
restored it on connectionImpl only.
Ports arrow-adbc commit 4079e1747 (Fix impersonation scopes not
respecting empty values, adbc-drivers#91).
…auth Adds three related auth features from the legacy dbt-labs/arrow-adbc bigquery driver, needed by the fs dbt-auth flow: - OptionValueAuthTypeTemporaryAccessToken auth type + OptionStringAuthAccessToken: fed with a short-lived bearer token (the dbt platform authenticator supplies these). Uses a static oauth2.TokenSource. - OptionStringAuthAccessTokenEndpoint / ...ServerName: overrides the OAuth token endpoint (default https://accounts.google.com/o/oauth2/token) and TLS ServerName used when refreshing user OAuth credentials. Needed to route token exchange through internal proxies. Ports from arrow-adbc commits 6a57518fe (initial dbt fork - adapter driver interface) and 5cebc7d1a (Make OAuth access token endpoint configurable, #5). rebase the latest upstream main, resolve conflicts; add TODOs to migrate the options to using the new format
Accepts both the legacy option name (adbc.bigquery.sql.api_endpoint) and the new one (adbc.bigquery.sql.endpoint) at the database level, so callers that still use the legacy constant from dbt-labs/arrow-adbc work unchanged. Ports arrow-adbc commit 1850d6103 (pipe api_endpoint through to Bigquery NewClient, adbc-drivers#120). rebase the latest upstream main, resolve conflicts: remove OptionStringTableID which is unused, keep OptionStringAPIEndpoint, add TODOs to migrate it to the newer format
Adds OptionStringQueryLabels ("adbc.bigquery.sql.query.labels") that
accepts a JSON object mapping label name to value; the map is set on
the QueryConfig.Labels field and thus attached to the resulting
BigQuery query job.
Used by fs/sa/crates/dbt-adapter/src/engine/adapter_engine.rs to
propagate dbt run metadata (invocation_id, model name, …) to BigQuery
so operators can filter jobs in the console.
Ports arrow-adbc commit 86e382780 (Add job labels to statement
options, adbc-drivers#76).
rebase the latest upstream main, resolve conflicts: add TODOs to migrate
the OptionStringQueryLabels to the newer format
Adds OptionBoolQueryLinkFailedJob ("adbc.bigquery.sql.query.link_failed_job").
When set on a statement, query errors are wrapped with a URL pointing at
the BigQuery web console job so operators can jump straight to the
failed run.
The project id, location, and job id are all URL-safe:
- Project id and job id can only contain URL-safe characters:
https://cloud.google.com/bigquery/docs/reference/rest/v2/JobReference
- Locations are also URL-safe:
https://cloud.google.com/bigquery/docs/locations
Threads the flag through runQuery, runPlainQuery,
queryRecordWithSchemaCallback, and newRecordReader, and wraps errors
from safeWaitForJob, job.Read, iter.IsAccelerated, and
iter.ArrowIterator with the console link.
Ports arrow-adbc commit d8b3f2a8b (Add option to link failed jobs, adbc-drivers#80).
After a query returns, publish the executing job's ID under the Arrow schema metadata key BIGQUERY:query_id (constant MetadataKeyBigqueryQueryID). Lets callers surface the BigQuery job that produced a result set without additional API calls. runQuery now sets query.JobID = job.ID() on the successful path. metadataFromJobStatistics takes a jobID param and, when non-empty, adds MetadataKeyBigqueryQueryID to the emitted metadata. The three callers (ipcReaderFromArrowIterator, makeDryRunReader, and statement.ExecuteSchema) thread the ID through. Ports arrow-adbc commit a0ac7f004 (Add query JobID to Arrow schema metadata, adbc-drivers#86).
Adds OptionBoolUseStorageApiDisabledClient
("adbc.bigquery.sql.query.use_storage_api_disabled_client"). When set,
the statement routes queries through a secondary bigquery.Client whose
Storage Read API is NOT enabled, so pseudo-columns like _PARTITIONDATE
and _PARTITIONTIME return values instead of nulls.
Wiring:
- ContextKeyUseStorageApiDisabledClient carries the flag from
statement.ExecuteQuery/ExecuteUpdate into runQuery.
- runQuery gains an alloc parameter and, when the ctx flag is set,
wraps bigquery.RowIterator in a RowBasedArrowIterator (new file
row_based_iterator.go) instead of failing on !iter.IsAccelerated().
- The row iterator batches 1000 rows at a time, converts to an Arrow
RecordBatch via rowsToArrowRecordBatch, and re-serializes as IPC
bytes so downstream code that expects bigquery.ArrowIterator sees
the same interface.
- Connection refactor from earlier commit (authOptions helper +
lazily-created clientStorageApiDisabled) is unchanged.
Type support in rowsToArrowRecordBatch is initially limited to DATE
and TIMESTAMP because those cover the intended pseudo-column use
cases; extend as more types come up.
Ports the pseudo-columns half of arrow-adbc commit a79555b0f (allow
fetching data selecting pseudo columns; support copy table, adbc-drivers#96).
The copy-table half is a separate commit.
Adds copy_table.{source,destination,write_disposition} options and
dispatches to a new executeCopyTable helper (table_ops.go) when
statement.copyTableSource is set. Uses the standard bigquery.Copier
API for a server-side table copy, honoring the write_disposition if
provided.
Ports the copy-table half of arrow-adbc commit a79555b0f (allow
fetching data selecting pseudo columns; support copy table, adbc-drivers#96) —
the storage-api-disabled half is a separate commit.
rebase the latest upstream main, resolve conflicts: add TODOs to migrate
the legacy options
Adds OptionJsonUpdateTableColumnsDescription
("adbc.bigquery.table.update_columns_description"). Value is a JSON
object mapping column name to new description; ExecuteQuery fetches
the destination table's current schema, overlays the descriptions,
and calls Table.Update. Columns absent from the map are untouched.
Ports arrow-adbc commit 13711b856 (support update a table's columns'
descriptions, adbc-drivers#34).
rebase the latest upstream main, resolve conflicts: add TODOs to migrate
the legacy options
Extends executeUpdateTableColumnsMetadata with support for
OptionJsonUpdateTableColumnsPolicyTags
("adbc.bigquery.table.update_columns_policy_tags"). Value is a JSON
object mapping column name to a list of BigQuery policy tag IDs.
RECORD-typed columns are skipped because BigQuery does not allow
policy tags on nested types.
The two updates share the same code path so descriptions and policy
tags can be applied in a single Table.Update call.
Ports arrow-adbc commit e8c77a4c5 (support per-column policy tag
updates, adbc-drivers#128).
Adds OptionStringUpdateTableDescriptionValue
("adbc.bigquery.table.update_description"). When set, ExecuteQuery
dispatches to executeUpdateTableDescription which issues a
TableMetadataToUpdate{Description: …} on the destination table. Only
the description is touched; schema and other metadata are unchanged.
Ports arrow-adbc commit 51466a2dc (supports update table description,
Adds OptionJsonAuthorizeViewToDatasets
("adbc.bigquery.dataset.authorize_view_to_datasets"). Value is a JSON
object mapping view reference to a list of {project, dataset} pairs;
each dataset gets an AccessEntry granting the view SELECT access to
its source data. The update is idempotent — views already authorized
on a dataset are skipped.
Adds a connection helper datasetInProject() so callers outside
connection.go can obtain a Dataset handle bound to the current client.
Ports arrow-adbc commit e51c22694 (authorized views support, adbc-drivers#37).
Adds the legacy CSV-file ingest path used by dbt seed ingestion, ported
from the legacy arrow-adbc bigquery driver. Triggered by setting
OptionStringIngestPath ("adbc.bigquery.ingest.csv_filepath"); the
driver opens the file, configures a BigQuery LoaderFrom against
QueryConfig.Dst with SKIP_LEADING_ROWS=1 and either an explicit
IPC-encoded schema (via OptionStringIngestSchema as bytes or string)
or AutoDetect.
This is separate from the new ADBC bulk-ingest API — new code should
prefer the bulk-ingest path. Kept for compatibility with existing
callers using dbt seeds through the legacy driver.
Ports arrow-adbc commit 400da64e5 (preliminary support of seed
ingestion in bigquery, adbc-drivers#29).
Adds statement-level support for the dbt python-models execution path,
ported from the legacy arrow-adbc bigquery driver. When any of the
Dataproc/GCS statement options are set, ExecuteQuery dispatches to a
new python_models.go path instead of running SQL:
- adbc.bigquery.create_batch.{parent,batch_yml,batch_id} — submit a
Dataproc serverless Batch (YAML payload parsed into dataprocpb.Batch).
- adbc.bigquery.dataproc.submit_job.{cluster_name,gcs_path} — submit a
PySpark job on an existing Dataproc cluster.
- adbc.bigquery.write_gcs.{bucket,object_name,content} — upload a
string to GCS (used to stage the model source).
- adbc.bigquery.dataproc.{compute_region,project,pooling_timeout} —
Dataproc addressing + operation timeout.
Adds three connection helpers (newDataprocBatchClient,
newJobControllerClient, newGCSClient) that reuse the connection's
authOptions() so credentials, quota project, and impersonation are
shared across BigQuery/Dataproc/GCS clients.
Ports arrow-adbc commit fad437ac7 (python models(serverless + cluster),
rebase the latest upstream main, resolve conflicts: go.mod is updated
with a few new packages
Adds the second execution mode for dbt python models: Vertex AI Notebook execution jobs, used by the bigframes path. When adbc.bigquery.notebook_execute_job.parent is set, ExecuteQuery dispatches to executeCreateNotebookExecutionJob which: - Resolves a notebook runtime template (using template_id if provided, else the ONE_CLICK template, else creating a default one). - Submits a NotebookExecutionJob targeting the GCS-hosted source notebook and staging outputs at gs://<gsc_bucket>/<model>/logs. - Populates ExecutionIdentity from the connection auth (service account JSON email, impersonation principal, or user email via the userinfo endpoint). - Polls for job completion with the dataproc pooling timeout and streams a summary of the notebook outputs to the connection Logger. adbc.bigquery.notebook_execute_job.* options cover parent/project/ region/template_id/model_name/model_file_name/gsc_bucket/gsc_path. Ports arrow-adbc commit 9ebc9900f (python models bigframes, adbc-drivers#97). merge the latest upstream main, resolve conflicts: add TODOs to migrate the legacy options
BigQuery's FloatFieldType is the string "FLOAT", which is not a valid literal in Standard SQL. Callers that read BIGQUERY:type from Arrow metadata and use it in generated SQL would produce invalid queries. Override to "FLOAT64" when the field is a bigquery.FloatFieldType. Ports arrow-adbc commit 0c7a2146e (Fix float rendering in Bigquery, adbc-drivers#127).
When Arrow schemas carry a BIGQUERY:type metadata annotation (set by this driver's buildField on the read path), consult it in arrowFieldToBigQueryField so we can round-trip BigQuery logical types that don't have a distinct Arrow physical type: DATETIME vs TIMESTAMP, NUMERIC vs BIGNUMERIC, JSON, GEOGRAPHY, INTERVAL. Parameter suffixes like NUMERIC(38,9) or ARRAY<INT64> are trimmed before lookup; composite ARRAY/RECORD/STRUCT are left to the Arrow-driven path. Fixes seed ingestion when the source of the data was BigQuery itself (reader -> writer round-trip previously downgraded logical types). Ports arrow-adbc commit ace9cd8a1 (add more type support for types in seeds, adbc-drivers#121), which supersedes the earlier adbc-drivers#118 metadata mapper.
Adds a new auth type OptionValueAuthTypeExternalAccount that lets the driver federate identity from an external OAuth2 IdP into Google via the STS token-exchange flow. The subject token is fetched from external_account.request_url with a client-credentials request_data body and exchanged at https://sts.googleapis.com/v1/token; if service_account_impersonation_url is set, Google impersonates that service account for downstream calls. Options: - adbc.bigquery.sql.auth_type.external_account (auth type value) - adbc.bigquery.sql.auth.external_account.audience - adbc.bigquery.sql.auth.external_account.impersonation_url (optional) - adbc.bigquery.sql.auth.external_account.request_url - adbc.bigquery.sql.auth.external_account.request_data Implementation: - idpTokenSupplier implements externalaccount.SubjectTokenSupplier, POSTing to the IdP with resty configured for capped exponential backoff on transient failures (POST retries explicitly allowed). - Error messages surface OAuth error/error_description bodies when present, cap large HTML bodies, and call out HTTP 429 rate-limits distinctly. - Uses default STS endpoint and JWT subject-token type (added as package-level constants). Ports arrow-adbc commit e7a344229 (add external-account (WIF) auth support, adbc-drivers#134). Corresponds to the auth_type::EXTERNAL_ACCOUNT path already referenced by fs/sa/crates/dbt-auth/src/bigquery/mod.rs. merge the latest upstream main, resolve conflicts: remove OptionValueAuthTypeAppDefaultCredentials, OptionValueAuthTypeJSONCredentials, OptionValueAuthTypeOAuthClientIDs since they're redeclared using values of the new format, and they are already coverted in the remapping map
Publish additional BigQuery table-metadata keys on the schema returned by
GetTableSchema. Downstream consumers (fs adapter) already depend on
these key names, so match them exactly.
Adds keys:
- ViewQuery, UseLegacySQL, UseStandardSQL
- Clustering.Fields (JSON-encoded via new encodeJson helper)
- ExpirationTime
- ExternalDataConfig.{SourceFormat, SourceURIs, AutoDetect, Compression,
IgnoreUnknownValues, MaxBadRecords, HivePartitioningOptions.*,
DecimalTargetTypes, ConnectionID, ReferenceFileSchemaURI,
MetadataCacheMode}
- EncryptionConfig.KMSKeyName
- StreamingBuffer.{EstimatedBytes, EstimatedRows, OldestEntryTime}
- TableConstraints.PrimaryKey.Columns
- ResourceTags
Behavior change: RequirePartitionFilter is now emitted unconditionally
(previously gated behind the bool being true), matching the legacy
driver so downstream code that checks for the key's presence still works
when the value is 'false'.
Refactors the ad-hoc inline JSON encoding for Labels to the new generic
encodeJson helper.
Ports arrow-adbc commit 735f692d7 (Add extra metadata, adbc-drivers#67).
Scopes are base OAuth scopes, not an impersonation trigger. Counting them in hasImpersonationOptions() makes any connection that sets scopes without a target principal fail with "`bigquery.impersonate.target_principal` parameter is empty for impersonation" — which is every connection from fs, since dbt-auth always sends the four default scopes (bigquery, cloud-platform, drive, userinfo.email). - hasImpersonationOptions() on connectionImpl and databaseImpl no longer counts impersonateScopes. - authOptions() applies the scopes as base OAuth scopes via option.WithScopes, so non-impersonation connections stop silently dropping them. - getAccessToken() forwards them as the `scope` form parameter on the refresh-token request. Ports arrow-adbc commit 9a400e9b8 (support service account impersonation, adbc-drivers#75).
ajhlee-dbt
left a comment
There was a problem hiding this comment.
My two commits look good. Just curious, why do we prefer putting JobID as another parameter in ipcReaderFromArrowIterator? I do like setting it under metadataFromJobStatistics
| // TODO (harry): add the old option values to optionRemapping map for backward-compatibility | ||
| // OptionBoolQueryLinkFailedJob, when set, instructs the driver to include the | ||
| // BigQuery web console link to a failed query job in the error message. | ||
| // (Option is stored but the error-wrapping path is a TODO; see FINDINGS.) |
There was a problem hiding this comment.
There's a reference to FINDINGS here that I don't understand. Where can I find it?
There was a problem hiding this comment.
good catch, comments are stale, I'll clean them up!
|
|
||
| // Wrap failed-query errors with a link to the BigQuery web console job. | ||
| // Currently option-storage only; see FINDINGS for the TODO. | ||
| linkFailedJob bool |
|
I just reviewed my commits and it looks good to me! The skipped commits also look good! |
@ajhlee-dbt |
stringToTable built its result as a &bigquery.Table{...} literal, which
leaves the type's unexported client reference nil. Every method that
reaches the API through it then panics with a nil pointer dereference:
- Dst.LoaderFrom(...).Run() on the CSV ingest path (dbt seeds)
- Dst.Metadata(ctx) on the update-columns path
Both surfaced as "Go panic in bigquery driver: invalid memory address or
nil pointer dereference" from AdbcStatementExecuteQuery, and the panic
poisons the driver so all later AdbcDatabaseNew calls fail too.
Route through connectionImpl.table, which was already present and does
this correctly, so the returned table carries its client.
Deviates from the upstream commit only in signature: stringToTable takes
*connectionImpl rather than *statement, since the connection is all it
needs. Upstream's other half — renaming OptionStringQueryDestination to
OptionStringQueryDestinationTable — is already in this driver, which is
why the commit was originally triaged as "already in new".
Ports arrow-adbc commit 883c5882a (bring back
OptionStringQueryDestinationTable, adbc-drivers#33).
| // SubjectToken is invoked by the externalaccount token source to fetch the | ||
| // IdP subject token. | ||
| // See https://pkg.go.dev/golang.org/x/oauth2/google/externalaccount#SubjectTokenSupplier | ||
| func (s *idpTokenSupplier) SubjectToken(ctx context.Context, _ externalaccount.SupplierOptions) (string, error) { |
There was a problem hiding this comment.
Are we not porting the tests connection_test.go?
Port over commits under go/adbc/driver/bigquery from https://github.com/dbt-labs/arrow-adbc to the foundry driver fork
This PR should use rebase and merge to merge.
Better to be reviewed by commit. Most of the conflicts when rebasing came from the adbc-drivers#215 that updated the format of the option values, I preserve most of the option values added by us using the legacy format, and added TODOs to address them in a separate PR
More details can be found in the PORT_REVIEW.md