Skip to content

Commit 50e3d1e

Browse files
Support Amazon S3 Tables in the SQLAlchemy dialect
Athena exposes S3 Tables through a catalog named s3tablescatalog/<table-bucket> whose namespace is the database, and such tables use managed storage. Athena rejects a three-part catalog.namespace.table identifier in DDL ("Unsupported ddl with 2 catalogs"), so the catalog is selected on the connection via catalog_name (matched case-insensitively) and the namespace is used as the table schema. - Emit only TBLPROPERTIES for CREATE TABLE on an S3 Tables catalog: LOCATION, ROW FORMAT, SERDEPROPERTIES, and STORED AS are not accepted by managed Iceberg tables. Raise clear compile-time errors when the table does not declare table_type ICEBERG or specifies an explicit awsathena_location, instead of emitting DDL Athena would reject. - Store connection options in _create_connect_args so subclass dialects (pandas, arrow, polars, s3fs) that call it directly also expose catalog_name to the DDL compiler. - Add no-AWS compiler unit tests and AWS-gated E2E tests (CREATE TABLE with reflection, partition transform, CTAS with is_external=false) with idempotent DROP TABLE IF EXISTS cleanup, an S3 Tables docs section, and optional AWS_ATHENA_S3_TABLES_* test env vars. - Provision the S3 Tables test infrastructure in CloudFormation: table bucket, namespace, the federated s3tablescatalog Glue catalog (IAM-based access), OIDC-role IAM permissions, and 1-day unreferenced file removal to bound storage cost. Wire the env vars into the test workflow. - Import the existing S3 staging bucket into the CloudFormation stack with DeletionPolicy: Retain, capturing its lifecycle, encryption, ownership, and public-access settings (drift-free import). - Cancel superseded workflow runs with a concurrency group: overlapping runs stampede the account's Athena concurrent-query quota (TooManyRequestsException) and fail each other. Refs #710 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent d2cc374 commit 50e3d1e

10 files changed

Lines changed: 546 additions & 16 deletions

File tree

.github/workflows/test-suite.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ jobs:
1818
AWS_ATHENA_WORKGROUP: pyathena
1919
AWS_ATHENA_SPARK_WORKGROUP: pyathena-spark
2020
AWS_ATHENA_MANAGED_WORKGROUP: pyathena-managed
21+
# Registered S3 Tables catalog (s3tablescatalog/<table-bucket>) and namespace
22+
# for the SQLAlchemy S3 Tables tests; the table bucket, namespace, and the
23+
# AWS analytics-services integration are provisioned out of band.
24+
AWS_ATHENA_S3_TABLES_CATALOG: s3tablescatalog/laughingman7743-pyathena-s3-tables
25+
AWS_ATHENA_S3_TABLES_NAMESPACE: pyathena
2126

2227
strategy:
2328
fail-fast: false

.github/workflows/test.yaml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ permissions:
1212
id-token: write
1313
contents: read
1414

15+
# Cancel superseded runs: overlapping runs stampede the account's Athena
16+
# concurrent-query quota (TooManyRequestsException) and fail each other.
17+
concurrency:
18+
group: ${{ github.workflow }}-${{ github.ref }}
19+
cancel-in-progress: true
20+
1521
jobs:
1622
test:
1723
uses: ./.github/workflows/test-suite.yaml

docs/sqlalchemy.md

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,63 @@ If you want to limit the column options to specific table names only, specify th
363363
awsathena+rest://:@athena.us-west-2.amazonaws.com:443/default?partition=table1.column1%2Ctable1.column2&cluster=table2.column1%2Ctable2.column2&...
364364
```
365365

366+
## Amazon S3 Tables
367+
368+
[Amazon S3 Tables](https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-tables.html) are Iceberg-backed
369+
tables stored in a dedicated table bucket. Once the table bucket is integrated with the AWS analytics services,
370+
Athena registers it as a catalog named `s3tablescatalog/<table-bucket>`, with the S3 Tables namespace as the
371+
database.
372+
373+
Athena does not accept a three-part `catalog.namespace.table` identifier in DDL, so select the catalog on the
374+
**connection** with `catalog_name` and use the namespace as the table `schema`. Because S3 Tables use managed
375+
storage, the dialect omits the `LOCATION`, `ROW FORMAT`, and `STORED AS` clauses automatically. Tables must
376+
declare `awsathena_tblproperties={"table_type": "ICEBERG"}` (S3 Tables support only Iceberg), and
377+
`awsathena_location` must not be set — the dialect raises a compile-time error otherwise.
378+
379+
```python
380+
engine = create_engine(
381+
"awsathena+rest://athena.us-west-2.amazonaws.com:443/my_namespace"
382+
"?s3_staging_dir=s3://my-bucket/athena-results/"
383+
"&catalog_name=s3tablescatalog/my-table-bucket"
384+
)
385+
386+
table = Table(
387+
"some_table",
388+
MetaData(schema="my_namespace"),
389+
Column("id", types.Integer),
390+
Column("dt", types.Date, awsathena_partition=True, awsathena_partition_transform="day"),
391+
awsathena_tblproperties={"table_type": "ICEBERG"},
392+
)
393+
with engine.connect() as conn:
394+
table.create(bind=conn)
395+
```
396+
397+
which builds the following statement:
398+
399+
```sql
400+
CREATE TABLE my_namespace.some_table (
401+
id INT,
402+
dt DATE
403+
)
404+
PARTITIONED BY (
405+
day(dt)
406+
)
407+
TBLPROPERTIES (
408+
'table_type' = 'ICEBERG'
409+
)
410+
```
411+
412+
All Iceberg partition transforms (`year`, `month`, `day`, `hour`, `bucket`, `truncate`) are supported, the same as
413+
for other Iceberg tables. CTAS (`CREATE TABLE ... AS SELECT`) is not modeled as a SQLAlchemy construct; issue it as
414+
raw SQL via `text()`. Managed Iceberg CTAS requires `is_external = false`:
415+
416+
```python
417+
conn.execute(text(
418+
'CREATE TABLE "my_namespace"."some_table" '
419+
"WITH (table_type = 'ICEBERG', is_external = false) AS SELECT 1 AS id"
420+
))
421+
```
422+
366423
## Temporal/Time-travel with Iceberg
367424

368425
Athena supports time-travel queries on Iceberg tables by either a version_id or a timestamp. The `FOR TIMESTAMP AS OF`

pyathena/sqlalchemy/base.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,8 +207,7 @@ def create_connect_args(self, url: URL) -> tuple[tuple[str], MutableMapping[str,
207207
# awsathena+rest://
208208
# {aws_access_key_id}:{aws_secret_access_key}@athena.{region_name}.amazonaws.com:443/
209209
# {schema_name}?s3_staging_dir={s3_staging_dir}&...
210-
self._connect_options = self._create_connect_args(url)
211-
return cast(tuple[str], ()), self._connect_options
210+
return cast(tuple[str], ()), self._create_connect_args(url)
212211

213212
def _create_connect_args(self, url: URL) -> dict[str, Any]:
214213
opts: dict[str, Any] = {
@@ -238,6 +237,12 @@ def _create_connect_args(self, url: URL) -> dict[str, Any]:
238237
opts.update({"result_reuse_enable": bool(strtobool(opts["result_reuse_enable"]))})
239238
if "result_reuse_minutes" in opts:
240239
opts.update({"result_reuse_minutes": int(opts["result_reuse_minutes"])})
240+
# Store on the dialect so compilers can consult connection options
241+
# (e.g. catalog_name for S3 Tables detection). Assigned here rather than
242+
# in create_connect_args because subclass dialects call this method
243+
# directly and mutate the returned dict afterwards; sharing the same
244+
# object keeps _connect_options in sync with their updates.
245+
self._connect_options = opts
241246
return opts
242247

243248
@reflection.cache

pyathena/sqlalchemy/compiler.py

Lines changed: 95 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@
3939
_DialectArgDict = Mapping[str, Any]
4040
CreateColumn = Any
4141

42+
# Prefix of the Athena data catalog name registered for an Amazon S3 Tables
43+
# table bucket (e.g. ``s3tablescatalog/my-bucket``). It is selected via the
44+
# connection ``catalog_name``. S3 Tables are Iceberg-backed and use managed
45+
# storage, so their CREATE TABLE statements must not include a LOCATION clause.
46+
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-tables-integrations-query-athena.html
47+
S3_TABLES_CATALOG_PREFIX = "s3tablescatalog/"
48+
4249

4350
class AthenaTypeCompiler(GenericTypeCompiler):
4451
"""Type compiler for Amazon Athena SQL types.
@@ -305,6 +312,10 @@ class AthenaDDLCompiler(DDLCompiler):
305312
306313
- External table creation (EXTERNAL keyword for Hive-style tables)
307314
- Iceberg table creation (managed tables with ACID support)
315+
- Amazon S3 Tables (Iceberg-backed, managed storage): set the connection
316+
``catalog_name`` to ``s3tablescatalog/<table-bucket>`` and use the
317+
namespace as the table ``schema``. The LOCATION clause is omitted since
318+
storage is managed.
308319
- File formats: PARQUET, ORC, TEXTFILE, JSON, AVRO, etc.
309320
- Row formats with SerDe specifications
310321
- Compression settings for various file formats
@@ -445,6 +456,74 @@ def _get_serde_properties_specification(
445456
text.append(")")
446457
return "\n".join(text)
447458

459+
@staticmethod
460+
def _is_s3_tables_catalog(connect_opts: Mapping[str, Any]) -> bool:
461+
"""Return whether the connection targets an Amazon S3 Tables catalog.
462+
463+
S3 Tables are queried by setting the connection ``catalog_name`` to
464+
``s3tablescatalog/<table-bucket>`` and using the namespace as the table
465+
``schema`` (a two-part ``namespace.table`` identifier). Athena rejects a
466+
three-part ``catalog.namespace.table`` identifier in DDL, so the catalog
467+
must be selected at the connection level. Such tables use managed
468+
storage, so their CREATE TABLE statement must omit the LOCATION clause.
469+
470+
Args:
471+
connect_opts: The dialect connection options.
472+
473+
Returns:
474+
True if ``catalog_name`` names an S3 Tables catalog.
475+
"""
476+
if not connect_opts:
477+
return False
478+
catalog = connect_opts.get("catalog_name") or ""
479+
# Athena resolves catalog names case-insensitively.
480+
return catalog.lower().startswith(S3_TABLES_CATALOG_PREFIX)
481+
482+
def _is_iceberg_table(
483+
self, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
484+
) -> bool:
485+
"""Return whether the table properties declare an Iceberg table.
486+
487+
Args:
488+
dialect_opts: The table's ``awsathena_*`` dialect options.
489+
connect_opts: The dialect connection options.
490+
491+
Returns:
492+
True if the rendered TBLPROPERTIES set ``table_type`` to Iceberg.
493+
"""
494+
table_properties = self._get_table_properties_specification(
495+
dialect_opts, connect_opts
496+
).lower()
497+
return ("table_type" in table_properties) and ("iceberg" in table_properties)
498+
499+
def _validate_s3_tables_create_table(
500+
self, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
501+
) -> None:
502+
"""Validate a CREATE TABLE compiled against an S3 Tables catalog.
503+
504+
S3 Tables support only Iceberg tables on managed storage, so the table
505+
must declare ``table_type`` ICEBERG and must not specify a location.
506+
Raising here surfaces a clear client-side error instead of emitting DDL
507+
that Athena would reject.
508+
509+
Args:
510+
dialect_opts: The table's ``awsathena_*`` dialect options.
511+
connect_opts: The dialect connection options.
512+
513+
Raises:
514+
exc.CompileError: If the table is not Iceberg or specifies a location.
515+
"""
516+
if not self._is_iceberg_table(dialect_opts, connect_opts):
517+
raise exc.CompileError(
518+
"S3 Tables support only Iceberg tables; specify the dialect keyword "
519+
"argument `awsathena_tblproperties={'table_type': 'ICEBERG'}`"
520+
)
521+
if dialect_opts["location"]:
522+
raise exc.CompileError(
523+
"S3 Tables use managed storage and do not accept a table location; "
524+
"remove the dialect keyword argument `awsathena_location`"
525+
)
526+
448527
def _get_table_location(
449528
self, table: Table, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
450529
) -> str | None:
@@ -662,12 +741,7 @@ def visit_create_table(self, create: CreateTable, **kwargs) -> str:
662741
dialect = cast("AthenaDialect", self.dialect)
663742
connect_opts = dialect._connect_options
664743

665-
table_properties = self._get_table_properties_specification(
666-
dialect_opts, connect_opts
667-
).lower()
668-
is_iceberg = False
669-
if ("table_type" in table_properties) and ("iceberg" in table_properties):
670-
is_iceberg = True
744+
is_iceberg = self._is_iceberg_table(dialect_opts, connect_opts)
671745

672746
# https://docs.aws.amazon.com/athena/latest/ug/querying-iceberg-creating-tables.html
673747
text = ["\nCREATE TABLE"] if is_iceberg else ["\nCREATE EXTERNAL TABLE"]
@@ -705,11 +779,19 @@ def post_create_table(self, table: Table) -> str:
705779
dialect_opts: _DialectArgDict = table.dialect_options["awsathena"]
706780
dialect = cast("AthenaDialect", self.dialect)
707781
connect_opts = dialect._connect_options
708-
text = [
709-
self._get_row_format_specification(dialect_opts, connect_opts),
710-
self._get_serde_properties_specification(dialect_opts, connect_opts),
711-
self._get_file_format_specification(dialect_opts, connect_opts),
712-
self._get_table_location_specification(table, dialect_opts, connect_opts),
713-
self._get_table_properties_specification(dialect_opts, connect_opts),
714-
]
782+
if self._is_s3_tables_catalog(connect_opts):
783+
# S3 Tables are managed Iceberg tables: ROW FORMAT, SERDEPROPERTIES,
784+
# STORED AS, and LOCATION are not accepted, so emit only TBLPROPERTIES.
785+
self._validate_s3_tables_create_table(dialect_opts, connect_opts)
786+
text = [
787+
self._get_table_properties_specification(dialect_opts, connect_opts),
788+
]
789+
else:
790+
text = [
791+
self._get_row_format_specification(dialect_opts, connect_opts),
792+
self._get_serde_properties_specification(dialect_opts, connect_opts),
793+
self._get_file_format_specification(dialect_opts, connect_opts),
794+
self._get_table_location_specification(table, dialect_opts, connect_opts),
795+
self._get_table_properties_specification(dialect_opts, connect_opts),
796+
]
715797
return "\n".join([t for t in text if t])

scripts/cloudformation/github_actions_oidc.yaml

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,14 @@ Parameters:
2020
ManagedWorkGroupName:
2121
Type: String
2222
Default: "pyathena-managed"
23+
S3TablesBucketName:
24+
Type: String
25+
Default: "laughingman7743-pyathena-s3-tables"
26+
Description: Name of the Amazon S3 Tables table bucket used by the SQLAlchemy S3 Tables tests.
27+
S3TablesNamespaceName:
28+
Type: String
29+
Default: "pyathena"
30+
Description: Namespace created in the S3 Tables table bucket for the SQLAlchemy S3 Tables tests.
2331
OIDCProviderArn:
2432
Type: String
2533
Default: ""
@@ -120,6 +128,30 @@ Resources:
120128
Resource: [
121129
!Sub "arn:aws:s3:::*"
122130
]
131+
- PolicyName: s3tables-access
132+
PolicyDocument:
133+
Version: "2012-10-17"
134+
Statement:
135+
# Amazon S3 Tables data plane: create/query/drop namespaces and
136+
# tables in the table bucket used by the SQLAlchemy S3 Tables tests.
137+
- Effect: Allow
138+
Action: [
139+
"s3tables:*"
140+
]
141+
Resource: [
142+
!Sub "arn:aws:s3tables:${AWS::Region}:${AWS::AccountId}:bucket/*"
143+
]
144+
# Resolve the federated `s3tablescatalog` Glue catalog and obtain
145+
# Lake Formation data access when querying S3 Tables from Athena.
146+
- Effect: Allow
147+
Action: [
148+
"glue:GetCatalog",
149+
"glue:GetCatalogs",
150+
"lakeformation:GetDataAccess"
151+
]
152+
Resource: [
153+
"*"
154+
]
123155
SparkRole:
124156
Type: AWS::IAM::Role
125157
Properties:
@@ -269,6 +301,83 @@ Resources:
269301
ManagedQueryResultsConfiguration:
270302
Enabled: true
271303

304+
# S3 staging bucket for Athena query results and test data. Created outside
305+
# CloudFormation and imported into the stack. Retain on delete/replace so the
306+
# bucket and its data survive stack operations.
307+
# NOTE: Because the bucket is retained, recreating the stack from scratch will
308+
# fail with an AlreadyExists error on this resource; re-import it instead
309+
# (create-change-set --change-set-type IMPORT --resources-to-import ...).
310+
S3StagingBucket:
311+
Type: AWS::S3::Bucket
312+
DeletionPolicy: Retain
313+
UpdateReplacePolicy: Retain
314+
Properties:
315+
BucketName: !Sub "${BucketName}"
316+
PublicAccessBlockConfiguration:
317+
BlockPublicAcls: true
318+
IgnorePublicAcls: true
319+
BlockPublicPolicy: true
320+
RestrictPublicBuckets: true
321+
OwnershipControls:
322+
Rules:
323+
- ObjectOwnership: BucketOwnerEnforced
324+
BucketEncryption:
325+
ServerSideEncryptionConfiguration:
326+
- ServerSideEncryptionByDefault:
327+
SSEAlgorithm: AES256
328+
BucketKeyEnabled: true
329+
LifecycleConfiguration:
330+
Rules:
331+
- Id: delete
332+
Status: Enabled
333+
ExpirationInDays: 1
334+
NoncurrentVersionExpiration:
335+
NoncurrentDays: 1
336+
NewerNoncurrentVersions: 1
337+
AbortIncompleteMultipartUpload:
338+
DaysAfterInitiation: 1
339+
340+
# Amazon S3 Tables infrastructure for the SQLAlchemy S3 Tables tests.
341+
# UnreferencedFileRemoval garbage-collects files left behind by dropped test
342+
# tables; 1 day is the most aggressive setting to keep test storage costs low.
343+
S3TablesBucket:
344+
Type: AWS::S3Tables::TableBucket
345+
Properties:
346+
TableBucketName: !Sub "${S3TablesBucketName}"
347+
UnreferencedFileRemoval:
348+
Status: Enabled
349+
UnreferencedDays: 1
350+
NoncurrentDays: 1
351+
352+
S3TablesNamespace:
353+
Type: AWS::S3Tables::Namespace
354+
Properties:
355+
TableBucketARN: !GetAtt S3TablesBucket.TableBucketARN
356+
Namespace: !Sub "${S3TablesNamespaceName}"
357+
358+
# One-time, per-Region integration: register S3 Tables into the Glue Data
359+
# Catalog as the federated `s3tablescatalog` catalog so Athena can query them.
360+
# IAM_ALLOWED_PRINCIPALS default permissions select IAM-based access control
361+
# (no Lake Formation grants needed; the OIDC role's IAM policy governs access).
362+
S3TablesCatalog:
363+
Type: AWS::Glue::Catalog
364+
Properties:
365+
Name: s3tablescatalog
366+
Description: Federated catalog that maps S3 Tables into the Glue Data Catalog
367+
FederatedCatalog:
368+
Identifier: !Sub "arn:aws:s3tables:${AWS::Region}:${AWS::AccountId}:bucket/*"
369+
ConnectionName: aws:s3tables
370+
CreateDatabaseDefaultPermissions:
371+
- Principal:
372+
DataLakePrincipalIdentifier: IAM_ALLOWED_PRINCIPALS
373+
Permissions:
374+
- ALL
375+
CreateTableDefaultPermissions:
376+
- Principal:
377+
DataLakePrincipalIdentifier: IAM_ALLOWED_PRINCIPALS
378+
Permissions:
379+
- ALL
380+
272381
GithubOidc:
273382
Type: AWS::IAM::OIDCProvider
274383
Condition: CreateOIDCProvider

tests/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,13 @@ def __init__(self):
3131
)
3232
self.default_work_group = os.getenv("AWS_ATHENA_DEFAULT_WORKGROUP", "primary")
3333
self.managed_work_group = os.getenv("AWS_ATHENA_MANAGED_WORKGROUP")
34+
# Optional Amazon S3 Tables configuration for the SQLAlchemy dialect tests.
35+
# `s3tables_catalog` is the registered table-bucket catalog, e.g.
36+
# "s3tablescatalog/<table-bucket>"; `s3tables_namespace` is a namespace that
37+
# already exists in it. These are optional so the suite stays runnable
38+
# without an S3 Tables bucket; the S3 Tables tests skip when they are unset.
39+
self.s3tables_catalog = os.getenv("AWS_ATHENA_S3_TABLES_CATALOG")
40+
self.s3tables_namespace = os.getenv("AWS_ATHENA_S3_TABLES_NAMESPACE")
3441
self.schema = "pyathena_test_" + "".join(
3542
random.choices(string.ascii_lowercase + string.digits, k=10)
3643
)

0 commit comments

Comments
 (0)