diff --git a/great_expectations/execution_engine/sqlalchemy_execution_engine.py b/great_expectations/execution_engine/sqlalchemy_execution_engine.py index cecefab33ede..5d54f5034fec 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: @@ -1003,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( - selectable.columns().subquery() - ) - 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()) @@ -1219,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(selectable.columns().subquery()) - 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) @@ -1370,8 +1399,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..a6c597157703 --- /dev/null +++ b/tests/integration/data_sources_and_expectations/data_sources/test_query_asset_sql_comments.py @@ -0,0 +1,134 @@ +"""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 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 +from tests.integration.test_utils.data_source_config import ( + PostgreSQLDatasourceTestConfig, + SqliteDatasourceTestConfig, +) + +if TYPE_CHECKING: + import pathlib + + 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 + + +@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, +) +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