Skip to content

Commit 64d6e84

Browse files
Support S3 Tables and multi-part identifiers in SQLAlchemy dialect
Athena addresses S3 Tables by a three-part identifier (catalog.namespace.table) whose catalog segment is s3tablescatalog/<table-bucket>, and such tables use managed storage with no LOCATION clause. The dialect previously quoted a dotted schema as a single identifier and always required a LOCATION, so S3 Tables could not be modeled. - Add AthenaBaseIdentifierPreparer.quote_schema to split a dotted schema and quote each part, so catalog.namespace.table round-trips in DDL (backtick) and DML (double-quote). - Omit the LOCATION clause for tables targeting an s3tablescatalog/ catalog instead of raising or building an s3_staging_dir fallback. - Add no-AWS compiler unit tests and gated E2E tests, an S3 Tables docs section, and optional S3 Tables test env vars. Refs #710 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent d2cc374 commit 64d6e84

6 files changed

Lines changed: 315 additions & 3 deletions

File tree

docs/sqlalchemy.md

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,54 @@ 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 with their own data catalog. Athena registers each table bucket as a
370+
catalog named `s3tablescatalog/<table-bucket>`, so an S3 Tables table is addressed by a three-part identifier
371+
`catalog.namespace.table`.
372+
373+
To target an S3 Tables table, set the `schema` to the dotted `s3tablescatalog/<table-bucket>.<namespace>` form.
374+
The dialect quotes the catalog, namespace, and table independently, and because S3 Tables use managed storage it
375+
omits the `LOCATION` clause automatically (so `awsathena_location` is not required).
376+
377+
```python
378+
table = Table(
379+
"some_table",
380+
MetaData(schema="s3tablescatalog/my-bucket.my_namespace"),
381+
Column("id", types.Integer),
382+
Column("dt", types.Date, awsathena_partition=True, awsathena_partition_transform="day"),
383+
awsathena_tblproperties={"table_type": "ICEBERG"},
384+
)
385+
table.create(bind=conn)
386+
```
387+
388+
which builds the following statement:
389+
390+
```sql
391+
CREATE TABLE `s3tablescatalog/my-bucket`.my_namespace.some_table (
392+
id INT,
393+
dt DATE
394+
)
395+
PARTITIONED BY (
396+
day(dt)
397+
)
398+
TBLPROPERTIES (
399+
'table_type' = 'ICEBERG'
400+
)
401+
```
402+
403+
All Iceberg partition transforms (`year`, `month`, `day`, `hour`, `bucket`, `truncate`) are supported, the same as
404+
for other Iceberg tables. CTAS (`CREATE TABLE ... AS SELECT`) is not modeled as a SQLAlchemy construct; issue it as
405+
raw SQL via `text()`, using the double-quoted three-part identifier:
406+
407+
```python
408+
conn.execute(text(
409+
'CREATE TABLE "s3tablescatalog/my-bucket"."my_namespace"."some_table" '
410+
"WITH (table_type = 'ICEBERG') AS SELECT 1 AS id"
411+
))
412+
```
413+
366414
## Temporal/Time-travel with Iceberg
367415

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

pyathena/sqlalchemy/compiler.py

Lines changed: 35 additions & 0 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``). S3 Tables are
44+
# Iceberg-backed and use managed storage, so their CREATE TABLE statements must
45+
# 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,11 @@ 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): targeted by a
316+
three-part ``catalog.namespace.table`` identifier whose catalog segment is
317+
``s3tablescatalog/<table-bucket>``; the LOCATION clause is omitted since
318+
storage is managed. Set the table's ``schema`` to
319+
``s3tablescatalog/<table-bucket>.<namespace>``.
308320
- File formats: PARQUET, ORC, TEXTFILE, JSON, AVRO, etc.
309321
- Row formats with SerDe specifications
310322
- Compression settings for various file formats
@@ -445,6 +457,26 @@ def _get_serde_properties_specification(
445457
text.append(")")
446458
return "\n".join(text)
447459

460+
@staticmethod
461+
def _is_s3_tables(table: Table) -> bool:
462+
"""Return whether the table targets an Amazon S3 Tables catalog.
463+
464+
S3 Tables are addressed by a three-part identifier whose catalog segment
465+
is ``s3tablescatalog/<table-bucket>`` (e.g.
466+
``s3tablescatalog/my-bucket.namespace.table``). Such tables use managed
467+
storage, so their DDL must omit the LOCATION clause.
468+
469+
Args:
470+
table: The table being compiled.
471+
472+
Returns:
473+
True if the table's schema catalog segment is an S3 Tables catalog.
474+
"""
475+
schema = table.schema
476+
if not schema:
477+
return False
478+
return schema.split(".", 1)[0].startswith(S3_TABLES_CATALOG_PREFIX)
479+
448480
def _get_table_location(
449481
self, table: Table, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
450482
) -> str | None:
@@ -466,6 +498,9 @@ def _get_table_location(
466498
def _get_table_location_specification(
467499
self, table: Table, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
468500
) -> str:
501+
if self._is_s3_tables(table):
502+
# S3 Tables use managed storage; CREATE TABLE must not set a LOCATION.
503+
return ""
469504
location = self._get_table_location(table, dialect_opts, connect_opts)
470505
text = []
471506
if location:

pyathena/sqlalchemy/preparer.py

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,46 @@
77
from pyathena.sqlalchemy.constants import DDL_RESERVED_WORDS, SELECT_STATEMENT_RESERVED_WORDS
88

99
if TYPE_CHECKING:
10+
from typing import Any
11+
1012
from sqlalchemy import Dialect
1113

1214

13-
class AthenaDMLIdentifierPreparer(IdentifierPreparer):
15+
class AthenaBaseIdentifierPreparer(IdentifierPreparer):
16+
"""Base identifier preparer shared by Athena's DML and DDL preparers.
17+
18+
Athena identifiers can be three-part (``catalog.namespace.table``), which is
19+
required to target S3 Tables catalogs such as ``s3tablescatalog/<bucket>``.
20+
SQLAlchemy's default :meth:`IdentifierPreparer.quote_schema` treats the whole
21+
schema string as a single identifier, so a dotted schema like
22+
``s3tablescatalog/bucket.ns`` collapses into one quoted token
23+
(catalog and namespace merged) instead of two separately quoted tokens.
24+
25+
See Also:
26+
:class:`AthenaDMLIdentifierPreparer`: Preparer for DML statements.
27+
:class:`AthenaDDLIdentifierPreparer`: Preparer for DDL statements.
28+
"""
29+
30+
def quote_schema(self, schema: str, force: Any = None) -> str:
31+
"""Quote a possibly multi-part (``catalog.namespace``) schema name.
32+
33+
Each dot-separated part is quoted independently so the catalog and
34+
namespace round-trip correctly in both DDL (backtick) and DML
35+
(double-quote) statements. Athena database and namespace names cannot
36+
contain a dot, so the separator is unambiguous; a schema without a dot
37+
is quoted exactly as before.
38+
39+
Args:
40+
schema: The schema name, optionally qualified as ``catalog.namespace``.
41+
force: Unused; kept for SQLAlchemy API compatibility.
42+
43+
Returns:
44+
The quoted, dot-joined schema identifier.
45+
"""
46+
return ".".join(self.quote(part) for part in schema.split("."))
47+
48+
49+
class AthenaDMLIdentifierPreparer(AthenaBaseIdentifierPreparer):
1450
"""Identifier preparer for Athena DML (SELECT, INSERT, etc.) statements.
1551
1652
This preparer handles quoting and escaping of identifiers in DML statements.
@@ -29,7 +65,7 @@ class AthenaDMLIdentifierPreparer(IdentifierPreparer):
2965
reserved_words: set[str] = SELECT_STATEMENT_RESERVED_WORDS
3066

3167

32-
class AthenaDDLIdentifierPreparer(IdentifierPreparer):
68+
class AthenaDDLIdentifierPreparer(AthenaBaseIdentifierPreparer):
3369
"""Identifier preparer for Athena DDL (CREATE, ALTER, DROP) statements.
3470
3571
This preparer handles quoting and escaping of identifiers in DDL statements.

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
)

tests/pyathena/sqlalchemy/test_base.py

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import contextlib
12
import re
23
import textwrap
34
import uuid
@@ -26,6 +27,13 @@
2627
)
2728
from tests.pyathena.conftest import ENV
2829

30+
# Amazon S3 Tables tests need a pre-provisioned table-bucket catalog and namespace.
31+
# Skip them unless AWS_ATHENA_S3_TABLES_CATALOG / AWS_ATHENA_S3_TABLES_NAMESPACE are set.
32+
requires_s3_tables = pytest.mark.skipif(
33+
not ENV.s3tables_catalog or not ENV.s3tables_namespace,
34+
reason="AWS_ATHENA_S3_TABLES_CATALOG / AWS_ATHENA_S3_TABLES_NAMESPACE are not configured",
35+
)
36+
2937

3038
class TestSQLAlchemyAthena:
3139
@pytest.mark.parametrize(
@@ -2019,6 +2027,104 @@ def test_create_iceberg_table_with_multiple_partition_transform(self, engine):
20192027
tblproperties = actual.dialect_options["awsathena"]["tblproperties"]
20202028
assert tblproperties["table_type"] == "ICEBERG"
20212029

2030+
@requires_s3_tables
2031+
def test_create_s3tables_iceberg_table(self, engine):
2032+
engine, conn = engine
2033+
# S3 Tables are addressed by a three-part identifier
2034+
# (catalog.namespace.table) and use managed storage, so the emitted DDL
2035+
# quotes the catalog/namespace/table separately and omits LOCATION.
2036+
schema = f"{ENV.s3tables_catalog}.{ENV.s3tables_namespace}"
2037+
table_name = "test_create_s3tables_iceberg_table"
2038+
table = Table(
2039+
table_name,
2040+
MetaData(schema=schema),
2041+
Column("col_1", types.String),
2042+
Column("col_2", types.Integer),
2043+
awsathena_tblproperties={"table_type": "ICEBERG"},
2044+
)
2045+
ddl = CreateTable(table).compile(bind=conn)
2046+
2047+
assert str(ddl) == textwrap.dedent(
2048+
f"""
2049+
CREATE TABLE `{ENV.s3tables_catalog}`.{ENV.s3tables_namespace}.{table_name} (
2050+
\tcol_1 STRING,
2051+
\tcol_2 INT
2052+
)
2053+
TBLPROPERTIES (
2054+
\t'table_type' = 'ICEBERG'
2055+
)
2056+
"""
2057+
)
2058+
2059+
table.create(bind=conn)
2060+
try:
2061+
# NOTE: SQLAlchemy reflection (autoload_with) of an S3 Tables table is
2062+
# not covered here. The dialect's introspection path passes the dotted
2063+
# schema straight through as the database name and does not split the
2064+
# catalog from the namespace, so reflecting a three-part identifier is
2065+
# a separate, currently-unsupported concern. Verify creation by querying
2066+
# the table directly with its double-quoted three-part identifier.
2067+
fqtn = f'"{ENV.s3tables_catalog}"."{ENV.s3tables_namespace}"."{table_name}"'
2068+
conn.execute(text(f"SELECT col_1, col_2 FROM {fqtn} LIMIT 0"))
2069+
finally:
2070+
with contextlib.suppress(Exception):
2071+
table.drop(bind=conn)
2072+
2073+
@requires_s3_tables
2074+
def test_create_s3tables_iceberg_table_with_partition_transform(self, engine):
2075+
engine, conn = engine
2076+
schema = f"{ENV.s3tables_catalog}.{ENV.s3tables_namespace}"
2077+
table_name = "test_create_s3tables_iceberg_table_with_partition_transform"
2078+
table = Table(
2079+
table_name,
2080+
MetaData(schema=schema),
2081+
Column("col_1", types.String),
2082+
Column("dt", types.Date, awsathena_partition=True, awsathena_partition_transform="day"),
2083+
awsathena_tblproperties={"table_type": "ICEBERG"},
2084+
)
2085+
ddl = CreateTable(table).compile(bind=conn)
2086+
2087+
assert str(ddl) == textwrap.dedent(
2088+
f"""
2089+
CREATE TABLE `{ENV.s3tables_catalog}`.{ENV.s3tables_namespace}.{table_name} (
2090+
\tcol_1 STRING,
2091+
\tdt DATE
2092+
)
2093+
PARTITIONED BY (
2094+
\tday(dt)
2095+
)
2096+
TBLPROPERTIES (
2097+
\t'table_type' = 'ICEBERG'
2098+
)
2099+
"""
2100+
)
2101+
2102+
table.create(bind=conn)
2103+
with contextlib.suppress(Exception):
2104+
table.drop(bind=conn)
2105+
2106+
@requires_s3_tables
2107+
def test_create_s3tables_table_as_select(self, engine):
2108+
engine, conn = engine
2109+
# CTAS is not modeled as a SQLAlchemy construct; exercise it as raw SQL.
2110+
# The three-part identifier uses double quotes (Presto/Trino convention).
2111+
table_name = "test_create_s3tables_table_as_select"
2112+
fqtn = f'"{ENV.s3tables_catalog}"."{ENV.s3tables_namespace}"."{table_name}"'
2113+
conn.execute(text(f"DROP TABLE IF EXISTS {fqtn}"))
2114+
try:
2115+
conn.execute(
2116+
text(
2117+
f"CREATE TABLE {fqtn} "
2118+
"WITH (table_type = 'ICEBERG') AS "
2119+
"SELECT 1 AS id, 'a' AS name"
2120+
)
2121+
)
2122+
rows = conn.execute(text(f"SELECT id, name FROM {fqtn}")).fetchall()
2123+
assert rows == [(1, "a")]
2124+
finally:
2125+
with contextlib.suppress(Exception):
2126+
conn.execute(text(f"DROP TABLE IF EXISTS {fqtn}"))
2127+
20222128
def test_insert_from_select_cte_follows_insert_one(self, engine):
20232129
engine, conn = engine
20242130
metadata = MetaData(schema=ENV.schema)

0 commit comments

Comments
 (0)