feat(go): assume role, per-service endpoints, API retries and polling options by fpiped · Pull Request #20 · dbt-labs/athena

fpiped · 2026-09-25T13:12:26Z

Connection options for the settings of a dbt-athena profile. Reworked after review: the query statistics are now #22, the cancellation fix #21, and the Iceberg commit retries are gone.

option	effect
`athena.aws.role_arn`, `athena.aws.role_external_id`, `athena.aws.role_session_name`, `athena.aws.role_duration`	AssumeRole on top of the auth type's credentials, behind a credentials cache that refreshes them before they expire; every AWS client of the driver uses the role
`athena.aws.max_api_retries`	retries of a failed AWS API request, the first attempt not included (SDK default 2); a query that ran and failed is never rerun
`athena.aws.athena_endpoint_url`, `athena.aws.sts_endpoint_url`	one endpoint per service the driver calls, e.g. VPC interface endpoints without private DNS; the SDK's `AWS_ENDPOINT_URL_ATHENA` / `AWS_ENDPOINT_URL_STS` settings still apply when they are unset
`athena.poll_interval`	interval between `GetQueryExecution` calls; 500 ms unless set

Durations are Go durations ("1h", "500ms"); invalid values fail SetOption with InvalidArgument. The driver's defaults do not change. service/sts becomes a direct requirement.

Reruns of a query that fails with ICEBERG_COMMIT_ERROR moved to dbt's engine (dbt-labs/dbt#16274), as the application's job.

Testing

athena_database_test.go, without a network: the retryer's attempts on the Athena and STS clients (unset 3, 0 retries 1, 5 retries 6); each client's base endpoint, from the option or the SDK environment variables, the option winning; AssumeRole through a mock STS client (role, external ID, session name, duration), the Athena client signing with the assumed credentials, and no AssumeRole without a role; the poll interval (three polls at 50 ms take 150 ms to 1 s). Dropping the +1 of the retries or restoring the fixed 500 ms fails them.
TestIntegration_AthenaEndpointURL on a real work group: the region's endpoint runs a query, 127.0.0.1:9 fails on it.
go test -race ./..., go vet, the other integration tests.
End to end through dbt on Athena: a role with an external-ID trust policy is assumed, and refused with a wrong external ID; an explicit endpoint runs models and a dead one fails.

#14 moves Open, GetOption and SetOption to the context-aware API and adds a Glue client: if it lands first, I'll rebase this onto it and add athena.aws.glue_endpoint_url.

iconara

Athena person here, not a maintainer.

General comment: this introduces a few different, unrelated things. As a reviewer, I would have preferred them to be separate PRs. I have no comments on some of the added features, but strong opinions on others (Iceberg retries).

Comment from the service team point of view: I understand that it's convenient when a driver adds functionality missing from a service from the point of view of a specific use case – to be concrete: being able to retry Iceberg commits more times than the service does, because this can be a problem in a data processing pipeline. However, this means that this can't be a general driver for Athena. When we in the service team build drivers, we need to think about all use cases, and the benefits and downsides of features to all of them. Something that is convenient in one use case may cause complications, or unexpected costs in another. They are also not always possible to continue to support as the service evolves. Adding features like this is fine if this driver is only ever going to be the layer between dbt and Athena, but it means that it would be hard for us to support it as a general ADBC driver for Athena.

iconara · 2026-09-29T12:57:10Z

+	// OptionIcebergCommitRetries is how many times a query that fails with
+	// ICEBERG_COMMIT_ERROR, a commit conflict with a concurrent Iceberg write, is
+	// run again (default 0).
+	OptionIcebergCommitRetries = "athena.iceberg_commit_retries"


This option is problematic, for two reasons:

The name is easily confused with the table property of the same name, which controls the number of commit retries that the engine does before failing. Retrying a commit in the engine is a cheap operation, retrying from the client side means recalculating the transaction in a new query.

Retrying Iceberg commits is the responsibility of the engine. If a transaction fails, it's most often because there are multiple clients attempting to commit at the same time. Retrying in a driver and not propagating the error to the application will in most cases make the problem worse. Retrying a failed Iceberg commit should be an application responsibility.

Agreed, removed. The reruns now live in dbt's engine (dbt-labs/dbt#16274), as dbt-athena does them in the adapter: up to num_iceberg_retries, after a random wait, so that writers that conflicted together do not retry together. The driver reports the ICEBERG_COMMIT_ERROR and stops there.

iconara · 2026-09-29T12:58:09Z

+	OptionRoleDuration = "athena.aws.role_duration"
+	// OptionMaxAttempts is the maximum number of attempts of each AWS API request,
+	// the first one included, for the SDK's standard retryer (default 3).
+	OptionMaxAttempts = "athena.aws.max_attempts"


This is a very generic name that could easily be confused for many different things, can you find a more specific name for it, like "athena.aws.max_api_retries"?

Renamed athena.aws.max_api_retries. It now counts retries, the first attempt not included, which is also what dbt-athena's num_retries counts, and it only covers the API calls: a query that ran and failed is not rerun.

iconara · 2026-09-29T13:00:30Z

+	OptionMaxAttempts = "athena.aws.max_attempts"
+	// OptionEndpointURL overrides the endpoint of the Athena API, e.g. a VPC
+	// interface endpoint. Glue, S3 and STS keep their default endpoints.
+	OptionEndpointURL = "athena.endpoint_url"


There needs to be an endpoint for each service, otherwise the VPC case will not work. If the driver is used from within a VPC, the likelihood is that all services are only reachable through VPC-specific endpoints.

Now one option per service the driver calls: athena.aws.athena_endpoint_url and athena.aws.sts_endpoint_url (STS for the assumed role). athena.aws.glue_endpoint_url comes with the Glue client of #14. The SDK's AWS_ENDPOINT_URL_<SERVICE> settings also apply when the options are unset; a test covers both.

iconara · 2026-09-29T13:02:41Z

+	// MetadataKeyQueryID is the query result schema metadata key holding the
+	// query execution ID.
+	MetadataKeyQueryID = "ATHENA:query_id"
+	// MetadataKeyDataScannedInBytes holds the bytes the query scanned
+	// (QueryExecutionStatistics.DataScannedInBytes).
+	MetadataKeyDataScannedInBytes = "ATHENA:Statistics:DataScannedInBytes"
+	// MetadataKeyUpdateCount holds the rows written by a CREATE TABLE AS SELECT,
+	// INSERT INTO, MERGE, DELETE or UNLOAD (GetQueryResults UpdateCount); absent
+	// when Athena reports none, as for DDL.
+	MetadataKeyUpdateCount = "ATHENA:UpdateCount"


Why the inconsistent naming of these options? query_id in snake case, Statistics:DataScannedInBytes in camel case, and UpdateCount without the statistics prefix (I guess that one is consistent in a way since the update count is not in the statistics of the QueryExecution data type).

Fair. The query statistics moved to #22, with the keys lowercase like the existing field metadata key ATHENA:type: ATHENA:query_execution_id, ATHENA:data_scanned_in_bytes, ATHENA:update_count. update_count is on its own because it comes from GetQueryResults, not from the execution statistics.

iconara · 2026-09-29T13:04:28Z

+
+// athenaServer answers every Athena API call with an error, recording the target
+// operation and the access key that signed the request.
+func athenaServer(t *testing.T) (*httptest.Server, *[]string) {


It seems excessive to me to start up a server to respond to API calls. Mocking the AWS SDK feels more sensible. Is there a benefit to this that I'm missing?

No real benefit: the server checked that the options reached the SDK clients, which their Options() show directly. Replaced by athena_database_test.go: assertions on the built Athena and STS clients (endpoint, retryer attempts) and a mock STS client for AssumeRole, no network. The only real-endpoint check left is an integration test.

… options - athena.aws.role_arn, role_external_id, role_session_name, role_duration: an IAM role assumed on top of the auth type's credentials, for every AWS client of the driver. - athena.aws.max_api_retries: retries of a failed AWS API request, the first attempt not included. A query that ran and failed is not rerun. - athena.aws.athena_endpoint_url and athena.aws.sts_endpoint_url: one endpoint per service, for VPC interface endpoints without private DNS. The SDK's AWS_ENDPOINT_URL_<SERVICE> settings still apply when they are unset. - athena.poll_interval: the interval between query status checks, 500ms by default as before.

fpiped · 2026-10-13T10:03:17Z

@iconara thanks for the review. Split as you suggested: this PR now has the connection options only, the query statistics are #22 and the cancellation fix #21. On the general-driver point: the Iceberg retries are gone (now in dbt), and what is left are settings any Athena client can use, with the SDK defaults unchanged when they are unset.

…now names them dbt-labs/athena#20 renames its options after review: athena.aws.max_api_retries counts retries, as dbt-athena's num_retries and botocore's max_attempts do, so the value passes through unchanged; endpoint_url maps to athena.aws.athena_endpoint_url, the endpoint dbt-athena points its Athena client at. num_iceberg_retries no longer reaches the driver: reruns of ICEBERG_COMMIT_ERROR belong to the engine.

fpiped requested review from a team as code owners September 25, 2026 13:12

This was referenced Sep 25, 2026

feat(athena): AthenaDbConfig profile schema dbt-labs/dbt#16274

Open

feat(athena): Part 6 — AthenaAdapter methods behind the dbt-athena macros dbt-labs/dbt#16376

Open

[Feature] Athena adapter for dbt Core 2.0 dbt-labs/dbt#16252

Open

iconara reviewed Sep 29, 2026

View reviewed changes

fpiped mentioned this pull request Sep 29, 2026

feat(go): run Glue, S3 and STS operations behind statement options #19

Draft

David-Banquet mentioned this pull request Oct 9, 2026

chore(go): add support for more data types #15

Open

This was referenced Oct 13, 2026

fix(go): stop the query when a cancellation fails the status call #21

Open

feat(go): report query statistics in the result schema metadata #22

Open

fpiped force-pushed the feat/connection-options branch from bc6e583 to 0df2075 Compare October 13, 2026 10:02

fpiped changed the title ~~feat(go): query statistics and dbt-athena connection options~~ feat(go): assume role, per-service endpoints, API retries and polling options Oct 13, 2026

Navigation Menu

Repository navigation

feat(go): assume role, per-service endpoints, API retries and polling options - #20

feat(go): assume role, per-service endpoints, API retries and polling options#20
fpiped wants to merge 1 commit into
dbt-labs:mainfrom
fpiped:feat/connection-options

fpiped commented Sep 25, 2026 •

edited

Loading

Uh oh!

iconara left a comment

Uh oh!

iconara Sep 29, 2026

Uh oh!

fpiped Oct 13, 2026

Uh oh!

iconara Sep 29, 2026

Uh oh!

fpiped Oct 13, 2026

Uh oh!

iconara Sep 29, 2026

Uh oh!

fpiped Oct 13, 2026

Uh oh!

iconara Sep 29, 2026

Uh oh!

fpiped Oct 13, 2026

Uh oh!

iconara Sep 29, 2026

Uh oh!

fpiped Oct 13, 2026

Uh oh!

fpiped commented Oct 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

fpiped commented Sep 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing

Uh oh!

iconara left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fpiped commented Oct 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

fpiped commented Sep 25, 2026 •

edited

Loading