feat(go): run Glue, S3 and STS operations behind statement options by fpiped · Pull Request #19 · dbt-labs/athena

fpiped · 2026-09-20T13:57:54Z

Closes #18.

What

Two statement options select an AWS API call in place of the SQL text:

stmt.SetOption("athena.operation", "glue.get_table")
stmt.SetOption("athena.operation.payload", `{"DatabaseName": "analytics", "Name": "orders"}`)
rdr, _, err := stmt.ExecuteQuery(ctx) // one row, utf8 column `result`: {"Table": {...}}

The payload is the JSON form of the AWS API input, the result the JSON form of the output; paginated calls return every page, batch calls are chunked (25 / 100 partitions, 1000 keys), a missing Glue table answers {"Table": null} / {"Deleted": false} and a missing database {"Database": null} instead of failing; glue.get_tables on a missing database answers an empty list, as dbt-athena's list_relations_without_caching treats it. ExecuteUpdate runs the operation and returns 0. Unknown operations and malformed payloads are InvalidArgument; AWS errors are IO, like SQL failures. The Glue / S3 / STS clients are built from the database's existing AWS config on first use and are mockable like athenaClientAPI. Six Athena operations run the Spark sessions and calculations of a Spark-enabled work group (athena.start_session, athena.get_session_status, athena.start_calculation_execution, athena.get_calculation_execution, athena.stop_calculation_execution, athena.terminate_session), which dbt-athena uses for Python models (the stop and the terminate run to completion, within 10 s, even when the statement is cancelled: dbt cancels every open statement on Ctrl-C, and a calculation left running still writes its table), and ten lakeformation.* operations for its LF-Tags, data cell filters and their grants; the single-call ones go through a generic callAPI. glue.get_tables, glue.get_database and glue.get_databases carry the metadata reads dbt-athena makes from Glue (relations, columns, schemas), so a lookup costs no Athena query. The operation table is in the README.

Why

See #18: dbt-athena's materializations manage Glue and S3 directly and run Python models on Athena Spark, and dbt Core 2.0 reaches AWS only through this driver.

Testing

make test-unit (22 new tests with mock Athena / Glue / S3 / STS / Lake Formation clients: option handling, each operation, pagination, chunking, not-found and error mapping) and make vet. Exercised end to end through the engine side (dbt-labs/dbt#16376) against a real work group: seeds, views, tables, incremental models and snapshots, with the Glue and S3 side effects checked, Python models (a Spark DataFrame to a Hive table, pandas to Iceberg, an incremental append) on a Spark work group, a model timeout and a Ctrl-C ending its calculation and session, Lake Formation tags and data cell filters read back from Lake Formation, and relations, columns and schemas read from Glue, S3 Tables catalogs and a Glue catalog registered in Athena included.

go get of the four service modules moved the aws-sdk-go-v2 core to v1.47.1 and smithy-go to v1.28. go mod tidy fails on main as well (the driverbase-go/testutil pseudo-version), so go.sum carries only what go get added.

Athena tables are Glue entries over S3 prefixes and Athena SQL manages neither, so dbt-athena drops tables, expires Glue versions, cleans partitions and stages seeds through the AWS APIs. A client that only has this driver had no way to do that with the connection's credentials. Two statement options select an operation in place of the SQL text: `athena.operation` names it, `athena.operation.payload` carries the JSON form of the AWS API input; the result is one record with a utf8 `result` column holding the JSON form of the API output. Paginated calls return all pages, batch calls are chunked to the API limits (25 / 100 partitions, 1000 keys), a missing Glue table is `{"Table": null}` / `{"Deleted": false}` rather than an error. Sixteen operations over Glue, S3, STS and the Athena control plane; the SDK clients are built from the database's existing AWS config on first use and mockable like the Athena client. Adds aws-sdk-go-v2 service/glue, service/s3 and service/sts, which moves the SDK core to v1.47.

Copilot

Copilot review overview

🟡 Changes recommended

Three unresolved moderate issues affect validation and documented payload behavior.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

Decode STS payload before making the request · New

What changed in this PR

Adds statement options for invoking AWS Glue, S3, and STS operations with JSON payloads and results, including pagination, batching, error mapping, client injection, tests, and documentation.

Changes:

Implements AWS operation dispatch and result readers.
Adds mockable AWS clients and operation tests.
Updates documentation and AWS SDK dependencies.

Review findings:

Unknown operations should be validated before AWS client initialization (moderate, 1 vote).
STS payloads must be decoded and malformed JSON rejected (moderate, 3 votes).
S3 delete-object payload handling must match or clearly document the AWS input contract (moderate, 1 vote).

File	Description
README.md	Documents operation options and supported APIs.
go/statement.go	Routes statements through operation execution.
go/operations.go	Implements Glue, S3, STS, and Athena operations.
go/operations_test.go	Adds operation unit tests and mocks.
go/go.sum	Adds dependency checksums.
go/go.mod	Updates AWS dependencies.
go/functional_test.go	Extends Athena client mocks.
go/connection.go	Caches auxiliary AWS clients.
go/client.go	Extends the Athena client interface.
go/athena_database.go	Injects auxiliary AWS clients.
CLAUDE.md	Documents the operations module.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…nt the S3 shapes An unknown operation now fails with InvalidArgument before any AWS client is built, and `sts.get_caller_identity` decodes its payload like the others, so a malformed one is rejected instead of ignored. The package doc and the README no longer claim every payload is the AWS API input: the three S3 operations take and return the simplified shapes in the table.

operations.go imports the three service packages directly.

`athena.start_session`, `athena.get_session_status`, `athena.start_calculation_execution`, `athena.get_calculation_execution` and `athena.stop_calculation_execution` pass the Athena API input and output as JSON, like the other operations. dbt-athena runs Python models this way: a calculation in a session of the profile's Spark work group.

`submit_python_job` follows dbt-athena's `AthenaPythonJobHelper`: the compiled model runs as a calculation in a Spark session of the profile's `spark_work_group`, through the driver's `athena.start_session`, `athena.get_session_status`, `athena.start_calculation_execution`, `athena.get_calculation_execution` and `athena.stop_calculation_execution` operations (dbt-labs/athena#19). - Model config as `AthenaSparkSessionConfig`: `timeout` (43200 s), `polling_interval` (the profile's `poll_interval`, else 1 s), `engine_config` over the defaults (1 / 2 / 1 DPU), Spark properties for `table_type` iceberg / hudi / delta_lake and the `spark_encryption`, `spark_cross_account_catalog`, `spark_requester_pays` flags; the same validation errors. - Sessions are shared by the process: a model takes a free session with the same engine configuration or starts one (10 minutes idle timeout). A session that ended is dropped and the calculation starts in a new one. - A failed or cancelled calculation reports its id, session, reason and stderr S3 path; a timeout or Ctrl-C stops the calculation.

…rations Ten `lakeformation.*` operations pass the API input and output as JSON (`list_data_cells_filter` and `list_permissions` return every page). dbt-athena manages LF-Tags on databases, tables and columns, and data cell filters with their SELECT grants, through these calls. The single-call operations, the Athena Spark ones included, go through a generic `callAPI`. Adds aws-sdk-go-v2 service/lakeformation (core v1.47.0 -> v1.47.1).

StopCalculationExecution is best effort; terminating the session is how a client makes sure an abandoned calculation stops.

dbt-athena reads relations, columns and schemas from Glue rather than information_schema. get_tables on a missing database returns an empty list, as dbt-athena's list_relations_without_caching does.

athena.stop_calculation_execution and athena.terminate_session are how a client ends Spark work it has given up on, often because it is being cancelled itself: dbt cancels every open statement on Ctrl-C, and the terminate never reached Athena, so the calculation ran to completion. Both now run on a context detached from the statement's, within a 10 s timeout.

`submit_python_job` follows dbt-athena's `AthenaPythonJobHelper`: the compiled model runs as a calculation in a Spark session of the profile's `spark_work_group`, through the driver's `athena.start_session`, `athena.get_session_status`, `athena.start_calculation_execution`, `athena.get_calculation_execution` and `athena.stop_calculation_execution` operations (dbt-labs/athena#19). - Model config as `AthenaSparkSessionConfig`: `timeout` (43200 s), `polling_interval` (the profile's `poll_interval`, else 1 s), `engine_config` over the defaults (1 / 2 / 1 DPU), Spark properties for `table_type` iceberg / hudi / delta_lake and the `spark_encryption`, `spark_cross_account_catalog`, `spark_requester_pays` flags; the same validation errors. - Sessions are shared by the process: a model takes a free session with the same engine configuration or starts one (10 minutes idle timeout). A session that ended is dropped and the calculation starts in a new one. - A failed or cancelled calculation reports its id, session, reason and stderr S3 path; a timeout or Ctrl-C stops the calculation.

iconara

Athena person here, not a maintainer.

I understand the convenience to be able to use the driver for things that are not Athena operations, but this will tie this driver tightly to dbt.

My recommendation is to not add any workarounds like these to the driver. Doing that would make it unsuitable to be a general Athena ADBC driver. Perhaps more importantly, it would mean that dbt would not be able to use any other Athena ADBC driver since they won't have the workarounds unique to this driver.

If there is no other way, I strongly suggest to at least:

Use the ADBC catalog API (GetObjects) as far as it is possible, e.g. looking up a table schema should not require going through Glue. There are limitations to the catalog API that we probably need to work with the ADBC community to fix, like no support for listing partitions
Use regular queries when the catalog API is not enough (e.g. listing partitions)
Don't add Athena Spark operations

fpiped · 2026-09-29T13:42:07Z

Thanks, that's clear, and it matches what you wrote on #20. I'll move #19 to draft and rework it in that direction. Metadata would go through GetObjects, with any gaps fixed in the driver in a general way. The driver already serves that from ListTableMetadata / GetTableMetadata, so a lookup doesn't cost a query. Where the catalog API isn't enough, it would use regular queries.

That leaves one problem. The aim is for dbt 2 on Athena to do everything dbt-athena does, and some of that can't be done with SQL or the catalog API:

uploading seeds to S3
cleaning up S3 when a table is dropped
Glue table versions
Lake Formation tags and data cell filters
Python models on Athena Spark

These need AWS API calls somewhere, and in the current series they have no place outside the driver.

`submit_python_job` follows dbt-athena's `AthenaPythonJobHelper`: the compiled model runs as a calculation in a Spark session of the profile's `spark_work_group`, through the driver's `athena.start_session`, `athena.get_session_status`, `athena.start_calculation_execution`, `athena.get_calculation_execution` and `athena.stop_calculation_execution` operations (dbt-labs/athena#19). - Model config as `AthenaSparkSessionConfig`: `timeout` (43200 s), `polling_interval` (the profile's `poll_interval`, else 1 s), `engine_config` over the defaults (1 / 2 / 1 DPU), Spark properties for `table_type` iceberg / hudi / delta_lake and the `spark_encryption`, `spark_cross_account_catalog`, `spark_requester_pays` flags; the same validation errors. - Sessions are shared by the process: a model takes a free session with the same engine configuration or starts one (10 minutes idle timeout). A session that ended is dropped and the calculation starts in a new one. - A failed or cancelled calculation reports its id, session, reason and stderr S3 path; a timeout or Ctrl-C stops the calculation.

fpiped requested review from a team as code owners September 20, 2026 13:57

This was referenced Sep 20, 2026

Expose the Glue / S3 / STS calls dbt-athena's materializations need #18

Open

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

Open

felipecrv requested a lite review from Copilot September 21, 2026 21:10

Copilot started reviewing on behalf of felipecrv September 21, 2026 21:10

View session

Copilot AI reviewed Sep 21, 2026

View reviewed changes

Comment thread go/operations.go Outdated

fpiped added 3 commits September 24, 2026 21:55

chore(go): list glue, s3 and sts as direct requirements

1e16297

operations.go imports the three service packages directly.

fpiped mentioned this pull request Sep 25, 2026

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

Open

feat(go): run athena.terminate_session

4af4ccf

StopCalculationExecution is best effort; terminating the session is how a client makes sure an abandoned calculation stops.

David-Banquet mentioned this pull request Sep 25, 2026

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

Open

3 tasks done

fpiped added 2 commits September 27, 2026 22:19

feat(go): glue.get_tables, glue.get_database and glue.get_databases

83969ff

dbt-athena reads relations, columns and schemas from Glue rather than information_schema. get_tables on a missing database returns an empty list, as dbt-athena's list_relations_without_caching does.

iconara reviewed Sep 29, 2026

View reviewed changes

fpiped marked this pull request as draft September 29, 2026 13:42

David-Banquet mentioned this pull request Oct 9, 2026

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

Open

Navigation Menu

Repository navigation

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

feat(go): run Glue, S3 and STS operations behind statement options#19
fpiped wants to merge 8 commits into
dbt-labs:mainfrom
fpiped:feat/glue-operations

fpiped commented Sep 20, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

iconara left a comment

Uh oh!

fpiped commented Sep 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

What

Why

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

iconara left a comment

Choose a reason for hiding this comment

Uh oh!

fpiped commented Sep 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fpiped commented Sep 20, 2026 •

edited

Loading