Skip to content

Add BigQuery job labels and job timeout support#7

Open
anujdh1 wants to merge 1 commit intomasterfrom
anujdh1-job-labels-timeout
Open

Add BigQuery job labels and job timeout support#7
anujdh1 wants to merge 1 commit intomasterfrom
anujdh1-job-labels-timeout

Conversation

@anujdh1
Copy link
Collaborator

@anujdh1 anujdh1 commented Mar 18, 2026

Summary

  • Parse account_id and job_timeout from DSN query parameters in configFromUri()
  • Apply account_id as a BigQuery job label in buildQuery() for server-side cost attribution; defaults to "UNSPECIFIED" when not set (best-effort, never fails)
  • Apply job_timeout as JobTimeout on queries for bounded execution; parsing is best-effort (silently ignores invalid values)
  • Split query.Read(ctx) into query.Run(ctx) + job.Read(ctx) in ExecContext/QueryContext to get a job handle
  • On context cancellation, fire go job.Cancel(context.Background()) to cancel the server-side BQ job (fire-and-forget, non-blocking)
  • Remove verbose debug/info logs from ExecContext, QueryContext, Exec, Query (keep pre-existing logrus.Debugf in buildParameter helpers)
  • Remove logrus dependency from driver.go

DSN format: bigquery://project/dataset?account_id=X&job_timeout=5m

Note on JobTimeout: BigQuery's JobTimeout (configuration.jobTimeoutMs) applies only to job execution time, not queue/pending time. If a job is waiting for slots, that wait does not count toward the timeout. The timeout clock starts when the job begins executing. BQ also floors the minimum timeout to 1 second.

Test Plan

E2E Validation on dev-113 (via sdmain vendor)

Deployment: dev-113 | Test Account: jtijl1035 | GCP Project: spark-dev-113

  1. Vendored these driver changes into sdmain and deployed turing-devops-service to dev-113
  2. Created test account jtijl1035, ingested 2 records, ran read query
  3. Verified BQ job metadata via bq show -j:
    • configuration.labels.account_id = jtijl1035
    • configuration.jobTimeoutMs = 300000 (5 min)
    • status.state = DONE

Timeout + Server-Side Cancellation Test

  1. Set job_timeout=1ms in DSN (BQ floors to 1 second minimum)
  2. Verified via bq show -j:
    • status.errorResult.reason = stopped
    • status.errorResult.message = Job execution was cancelled: Job timed out after 1 sec
  3. Confirmed server-side job cancellation works correctly

Context Cancellation

  • ExecContext/QueryContext split query execution into Run + Read phases
  • If job.Read(ctx) fails and ctx.Err() != nil, fires go job.Cancel(context.Background()) to cancel the BQ job server-side
  • Uses fire-and-forget goroutine to avoid blocking the caller
Feature Expected Actual Status
Job label account_id Set from DSN Confirmed in BQ job metadata PASS
Job timeout Set from DSN Confirmed (jobTimeoutMs: 300000) PASS
Server-side cancellation on timeout Job status stopped Confirmed PASS
Default account_id UNSPECIFIED Code verified OK
Invalid job_timeout Silently ignored Code verified OK
Context cancellation → job.Cancel Fires async Code verified OK

query.Labels = map[string]string{
"account_id": statement.connection.config.accountID,
}
logrus.Infof("bq driver: set job label account_id=%s",
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove this log

}
if statement.connection.config.jobTimeout > 0 {
query.JobTimeout = statement.connection.config.jobTimeout
logrus.Infof("bq driver: set job timeout=%s",
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove this log

driver/driver.go Outdated
config.location = fields[0]
}

logrus.Infof(
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove the comment

driver/driver.go Outdated
scopes: getScopes(u.Query()),
endpoint: u.Query().Get("endpoint"),
disableAuth: u.Query().Get("disable_auth") == "true",
accountID: u.Query().Get("account_id"),
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what happens if account ID is not set? we should fetch it best-effort

@anujdh1 anujdh1 force-pushed the anujdh1-job-labels-timeout branch 3 times, most recently from 26b4d65 to 63821dc Compare March 19, 2026 22:20
Parse account_id and job_timeout from DSN query parameters and apply
them to BigQuery queries:
- account_id is set as a job label for server-side cost attribution
- job_timeout sets JobTimeout on the query for bounded execution
- Fix QueryContext to use caller's context instead of context.Background()

DSN format: bigquery://project/dataset?account_id=X&job_timeout=5m
@anujdh1 anujdh1 force-pushed the anujdh1-job-labels-timeout branch from 63821dc to f212971 Compare March 19, 2026 22:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant