[FEATURE] infer primary keys during column_types metric fetch (#11554)

Author: shiplet

great_expectations/expectations/metrics/util.py

@@ -376,87 +376,132 @@ def __getitem__(self, key: Any) -> Any:
         return item
 
 
-def get_sqlalchemy_column_metadata(  # noqa: C901 # FIXME CoP
+def get_sqlalchemy_column_metadata(
     execution_engine: SqlAlchemyExecutionEngine,
     table_selectable: sqlalchemy.Select,
     schema_name: Optional[str] = None,
 ) -> Sequence[Mapping[str, Any]] | None:
     try:
-        columns: Sequence[Dict[str, Any]]
-
         engine = execution_engine.engine
         inspector = execution_engine.get_inspector()
-        try:
-            # if a custom query was passed
-            if sqlalchemy.TextClause and isinstance(table_selectable, sqlalchemy.TextClause):  # type: ignore[truthy-function]
-                if hasattr(table_selectable, "selected_columns"):
-                    # New in version 1.4.
-                    columns = table_selectable.selected_columns.columns
-                else:
-                    # Implicit subquery for columns().column was deprecated in SQLAlchemy 1.4
-                    # We must explicitly create a subquery
-                    columns = table_selectable.columns().subquery().columns
-            elif sqlalchemy.quoted_name and isinstance(table_selectable, sqlalchemy.quoted_name):  # type: ignore[truthy-function]
-                columns = inspector.get_columns(
-                    table_name=table_selectable,
-                    schema=schema_name,
-                )
-            else:
-                logger.warning("unexpected table_selectable type")
-                columns = inspector.get_columns(  # type: ignore[assignment]
-                    table_name=str(table_selectable),
+
+        # Determine selectable type once
+        is_text_clause = sqlalchemy.TextClause and isinstance(  # type: ignore[truthy-function]
+            table_selectable, sqlalchemy.TextClause
+        )
+        is_quoted_name = sqlalchemy.quoted_name and isinstance(  # type: ignore[truthy-function]
+            table_selectable, sqlalchemy.quoted_name
+        )
+        table_name = str(table_selectable)
+
+        # Fetch primary key info (skip for custom queries/TextClause)
+        primary_key_columns: set[str] = set()
+        if not is_text_clause:
+            try:
+                pk_constraint = inspector.get_pk_constraint(
+                    table_name=table_name,
                     schema=schema_name,
                 )
-        except (
-            KeyError,
-            AttributeError,
-            sa.exc.NoSuchTableError,
-            sa.exc.ProgrammingError,
-        ) as exc:
-            logger.debug(f"{type(exc).__name__} while introspecting columns", exc_info=exc)
-            logger.info(f"While introspecting columns {exc!r}; attempting reflection fallback")
-            # we will get a KeyError for temporary tables, since
-            # reflection will not find the temporary schema
-            columns = column_reflection_fallback(
-                selectable=table_selectable,
-                dialect=engine.dialect,
-                sqlalchemy_engine=engine,
-            )
+                primary_key_columns = set(pk_constraint.get("constrained_columns", []))
+            except (
+                sa.exc.NoSuchTableError,
+                sa.exc.ProgrammingError,
+                NotImplementedError,
+                AttributeError,
+            ) as e:
+                logger.debug(f"Could not fetch primary key info for {table_name}: {e!r}")
+
+        # Fetch column metadata
+        columns = _get_columns_from_selectable(
+            table_selectable,
+            inspector,
+            schema_name,
+            is_text_clause,
+            is_quoted_name,
+        )
 
-        # Use fallback because for mssql and trino reflection mechanisms do not throw an error but return an empty list  # noqa: E501 # FIXME CoP
-        if len(columns) == 0:
+        # Use fallback for mssql/trino or when primary introspection fails
+        if not columns:
             columns = column_reflection_fallback(
                 selectable=table_selectable,
                 dialect=engine.dialect,
                 sqlalchemy_engine=engine,
             )
 
-        dialect_name = execution_engine.dialect.name
-        if dialect_name in [
-            GXSqlDialect.DATABRICKS,
-            GXSqlDialect.POSTGRESQL,
-            GXSqlDialect.SNOWFLAKE,
-            GXSqlDialect.TRINO,
-        ]:
-            # WARNING: Do not alter columns in place, as they are cached on the inspector
-            columns_copy = [column.copy() for column in columns]
-            for column in columns_copy:
-                if column.get("type"):
-                    # When using column_reflection_fallback, we might not be able to
-                    # extract the column type, and only have the column name
-                    compiled_type = column["type"].compile(dialect=execution_engine.dialect)
-                    # Make the type case-insensitive
-                    column["type"] = CaseInsensitiveString(str(compiled_type))
-
-            # Wrap all columns in CaseInsensitiveNameDict for all three dialects
-            return [CaseInsensitiveNameDict(column) for column in columns_copy]
-
-        return columns
+        # Build result: copy columns, add PK info, apply dialect-specific formatting
+        return _build_column_metadata_result(columns, primary_key_columns, execution_engine)
     except AttributeError as e:
         logger.debug(f"Error while introspecting columns: {e!r}", exc_info=e)
         return None
 
 
+def _get_columns_from_selectable(
+    table_selectable: sqlalchemy.Select,
+    inspector: Any,
+    schema_name: Optional[str],
+    is_text_clause: bool,
+    is_quoted_name: bool,
+) -> Sequence[Dict[str, Any]]:
+    """Extract column metadata from a selectable, using reflection fallback on failure."""
+    try:
+        if is_text_clause:
+            # Custom SQL query - extract columns from the clause itself (SQLAlchemy 1.4+)
+            if hasattr(table_selectable, "selected_columns"):
+                return [
+                    {"name": col.name, "type": col.type}
+                    for col in table_selectable.selected_columns.values()
+                ]
+            # Pre-1.4 SQLAlchemy is no longer supported; fall back to reflection
+            logger.debug("TextClause without selected_columns; using reflection fallback")
+            return []
+
+        if not is_quoted_name:
+            logger.warning("unexpected table_selectable type")
+
+        return inspector.get_columns(
+            table_name=table_selectable if is_quoted_name else str(table_selectable),
+            schema=schema_name,
+        )
+    except (KeyError, AttributeError, sa.exc.NoSuchTableError, sa.exc.ProgrammingError) as exc:
+        logger.debug(f"{type(exc).__name__} while introspecting columns", exc_info=exc)
+        logger.info(f"While introspecting columns {exc!r}; attempting reflection fallback")
+        return []  # Caller will use column_reflection_fallback
+
+
+def _build_column_metadata_result(
+    columns: Sequence[Dict[str, Any]],
+    primary_key_columns: set[str],
+    execution_engine: SqlAlchemyExecutionEngine,
+) -> Sequence[Mapping[str, Any]]:
+    """Build final column metadata with PK info and dialect-specific formatting."""
+    # Copy columns to avoid mutating cached inspector data
+    pk_columns_lower = {pk.casefold() for pk in primary_key_columns}
+    result = [
+        {
+            **col,
+            "primary_key": col.get("name", "").casefold() in pk_columns_lower,
+        }
+        for col in (c.copy() for c in columns)
+    ]
+
+    # Apply case-insensitive formatting for specific dialects
+    dialect_name = execution_engine.dialect.name
+    case_insensitive_dialects = {
+        GXSqlDialect.DATABRICKS,
+        GXSqlDialect.POSTGRESQL,
+        GXSqlDialect.SNOWFLAKE,
+        GXSqlDialect.TRINO,
+    }
+    if dialect_name in case_insensitive_dialects:
+        for col in result:
+            if col.get("type"):
+                compiled_type = col["type"].compile(dialect=execution_engine.dialect)
+                col["type"] = CaseInsensitiveString(str(compiled_type))
+        return [CaseInsensitiveNameDict(col) for col in result]
+
+    return result
+
+
 def column_reflection_fallback(  # noqa: C901, PLR0912, PLR0915 # FIXME CoP
     selectable: sqlalchemy.Select,
     dialect: sqlalchemy.Dialect,

great_expectations/experimental/metric_repository/metric_retriever.py

@@ -316,10 +316,16 @@ def _get_table_column_types(self, batch_request: BatchRequest) -> Metric:
                     {
                         "name": raw_column_type["name"],
                         "type": str(raw_column_type["type"]),
+                        "primary_key": raw_column_type.get("primary_key", False),
                     }
                 )
             else:
-                column_types_converted_to_str.append({"name": raw_column_type["name"]})
+                column_types_converted_to_str.append(
+                    {
+                        "name": raw_column_type["name"],
+                        "primary_key": raw_column_type.get("primary_key", False),
+                    }
+                )
 
         return TableMetric[List[str]](
             batch_id=batch_id,

great_expectations/metrics/metric_results.py

@@ -79,6 +79,7 @@ class Config:
 
     name: str
     type: str
+    primary_key: bool
 
 
 class TableColumnTypesResult(MetricResult[list[ColumnType]]): ...

tests/expectations/metrics/conftest.py

@@ -1,11 +1,12 @@
-from typing import Iterable, Optional, Union
+from typing import Generator, Iterable, Optional, Union
 
 import pytest
 
 from great_expectations.compatibility.sqlalchemy import (
     sqlalchemy as sa,
 )
 from great_expectations.core.metric_domain_types import MetricDomainTypes
+from great_expectations.data_context.util import file_relative_path
 from great_expectations.execution_engine import SqlAlchemyExecutionEngine
 from great_expectations.execution_engine.sqlalchemy_batch_data import SqlAlchemyBatchData
 
@@ -72,3 +73,25 @@ def mock_sqlalchemy_execution_engine():
     execution_engine = MockSqlAlchemyExecutionEngine()
     execution_engine._batch_manager = MockBatchManager()
     return execution_engine
+
+
+@pytest.fixture
+def sql_data_connector_test_db_execution_engine() -> Generator[
+    SqlAlchemyExecutionEngine, None, None
+]:
+    """Provide a sqlite ExecutionEngine pointing to the SQL data connector test database.
+
+    The engine and its underlying connections are explicitly closed after use to avoid
+    leaking sqlite3.Connection objects (which surface as ResourceWarning in CI).
+    """
+    db_file = file_relative_path(
+        __file__,
+        "../../test_sets/test_cases_for_sql_data_connector.db",
+    )
+    engine: sa.engine.Engine = sa.create_engine(f"sqlite:///{db_file}")
+    execution_engine = SqlAlchemyExecutionEngine(engine=engine)
+
+    try:
+        yield execution_engine
+    finally:
+        execution_engine.close()

tests/expectations/metrics/test_metrics_util.py

@@ -632,6 +632,140 @@ def test__eq__(
             assert input_case_insensitive != other
 
 
+@pytest.mark.unit
+def test_get_sqlalchemy_column_metadata_includes_primary_key_field(
+    sql_data_connector_test_db_execution_engine,
+):
+    """Test that get_sqlalchemy_column_metadata includes primary_key field for all columns."""
+    from great_expectations.execution_engine.sqlalchemy_batch_data import SqlAlchemyBatchData
+    from great_expectations.expectations.metrics.util import get_sqlalchemy_column_metadata
+
+    engine = sql_data_connector_test_db_execution_engine
+
+    # Test table with single primary key
+    batch_data = SqlAlchemyBatchData(execution_engine=engine, table_name="table_with_single_pk")
+    engine.load_batch_data("__test_single_pk", batch_data)
+
+    columns = get_sqlalchemy_column_metadata(
+        execution_engine=engine,
+        table_selectable=sqlalchemy.quoted_name("table_with_single_pk", quote=False),
+        schema_name=None,
+    )
+
+    assert columns is not None
+    assert len(columns) == 3  # id, name, value
+
+    # All columns should have primary_key field
+    assert all("primary_key" in col for col in columns)
+
+    # Only 'id' should be marked as primary key
+    pk_columns = [col["name"] for col in columns if col["primary_key"]]
+    assert pk_columns == ["id"]
+
+    # Other columns should not be primary keys
+    non_pk_columns = [col["name"] for col in columns if not col["primary_key"]]
+    assert set(non_pk_columns) == {"name", "value"}
+
+
+@pytest.mark.unit
+def test_get_sqlalchemy_column_metadata_composite_primary_key(
+    sql_data_connector_test_db_execution_engine,
+):
+    """Test that composite primary keys are correctly identified."""
+    from great_expectations.execution_engine.sqlalchemy_batch_data import SqlAlchemyBatchData
+    from great_expectations.expectations.metrics.util import get_sqlalchemy_column_metadata
+
+    engine = sql_data_connector_test_db_execution_engine
+
+    batch_data = SqlAlchemyBatchData(execution_engine=engine, table_name="table_with_composite_pk")
+    engine.load_batch_data("__test_composite_pk", batch_data)
+
+    columns = get_sqlalchemy_column_metadata(
+        execution_engine=engine,
+        table_selectable=sqlalchemy.quoted_name("table_with_composite_pk", quote=False),
+        schema_name=None,
+    )
+
+    assert columns is not None
+    assert len(columns) == 4  # user_id, order_id, product, quantity
+
+    # All columns should have primary_key field
+    assert all("primary_key" in col for col in columns)
+
+    # Both user_id and order_id should be marked as primary keys
+    pk_columns = sorted([col["name"] for col in columns if col["primary_key"]])
+    assert pk_columns == ["order_id", "user_id"]
+
+    # Other columns should not be primary keys
+    non_pk_columns = sorted([col["name"] for col in columns if not col["primary_key"]])
+    assert non_pk_columns == ["product", "quantity"]
+
+
+@pytest.mark.unit
+def test_get_sqlalchemy_column_metadata_no_primary_key(
+    sql_data_connector_test_db_execution_engine,
+):
+    """Test that tables without primary keys don't break."""
+    from great_expectations.execution_engine.sqlalchemy_batch_data import SqlAlchemyBatchData
+    from great_expectations.expectations.metrics.util import get_sqlalchemy_column_metadata
+
+    engine = sql_data_connector_test_db_execution_engine
+
+    batch_data = SqlAlchemyBatchData(execution_engine=engine, table_name="table_without_pk")
+    engine.load_batch_data("__test_no_pk", batch_data)
+
+    columns = get_sqlalchemy_column_metadata(
+        execution_engine=engine,
+        table_selectable=sqlalchemy.quoted_name("table_without_pk", quote=False),
+        schema_name=None,
+    )
+
+    assert columns is not None
+    assert len(columns) == 2  # description, amount
+
+    # All columns should have primary_key field
+    assert all("primary_key" in col for col in columns)
+
+    # No columns should be marked as primary keys
+    pk_columns = [col["name"] for col in columns if col["primary_key"]]
+    assert pk_columns == []
+
+    # All columns should have primary_key=False
+    assert all(not col["primary_key"] for col in columns)
+
+
+@pytest.mark.unit
+def test_get_sqlalchemy_column_metadata_quoted_pk_column(
+    sql_data_connector_test_db_execution_engine,
+):
+    """Test that quoted column names as primary keys work correctly."""
+    from great_expectations.execution_engine.sqlalchemy_batch_data import SqlAlchemyBatchData
+    from great_expectations.expectations.metrics.util import get_sqlalchemy_column_metadata
+
+    engine = sql_data_connector_test_db_execution_engine
+
+    batch_data = SqlAlchemyBatchData(execution_engine=engine, table_name="table_with_quoted_pk")
+    engine.load_batch_data("__test_quoted_pk", batch_data)
+
+    columns = get_sqlalchemy_column_metadata(
+        execution_engine=engine,
+        table_selectable=sqlalchemy.quoted_name("table_with_quoted_pk", quote=False),
+        schema_name=None,
+    )
+
+    assert columns is not None
+    assert len(columns) == 2  # UserId, UserName
+
+    # All columns should have primary_key field
+    assert all("primary_key" in col for col in columns)
+
+    # UserId should be marked as primary key
+    pk_columns = [col["name"] for col in columns if col["primary_key"]]
+    assert len(pk_columns) == 1
+    # Case-insensitive check
+    assert pk_columns[0].lower() == "userid"
+
+
 @pytest.mark.unit
 @patch("great_expectations.expectations.metrics.util.sa")
 def test_get_dialect_like_pattern_expression_is_resilient_to_missing_dialects(mock_sqlalchemy):