From f6ab9355ac3af0c218e3124e81ae90f4b1e4447c Mon Sep 17 00:00:00 2001 From: kshitij1439 Date: Tue, 28 Jul 2026 02:19:56 +0530 Subject: [PATCH 1/3] fix(profiler): sanitize BigQuery uniqueCount SQL for TIMESTAMP/DATE and STRUCT-subfield columns (#30152) --- .../sqlalchemy/columnValuesToBeUnique.py | 5 +- .../sqlalchemy/profiler_interface.py | 3 +- .../profiler/metrics/static/unique_count.py | 9 +- .../bigquery/test_unique_count_bigquery.py | 221 ++++++++++++++++++ 4 files changed, 230 insertions(+), 8 deletions(-) create mode 100644 ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py diff --git a/ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py b/ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py index 495bdded5c84..5a01f7667b28 100644 --- a/ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py +++ b/ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py @@ -38,6 +38,7 @@ from metadata.generated.schema.entity.data.table import TableData from metadata.generated.schema.tests.dimensionResult import DimensionResult from metadata.profiler.metrics.registry import Metrics +from metadata.profiler.metrics.static.unique_count import VALUE_COUNT_ALIAS from metadata.profiler.orm.functions.unique_count import _unique_count_dimensional_cte from metadata.profiler.orm.registry import Dialects from metadata.profiler.processor.runner import QueryRunner @@ -75,7 +76,7 @@ def _run_results(self, metric: Metrics, column: Column) -> Optional[int]: # noq """ count = Metrics.valuesCount.value(column).fn() grouped_cte = ( - select(count.label(column.name)).select_from(self.runner.dataset).group_by(column).cte("grouped_cte") # type: ignore + select(count.label(VALUE_COUNT_ALIAS)).select_from(self.runner.dataset).group_by(column).cte("grouped_cte") # type: ignore ) unique_count = Metrics.uniqueCount.value(column).query( sample=self.runner.dataset, @@ -90,7 +91,7 @@ def _run_results(self, metric: Metrics, column: Column) -> Optional[int]: # noq row = self.runner._select_from_dataset( grouped_cte, - func.sum(grouped_cte.c[column.name]).label(Metrics.valuesCount.name), + func.sum(grouped_cte.c[VALUE_COUNT_ALIAS]).label(Metrics.valuesCount.name), unique_count.label(Metrics.uniqueCount.name), query_group_by_=query_group_by_, ).first() diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py index 96e546b5457a..93b0f9e870a1 100644 --- a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py @@ -54,6 +54,7 @@ from metadata.profiler.metrics.static.mean import Mean from metadata.profiler.metrics.static.stddev import StdDev from metadata.profiler.metrics.static.sum import Sum +from metadata.profiler.metrics.static.unique_count import VALUE_COUNT_ALIAS from metadata.profiler.metrics.system.system import System, SystemMetricsRegistry from metadata.profiler.orm.functions.table_metric_computer import TableMetricComputer from metadata.profiler.orm.registry import Dialects @@ -290,7 +291,7 @@ def _compute_query_metrics( # hotfix to handle transition of unique count implementation sample_column = sample.__table__.c[column.key] if hasattr(sample, "__table__") else sample.c[column.key] subquery = ( - self.session.query(Count(sample_column).fn().label(column.name)) + self.session.query(Count(sample_column).fn().label(VALUE_COUNT_ALIAS)) .select_from(sample) .group_by(sample_column) .subquery() diff --git a/ingestion/src/metadata/profiler/metrics/static/unique_count.py b/ingestion/src/metadata/profiler/metrics/static/unique_count.py index 6c4906e1bb21..cbccff78ed36 100644 --- a/ingestion/src/metadata/profiler/metrics/static/unique_count.py +++ b/ingestion/src/metadata/profiler/metrics/static/unique_count.py @@ -17,7 +17,7 @@ from collections import Counter from typing import TYPE_CHECKING, Optional -from sqlalchemy import column, func +from sqlalchemy import Integer, column, func from sqlalchemy.orm import Session from metadata.generated.schema.configuration.profilerConfiguration import MetricType @@ -34,6 +34,8 @@ logger = profiler_logger() +VALUE_COUNT_ALIAS = "value_count" + class UniqueCount(QueryMetric): """ @@ -67,10 +69,7 @@ def query(self, sample: Optional[type], session: Optional[Session] = None): # n # TODO: Move all connectors from subquery to COUNT(IF) or COUNTIF for performance if session.get_bind().dialect.name == Dialects.BigQuery: - # We are querying against the subquery output (which is a COUNT), so the type is numeric. - # Use an untyped column to avoid passing the original metric type (like STRING or BYTES) into the COUNTIF comparison. - count_col = column(col.name) - return func.countif(count_col == 1).label(self.name()) + return func.countif(column(VALUE_COUNT_ALIAS, Integer()) == 1).label(self.name()) unique_count_query = _unique_count_query_mapper[session.get_bind().dialect.name](col, session, sample) only_once_sub = unique_count_query.subquery("only_once") diff --git a/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py new file mode 100644 index 000000000000..47c82175957a --- /dev/null +++ b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py @@ -0,0 +1,221 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Test UniqueCount metric and validator for BigQuery dialect +""" + +from unittest.mock import Mock, patch + +from sqlalchemy import Column, Date, DateTime, Integer, String, create_engine +from sqlalchemy.orm import DeclarativeBase, Session +from sqlalchemy_bigquery.base import BigQueryDialect + +from metadata.data_quality.validations.column.sqlalchemy.columnValuesToBeUnique import ( + ColumnValuesToBeUniqueValidator, +) +from metadata.profiler.interface.sqlalchemy.profiler_interface import ( + SQAProfilerInterface, +) +from metadata.profiler.metrics.registry import Metrics +from metadata.profiler.metrics.static.unique_count import VALUE_COUNT_ALIAS, UniqueCount +from metadata.profiler.orm.registry import Dialects +from metadata.profiler.processor.runner import QueryRunner + + +class Base(DeclarativeBase): + pass + + +class Orders(Base): + __tablename__ = "orders" + id = Column(Integer, primary_key=True) + created_at = Column("created_at", DateTime, key="created_at") + created_date = Column("created_date", Date, key="created_date") + customer_email = Column("customer.email", String(256), key="customer_email") + + +def test_unique_count_bigquery_timestamp_and_struct(): + """Test BigQuery uniqueCount SQL generation via SQAProfilerInterface for TIMESTAMP, DATE, and nested STRUCT columns""" + dialect = BigQueryDialect() + assert dialect.name == Dialects.BigQuery + + engine = create_engine("sqlite:///:memory:") + session = Session(engine) + + mock_bind = Mock() + mock_bind.dialect = dialect + mock_session = Mock(spec=Session) + mock_session.get_bind.return_value = mock_bind + mock_session.query = session.query + + runner = QueryRunner( + session=mock_session, + dataset=Orders, + raw_dataset=Orders, + partition_details=None, + profile_sample_query=None, + ) + + interface = SQAProfilerInterface.__new__(SQAProfilerInterface) + interface.session = mock_session + + # 1. Test TIMESTAMP column created_at via profiler_interface + col_timestamp = Orders.__table__.c["created_at"] + + captured_query_ts = None + + def mock_select_ts(q): + nonlocal captured_query_ts + captured_query_ts = q + mock_row = Mock() + mock_row._asdict.return_value = {Metrics.uniqueCount.name: 1} + return mock_row + + with patch.object(runner, "select_first_from_query", side_effect=mock_select_ts): + interface._compute_query_metrics( + metric=UniqueCount, + runner=runner, + column=col_timestamp, + session=mock_session, + sample=Orders, + ) + + sql_ts = str( + captured_query_ts.statement.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + ) + + # Verify outer countif uses value_count integer alias and literal 1 (no timestamp type mismatch) + assert "countif(`value_count` = 1)" in sql_ts + assert "timestamp" not in sql_ts.lower() + + # 2. Test DATE column created_date via profiler_interface + col_date = Orders.__table__.c["created_date"] + + captured_query_date = None + + def mock_select_date(q): + nonlocal captured_query_date + captured_query_date = q + mock_row = Mock() + mock_row._asdict.return_value = {Metrics.uniqueCount.name: 1} + return mock_row + + with patch.object(runner, "select_first_from_query", side_effect=mock_select_date): + interface._compute_query_metrics( + metric=UniqueCount, + runner=runner, + column=col_date, + session=mock_session, + sample=Orders, + ) + + sql_date = str( + captured_query_date.statement.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + ) + + # Verify outer countif uses value_count integer alias and literal 1 (no date type mismatch) + assert "countif(`value_count` = 1)" in sql_date + assert "SELECT count(`created_date`) AS `value_count`" in sql_date + + # 3. Test nested STRUCT column customer.email via profiler_interface + col_struct = Orders.__table__.c["customer_email"] + + captured_query_struct = None + + def mock_select_struct(q): + nonlocal captured_query_struct + captured_query_struct = q + mock_row = Mock() + mock_row._asdict.return_value = {Metrics.uniqueCount.name: 1} + return mock_row + + with patch.object(runner, "select_first_from_query", side_effect=mock_select_struct): + interface._compute_query_metrics( + metric=UniqueCount, + runner=runner, + column=col_struct, + session=mock_session, + sample=Orders, + ) + + sql_struct = str( + captured_query_struct.statement.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + ) + + # Verify outer expression isolates countif(`value_count` = 1) without dotted path in outer query + assert "countif(`value_count` = 1) AS `uniqueCount`" in sql_struct + assert "SELECT count(`customer`.`email`) AS `value_count`" in sql_struct + assert "GROUP BY `orders`.`customer`.`email`" in sql_struct + + +def test_column_values_to_be_unique_validator_bigquery(): + """Test BigQuery SQL generation for ColumnValuesToBeUniqueValidator""" + dialect = BigQueryDialect() + assert dialect.name == Dialects.BigQuery + + engine = create_engine("sqlite:///:memory:") + session = Session(engine) + + mock_bind = Mock() + mock_bind.dialect = dialect + mock_session = Mock(spec=Session) + mock_session.get_bind.return_value = mock_bind + mock_session.query = session.query + + validator = ColumnValuesToBeUniqueValidator.__new__(ColumnValuesToBeUniqueValidator) + validator.runner = Mock() + validator.runner.dataset = Orders + validator.runner._session = mock_session + validator.runner.dialect = Dialects.BigQuery + + col_struct = Orders.__table__.c["customer_email"] + + captured_args = None + + def mock_select_from_dataset(grouped_cte, *entities, **kwargs): + nonlocal captured_args + captured_args = (grouped_cte, entities) + mock_res = Mock() + mock_row = Mock() + mock_row._mapping = {Metrics.valuesCount.name: 3, Metrics.uniqueCount.name: 3} + mock_res.first.return_value = mock_row + return mock_res + + validator.runner._select_from_dataset = mock_select_from_dataset + + validator._run_results(Metrics.uniqueCount, col_struct) + + grouped_cte, entities = captured_args + + # Check CTE definition uses VALUE_COUNT_ALIAS + cte_sql = str( + grouped_cte.element.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + ) + assert "SELECT count(`customer`.`email`) AS `value_count`" in cte_sql + + # Note: unique_count expression is wrapped with .label(Metrics.uniqueCount.name) + # in validator._run_results for result mapping keys. + unique_count_expr = entities[1] + expr_sql = str( + unique_count_expr.element.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + ) + assert "countif(`value_count` = 1)" in expr_sql From 0ce5ca41d3e82a0ab207028cec82735652858fef Mon Sep 17 00:00:00 2001 From: kshitij1439 Date: Tue, 28 Jul 2026 02:38:40 +0530 Subject: [PATCH 2/3] fix(profiler): resolve unused VALUE_COUNT_ALIAS import in BigQuery test --- .../bigquery/test_unique_count_bigquery.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py index 47c82175957a..d216a6a2118b 100644 --- a/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py +++ b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py @@ -96,7 +96,7 @@ def mock_select_ts(q): ) # Verify outer countif uses value_count integer alias and literal 1 (no timestamp type mismatch) - assert "countif(`value_count` = 1)" in sql_ts + assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in sql_ts assert "timestamp" not in sql_ts.lower() # 2. Test DATE column created_date via profiler_interface @@ -127,8 +127,8 @@ def mock_select_date(q): ) # Verify outer countif uses value_count integer alias and literal 1 (no date type mismatch) - assert "countif(`value_count` = 1)" in sql_date - assert "SELECT count(`created_date`) AS `value_count`" in sql_date + assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in sql_date + assert f"SELECT count(`created_date`) AS `{VALUE_COUNT_ALIAS}`" in sql_date # 3. Test nested STRUCT column customer.email via profiler_interface col_struct = Orders.__table__.c["customer_email"] @@ -158,8 +158,8 @@ def mock_select_struct(q): ) # Verify outer expression isolates countif(`value_count` = 1) without dotted path in outer query - assert "countif(`value_count` = 1) AS `uniqueCount`" in sql_struct - assert "SELECT count(`customer`.`email`) AS `value_count`" in sql_struct + assert f"countif(`{VALUE_COUNT_ALIAS}` = 1) AS `uniqueCount`" in sql_struct + assert f"SELECT count(`customer`.`email`) AS `{VALUE_COUNT_ALIAS}`" in sql_struct assert "GROUP BY `orders`.`customer`.`email`" in sql_struct @@ -208,7 +208,7 @@ def mock_select_from_dataset(grouped_cte, *entities, **kwargs): dialect=dialect, compile_kwargs={"literal_binds": True} ) ) - assert "SELECT count(`customer`.`email`) AS `value_count`" in cte_sql + assert f"SELECT count(`customer`.`email`) AS `{VALUE_COUNT_ALIAS}`" in cte_sql # Note: unique_count expression is wrapped with .label(Metrics.uniqueCount.name) # in validator._run_results for result mapping keys. @@ -218,4 +218,4 @@ def mock_select_from_dataset(grouped_cte, *entities, **kwargs): dialect=dialect, compile_kwargs={"literal_binds": True} ) ) - assert "countif(`value_count` = 1)" in expr_sql + assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in expr_sql From ddb9819f5b76c632e3935ee762a7866ad4ef1085 Mon Sep 17 00:00:00 2001 From: kshitij1439 Date: Wed, 29 Jul 2026 18:54:03 +0530 Subject: [PATCH 3/3] fix(profiler): gracefully skip BigQuery uniqueCount tests if sqlalchemy-bigquery not installed --- .../bigquery/test_unique_count_bigquery.py | 37 ++++++------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py index d216a6a2118b..655e13528e27 100644 --- a/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py +++ b/ingestion/tests/unit/observability/profiler/sqlalchemy/bigquery/test_unique_count_bigquery.py @@ -13,6 +13,13 @@ Test UniqueCount metric and validator for BigQuery dialect """ +import pytest + +pytest.importorskip( + "sqlalchemy_bigquery", + reason="sqlalchemy-bigquery not installed — skipping BigQuery SQL compilation tests", +) + from unittest.mock import Mock, patch from sqlalchemy import Column, Date, DateTime, Integer, String, create_engine @@ -89,11 +96,7 @@ def mock_select_ts(q): sample=Orders, ) - sql_ts = str( - captured_query_ts.statement.compile( - dialect=dialect, compile_kwargs={"literal_binds": True} - ) - ) + sql_ts = str(captured_query_ts.statement.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) # Verify outer countif uses value_count integer alias and literal 1 (no timestamp type mismatch) assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in sql_ts @@ -120,11 +123,7 @@ def mock_select_date(q): sample=Orders, ) - sql_date = str( - captured_query_date.statement.compile( - dialect=dialect, compile_kwargs={"literal_binds": True} - ) - ) + sql_date = str(captured_query_date.statement.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) # Verify outer countif uses value_count integer alias and literal 1 (no date type mismatch) assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in sql_date @@ -151,11 +150,7 @@ def mock_select_struct(q): sample=Orders, ) - sql_struct = str( - captured_query_struct.statement.compile( - dialect=dialect, compile_kwargs={"literal_binds": True} - ) - ) + sql_struct = str(captured_query_struct.statement.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) # Verify outer expression isolates countif(`value_count` = 1) without dotted path in outer query assert f"countif(`{VALUE_COUNT_ALIAS}` = 1) AS `uniqueCount`" in sql_struct @@ -203,19 +198,11 @@ def mock_select_from_dataset(grouped_cte, *entities, **kwargs): grouped_cte, entities = captured_args # Check CTE definition uses VALUE_COUNT_ALIAS - cte_sql = str( - grouped_cte.element.compile( - dialect=dialect, compile_kwargs={"literal_binds": True} - ) - ) + cte_sql = str(grouped_cte.element.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) assert f"SELECT count(`customer`.`email`) AS `{VALUE_COUNT_ALIAS}`" in cte_sql # Note: unique_count expression is wrapped with .label(Metrics.uniqueCount.name) # in validator._run_results for result mapping keys. unique_count_expr = entities[1] - expr_sql = str( - unique_count_expr.element.compile( - dialect=dialect, compile_kwargs={"literal_binds": True} - ) - ) + expr_sql = str(unique_count_expr.element.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) assert f"countif(`{VALUE_COUNT_ALIAS}` = 1)" in expr_sql