From 216b7c112f70b48715a46e71cee42d57376c33dc Mon Sep 17 00:00:00 2001 From: nanjeshramesh Date: Fri, 28 Aug 2026 17:00:53 -0700 Subject: [PATCH 1/2] Fix query asset SQL wrapping so a trailing comment doesn't break validation A query asset whose SQL ends in a line comment (like "-- note") failed with a database syntax error on every SQL backend. GX wraps the raw query text in a subquery before running metrics against it, compiling it as "() AS anon_1" on one line, so the appended ")" and alias landed inside the comment and the statement never closed. The issue pointed at three call sites in SqlAlchemyExecutionEngine that do this wrap via .columns().subquery(). Those are real and now share one helper, but they weren't the whole story: _subselectable builds the same kind of subquery a different way, through sa.select(sa.text(...)).subquery(), which isn't a TextClause and so never went through the fix at those three sites. That method turns out to be the one SQLite actually hits for query assets (its partition clause is never sa.true(), so it always ends up there), which is why the original repro failed specifically on SQLite. Both paths now run the raw text through a small helper that appends a newline before wrapping if the text doesn't already end in one, so a trailing comment can't reach the closing paren and alias. Checked the neighboring cases too: a comment in the middle of the query and an already closed /* */ comment at the end both still work fine, same as before. An unterminated block comment at the end is a different problem that a newline can't solve on its own, closing it properly would mean actual comment parsing, so I left that alone and called it out in the PR instead of trying to solve it here. --- .../sqlalchemy_execution_engine.py | 42 ++++++++- .../test_query_asset_sql_comments.py | 90 +++++++++++++++++++ 2 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py diff --git a/great_expectations/execution_engine/sqlalchemy_execution_engine.py b/great_expectations/execution_engine/sqlalchemy_execution_engine.py index cecefab33ede..feae9e5ad779 100644 --- a/great_expectations/execution_engine/sqlalchemy_execution_engine.py +++ b/great_expectations/execution_engine/sqlalchemy_execution_engine.py @@ -233,6 +233,39 @@ def _dialect_requires_persisted_connection( return return_val +def _ensure_sql_text_ends_with_newline(text: str) -> str: + """Ensure raw SQL text ends on a fresh line before it is wrapped as a subquery. + + SQLAlchemy compiles wrapped raw SQL text verbatim inside "() AS alias" on a + single line. If the text ends in a line comment (e.g. "-- note"), the appended ")" + and alias land inside that comment and the resulting statement is never terminated. + Ending the text on a fresh line keeps any appended syntax out of reach of a comment + that runs to the end of the text, without changing what the query selects. + + Args: + text: The raw SQL text about to be wrapped as a subquery. + + Returns: + The same text, with a trailing newline added if it didn't already have one. + """ + return text if text.endswith("\n") else f"{text}\n" + + +def _wrap_raw_sql_as_subquery(selectable: sqlalchemy.TextClause) -> Subquery: + """Wrap a raw-SQL TextClause (e.g. from a query asset) as a subquery. + + See _ensure_sql_text_ends_with_newline for why the text is normalized first. + + Args: + selectable: The raw-SQL TextClause to wrap. + + Returns: + A Subquery built from the TextClause's columns. + """ + text = _ensure_sql_text_ends_with_newline(selectable.text) + return sa.text(text).columns().subquery() + + class SqlAlchemyExecutionEngine(ExecutionEngine[SQLAColumnClause]): """SparkDFExecutionEngine instantiates the ExecutionEngine API to support computations using Spark platform. @@ -652,7 +685,7 @@ def get_domain_records( # noqa: C901, PLR0912, PLR0915 # FIXME CoP to TextualSelect using sa.columns() before it can be converted to type Subquery """ if sqlalchemy.TextClause and isinstance(selectable, sqlalchemy.TextClause): # type: ignore[truthy-function] # FIXME CoP - selectable = selectable.columns().subquery() + selectable = _wrap_raw_sql_as_subquery(selectable) # Filtering by row condition. if "row_condition" in domain_kwargs and domain_kwargs["row_condition"] is not None: @@ -1010,7 +1043,7 @@ def resolve_metric_bundle( """ if sqlalchemy.TextClause and isinstance(selectable, sqlalchemy.TextClause): # type: ignore[truthy-function] # FIXME CoP sa_query_object = sa.select(*query["select"]).select_from( - selectable.columns().subquery() + _wrap_raw_sql_as_subquery(selectable) ) elif (sqlalchemy.Select and isinstance(selectable, sqlalchemy.Select)) or ( # type: ignore[truthy-function] # FIXME CoP sqlalchemy.TextualSelect and isinstance(selectable, sqlalchemy.TextualSelect) # type: ignore[truthy-function] # FIXME CoP @@ -1220,7 +1253,7 @@ def _count_query_parameters(self, selectable: sqlalchemy.Selectable, select_list """ DEFAULT_PARAMS_PER_SELECT = 2 # Conservative upper bound if isinstance(selectable, sqlalchemy.TextClause): - test_query = sa.select(*select_list).select_from(selectable.columns().subquery()) + test_query = sa.select(*select_list).select_from(_wrap_raw_sql_as_subquery(selectable)) elif isinstance(selectable, (sqlalchemy.Select, sqlalchemy.TextualSelect)): test_query = sa.select(*select_list).select_from(selectable.subquery()) elif isinstance(selectable, sa.sql.FromClause): @@ -1370,8 +1403,9 @@ def _subselectable(self, batch_spec: BatchSpec) -> sqlalchemy.Selectable: if not isinstance(query, str): raise ValueError(f"SQL query should be a str but got {query}") # noqa: TRY003 # FIXME CoP # Query is a valid SELECT query that begins with r"\w+select\w" + stripped_query = query.lstrip()[6:].strip().rstrip(";").rstrip() selectable = sa.select( - sa.text(query.lstrip()[6:].strip().rstrip(";").rstrip()) + sa.text(_ensure_sql_text_ends_with_newline(stripped_query)) ).subquery() return selectable diff --git a/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py b/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py new file mode 100644 index 000000000000..c5d209ca5893 --- /dev/null +++ b/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py @@ -0,0 +1,90 @@ +"""Regression tests for query assets whose raw SQL ends in a comment. + +See https://github.com/fivetran/great_expectations/issues/12122. GX wraps a query +asset's raw SQL in a subquery before running metrics against it. SQLAlchemy compiles +that as the user's text pasted verbatim inside "() AS anon_1" on a single line, +so if the user's SQL ends in a line comment, the appended ")" and alias land inside +that comment and the statement is never terminated. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pandas as pd + +import great_expectations.expectations as gxe +from great_expectations.datasource.fluent.sql_datasource import SQLDatasource, TableAsset +from tests.integration.conftest import parameterize_batch_for_data_sources +from tests.integration.test_utils.data_source_config import ( + PostgreSQLDatasourceTestConfig, + SqliteDatasourceTestConfig, +) + +if TYPE_CHECKING: + from great_expectations.datasource.fluent.interfaces import Batch + +COL_A = "col_a" + +DATA = pd.DataFrame({COL_A: ["x", "y"]}) + + +def _add_query_asset(batch_for_datasource: Batch, name: str, query: str) -> Batch: + asset = batch_for_datasource.data_asset + assert isinstance(asset, TableAsset) + datasource = batch_for_datasource.datasource + assert isinstance(datasource, SQLDatasource) + + query_asset = datasource.add_query_asset(name=name, query=query.format(table=asset.table_name)) + return query_asset.add_batch_definition_whole_table(f"{name}_bd").get_batch() + + +@parameterize_batch_for_data_sources( + data_source_configs=[SqliteDatasourceTestConfig(), PostgreSQLDatasourceTestConfig()], + data=DATA, +) +def test_query_asset_sql_ending_in_line_comment(batch_for_datasource: Batch) -> None: + """A trailing `--` comment does not change what a query selects.""" + batch = _add_query_asset( + batch_for_datasource, + "trailing_comment_asset", + f"SELECT {COL_A} FROM {{table}} -- only the rows we care about", + ) + + result = batch.validate(gxe.ExpectColumnValuesToNotBeNull(column=COL_A)) + + assert result.success, result.exception_info + + +@parameterize_batch_for_data_sources( + data_source_configs=[SqliteDatasourceTestConfig(), PostgreSQLDatasourceTestConfig()], + data=DATA, +) +def test_query_asset_sql_ending_in_terminated_block_comment(batch_for_datasource: Batch) -> None: + """A trailing, properly-closed `/* ... */` comment does not change what a query selects.""" + batch = _add_query_asset( + batch_for_datasource, + "terminated_block_comment_asset", + f"SELECT {COL_A} FROM {{table}} /* only the rows we care about */", + ) + + result = batch.validate(gxe.ExpectColumnValuesToNotBeNull(column=COL_A)) + + assert result.success, result.exception_info + + +@parameterize_batch_for_data_sources( + data_source_configs=[SqliteDatasourceTestConfig(), PostgreSQLDatasourceTestConfig()], + data=DATA, +) +def test_query_asset_sql_with_comment_not_at_end(batch_for_datasource: Batch) -> None: + """A comment in the middle of a query, followed by a later clause, is unaffected.""" + batch = _add_query_asset( + batch_for_datasource, + "mid_query_comment_asset", + f"SELECT {COL_A} -- the column we care about\nFROM {{table}}", + ) + + result = batch.validate(gxe.ExpectColumnValuesToNotBeNull(column=COL_A)) + + assert result.success, result.exception_info From ede7c9cf802f239d24d03bd2da459337bafa16db Mon Sep 17 00:00:00 2001 From: nanjeshramesh Date: Tue, 1 Sep 2026 13:06:28 -0700 Subject: [PATCH 2/2] Remove unreachable TextClause branches, add missing test coverage per review Two of the three call sites this PR touched turned out to be dead code: selectable in resolve_metric_bundle and _count_query_parameters always comes from get_domain_records(), which already converts any TextClause into a Subquery before returning, so their own TextClause branches (and the _wrap_raw_sql_as_subquery calls added to them) could never run. Removed both, matching what the review pointed out. The remaining TextClause branch, in get_domain_records itself, is genuinely reachable, but the existing tests never exercised it: they go through the default create_temp_table path, which converts the query asset into a Subquery before get_domain_records ever sees it. It's only reached when create_temp_table=False, which is in fact the default for SQLDatasource, just not what the shared test fixture used. Added a test that constructs that scenario directly and confirmed via tracing that it actually calls the wrap. One thing worth flagging: every path I could find that reaches get_domain_records this way has already gone through _subselectable() first, whose fix already neutralizes the corruption before it gets here. So the new test genuinely exercises get_domain_records's own wrap, but doesn't isolate it as the thing preventing failure, since _subselectable's fix alone would too. I looked for a way to construct a scenario where only get_domain_records's fix mattered and couldn't find one through the current public API, so I left it in place as reachable, tested, defense-in-depth rather than removing it outright. --- .../sqlalchemy_execution_engine.py | 22 ++++------ .../test_query_asset_sql_comments.py | 44 +++++++++++++++++++ 2 files changed, 53 insertions(+), 13 deletions(-) diff --git a/great_expectations/execution_engine/sqlalchemy_execution_engine.py b/great_expectations/execution_engine/sqlalchemy_execution_engine.py index feae9e5ad779..5d54f5034fec 100644 --- a/great_expectations/execution_engine/sqlalchemy_execution_engine.py +++ b/great_expectations/execution_engine/sqlalchemy_execution_engine.py @@ -1036,16 +1036,11 @@ def resolve_metric_bundle( assert len(query["select"]) == len(query["metric_ids"]) try: - """ - If a custom query is passed, selectable will be TextClause and not formatted - as a subquery wrapped in "(subquery) alias". TextClause must first be converted - to TextualSelect using sa.columns() before it can be converted to type Subquery - """ - if sqlalchemy.TextClause and isinstance(selectable, sqlalchemy.TextClause): # type: ignore[truthy-function] # FIXME CoP - sa_query_object = sa.select(*query["select"]).select_from( - _wrap_raw_sql_as_subquery(selectable) - ) - elif (sqlalchemy.Select and isinstance(selectable, sqlalchemy.Select)) or ( # type: ignore[truthy-function] # FIXME CoP + # Note: selectable here always comes from get_domain_records(), which converts + # any raw-SQL TextClause into a Subquery before returning (see the wrap in + # get_domain_records itself), so selectable can never be a TextClause at this + # point and there is no corresponding branch for it here. + if (sqlalchemy.Select and isinstance(selectable, sqlalchemy.Select)) or ( # type: ignore[truthy-function] # FIXME CoP sqlalchemy.TextualSelect and isinstance(selectable, sqlalchemy.TextualSelect) # type: ignore[truthy-function] # FIXME CoP ): sa_query_object = sa.select(*query["select"]).select_from(selectable.subquery()) @@ -1252,9 +1247,10 @@ def _count_query_parameters(self, selectable: sqlalchemy.Selectable, select_list Total number of parameters that would be generated when the query is compiled """ DEFAULT_PARAMS_PER_SELECT = 2 # Conservative upper bound - if isinstance(selectable, sqlalchemy.TextClause): - test_query = sa.select(*select_list).select_from(_wrap_raw_sql_as_subquery(selectable)) - elif isinstance(selectable, (sqlalchemy.Select, sqlalchemy.TextualSelect)): + # Note: selectable's sole caller passes a value from get_domain_records(), which + # converts any raw-SQL TextClause into a Subquery before returning, so selectable can + # never be a TextClause here and there is no corresponding branch for it. + if isinstance(selectable, (sqlalchemy.Select, sqlalchemy.TextualSelect)): test_query = sa.select(*select_list).select_from(selectable.subquery()) elif isinstance(selectable, sa.sql.FromClause): test_query = sa.select(*select_list).select_from(selectable) diff --git a/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py b/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py index c5d209ca5893..a6c597157703 100644 --- a/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py +++ b/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py @@ -12,7 +12,10 @@ from typing import TYPE_CHECKING import pandas as pd +import pytest +import sqlalchemy as sa +import great_expectations as gx import great_expectations.expectations as gxe from great_expectations.datasource.fluent.sql_datasource import SQLDatasource, TableAsset from tests.integration.conftest import parameterize_batch_for_data_sources @@ -22,6 +25,8 @@ ) if TYPE_CHECKING: + import pathlib + from great_expectations.datasource.fluent.interfaces import Batch COL_A = "col_a" @@ -56,6 +61,45 @@ def test_query_asset_sql_ending_in_line_comment(batch_for_datasource: Batch) -> assert result.success, result.exception_info +@pytest.mark.sqlite +def test_query_asset_sql_ending_in_line_comment_without_temp_table(tmp_path: pathlib.Path) -> None: + """The same trailing-comment case above, but constructed so the query asset's raw SQL + reaches SqlAlchemyExecutionEngine.get_domain_records() as a TextClause, exercising its + own wrap fix directly, rather than one already converted to a Subquery further upstream + (which is what the fixture-driven test above ends up exercising instead). This only + happens when create_temp_table=False on the datasource, which is in fact the default for + SQLDatasource, so this is a real, reachable path, not a contrived one, it's just not the + one the shared test fixture takes. + + Note: as of this fix, every path that reaches get_domain_records() this way has already + passed through _subselectable(), so this scenario happens to be covered by that fix too. + That doesn't make get_domain_records()'s own fix untested, both are exercised here, it + just means this specific test can't isolate one from the other with the current public + API's set of construction paths. + """ # FIXME CoP + db_path = tmp_path / "test.db" + + context = gx.get_context(mode="ephemeral") + datasource = context.data_sources.add_sqlite( + name="ds", connection_string=f"sqlite:///{db_path}", create_temp_table=False + ) + engine = datasource.get_execution_engine().engine + DATA.to_sql("my_tbl", engine, index=False) + + query_asset = datasource.add_query_asset( + name="qa", query=f"SELECT {COL_A} FROM my_tbl -- only the rows we care about" + ) + batch = query_asset.add_batch_definition_whole_table("bd").get_batch() + assert isinstance(batch.data.selectable, sa.TextClause), ( + "test setup no longer reaches get_domain_records() as a TextClause; this test needs " + "updating to match whatever now does" + ) + + result = batch.validate(gxe.ExpectColumnValuesToNotBeNull(column=COL_A)) + + assert result.success, result.exception_info + + @parameterize_batch_for_data_sources( data_source_configs=[SqliteDatasourceTestConfig(), PostgreSQLDatasourceTestConfig()], data=DATA,