From 9ae0bd27e92b4b1a0400873b67d81f8cdaed1ef5 Mon Sep 17 00:00:00 2001 From: sean-eyre Date: Wed, 13 Aug 2025 20:31:49 -0400 Subject: [PATCH 1/3] feat: quicker _get_data_objects query for athena --- sqlmesh/core/engine_adapter/athena.py | 53 +++++++++++++++++++-------- 1 file changed, 37 insertions(+), 16 deletions(-) diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py index abaf7ba281..ce9f8bdb1c 100644 --- a/sqlmesh/core/engine_adapter/athena.py +++ b/sqlmesh/core/engine_adapter/athena.py @@ -103,24 +103,45 @@ def _get_data_objects( """ schema_name = to_schema(schema_name) schema = schema_name.db - query = ( - exp.select( - exp.column("table_catalog").as_("catalog"), - exp.column("table_schema", table="t").as_("schema"), - exp.column("table_name", table="t").as_("name"), - exp.case() - .when( - exp.column("table_type", table="t").eq("BASE TABLE"), - exp.Literal.string("table"), - ) - .else_(exp.column("table_type", table="t")) - .as_("type"), + + base_query = exp.select( + exp.column("table_catalog", table="t").as_("catalog"), + exp.column("table_schema", table="t").as_("schema"), + exp.column("table_name", table="t").as_("name"), + exp.case() + .when( + exp.column("table_type", table="t").eq("BASE TABLE"), + exp.Literal.string("table"), ) - .from_(exp.to_table("information_schema.tables", alias="t")) - .where(exp.column("table_schema", table="t").eq(schema)) - ) + .else_(exp.column("table_type", table="t")) + .as_("type"), + ).from_(exp.to_table("information_schema.tables", alias="t")) + if object_names: - query = query.where(exp.column("table_name", table="t").isin(*object_names)) + # Use CTE with VALUES for better performance when filtering specific tables + query = ( + base_query.join( + exp.to_identifier("object_names"), + on=exp.column("table_name", table="t").eq( + exp.column("name", table="object_names") + ), + ) + .where(exp.column("table_schema", table="t").eq(schema)) + .with_( + "object_names", + exp.select(exp.Literal("name")).from_( + exp.Values( + expressions=[ + exp.Tuple(expressions=[exp.Literal.string(name)]) + for name in object_names + ], + ).as_("t(name)"), + ), + ) + ) + else: + # Simple query when no specific object names are provided + query = base_query.where(exp.column("table_schema", table="t").eq(schema)) df = self.fetchdf(query) From 2bba7bb06b141214dbce6f024e8869357891dbf2 Mon Sep 17 00:00:00 2001 From: sean-eyre Date: Wed, 13 Aug 2025 20:51:24 -0400 Subject: [PATCH 2/3] fix column aliasing --- sqlmesh/core/engine_adapter/athena.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py index ce9f8bdb1c..97c4c89037 100644 --- a/sqlmesh/core/engine_adapter/athena.py +++ b/sqlmesh/core/engine_adapter/athena.py @@ -129,13 +129,16 @@ def _get_data_objects( .where(exp.column("table_schema", table="t").eq(schema)) .with_( "object_names", - exp.select(exp.Literal("name")).from_( + exp.select(exp.column("name")).from_( exp.Values( expressions=[ exp.Tuple(expressions=[exp.Literal.string(name)]) for name in object_names ], - ).as_("t(name)"), + alias=exp.TableAlias( + this="t", columns=[exp.column("name")], + ) + ) ), ) ) From 9fbed7035f505102fe8d3e5d2e09f3b7f8006a38 Mon Sep 17 00:00:00 2001 From: sean-eyre Date: Wed, 13 Aug 2025 20:56:26 -0400 Subject: [PATCH 3/3] run make style --- sqlmesh/core/engine_adapter/athena.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py index 97c4c89037..adab019cc6 100644 --- a/sqlmesh/core/engine_adapter/athena.py +++ b/sqlmesh/core/engine_adapter/athena.py @@ -136,8 +136,9 @@ def _get_data_objects( for name in object_names ], alias=exp.TableAlias( - this="t", columns=[exp.column("name")], - ) + this="t", + columns=[exp.column("name")], + ), ) ), )