diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index ddf90d338..afe6a25c9 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -166,6 +166,7 @@ SECONDARY_ROLES_ALL, SECONDARY_ROLES_NONE, SHOW_LIKE_ESCAPE_CHARACTER, + SHOW_OBJECTS_TABLE_KIND, STATUS_COLUMN_NAME, SUCCESS_STATUS_TOKENS, TEMPORARY_FUNCTION_PREFIX, @@ -3561,15 +3562,23 @@ def estimate_relation_row_count( name: str, metadata_only: bool = False, ) -> RelationRowCountEstimate: - """Read ``ROW_COUNT`` from the information schema; views report no count.""" + """Read rows and bytes from SHOW OBJECTS, which needs no running warehouse.""" - row: tuple[Any, ...] | None = self._information_schema_table_row( - connection=connection, - database=database, - schema=schema, - name=name, - projection="row_count,bytes" if metadata_only else "row_count", - ) + basis: str = "show_objects_bytes" + row: tuple[Any, ...] | None + if schema is None or SHOW_LIKE_ESCAPE_CHARACTER in name: + basis = "information_schema_bytes" + row = self._information_schema_table_row( + connection=connection, + database=database, + schema=schema, + name=name, + projection="row_count,bytes", + ) + else: + row = self._show_object_row_count( + connection=connection, database=database, schema=schema, name=name + ) if row is None: return RelationRowCountEstimate(row_count=None, detail="relation not found") byte_count: int | None = int(row[1]) if metadata_only and row[1] is not None else None @@ -3577,8 +3586,35 @@ def estimate_relation_row_count( row_count=int(row[0]) if row[0] is not None else None, detail="no row count (view)" if row[0] is None else None, byte_count=byte_count, - byte_count_basis="information_schema_bytes" if byte_count is not None else None, + byte_count_basis=basis if byte_count is not None else None, + ) + + def _show_object_row_count( + self, *, connection: Any, database: str | None, schema: str, name: str + ) -> tuple[Any, Any] | None: + """Rows and bytes of one table; a view or other object lists neither.""" + + stored_name: str = self._information_schema_identifier(name) + rows: list[dict[str, object]] = self._fetch_schema_show_rows( + connection=connection, + query=( + f"SHOW OBJECTS LIKE '{show_like_pattern(stored_name)}' IN SCHEMA " + + self._show_schema_scope(connection=connection, database=database, schema=schema) + ), + show_database=self._show_database( + connection=connection, database=database, schema=schema + ), + schema=schema, ) + row: dict[str, object] + for row in rows: + if str(row.get("name")) == stored_name: + return ( + (row.get("rows"), row.get("bytes")) + if str(row.get("kind")) == SHOW_OBJECTS_TABLE_KIND + else (None, None) + ) + return None def inspect_row_diff_coverage( self, diff --git a/src/sqlbuild/adapters/snowflake/constants.py b/src/sqlbuild/adapters/snowflake/constants.py index 9d8e1fe45..78b6e350a 100644 --- a/src/sqlbuild/adapters/snowflake/constants.py +++ b/src/sqlbuild/adapters/snowflake/constants.py @@ -17,6 +17,7 @@ TEXT_TYPE_NAMES: frozenset[str] = frozenset({"TEXT", "VARCHAR", "CHAR", "CHARACTER"}) TRUE_METADATA_VALUE: str = "YES" SHOW_LIKE_ESCAPE_CHARACTER: str = "\\" +SHOW_OBJECTS_TABLE_KIND: str = "TABLE" CURRENT_DATABASE_ATTRIBUTE: str = "_sqlbuild_current_database" STATE_TABLE_RETENTION_DAYS: int = 90 STATE_TABLE_FALLBACK_RETENTION_DAYS: int = 1 diff --git a/tests/unit/src/sqlbuild/adapters/_test_types.py b/tests/unit/src/sqlbuild/adapters/_test_types.py index 7e4be3ec4..1ef3802c5 100644 --- a/tests/unit/src/sqlbuild/adapters/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/_test_types.py @@ -148,6 +148,8 @@ class RelationRowCountEstimateTestCase: metadata_only: bool = False expected_byte_count: int | None = None expected_byte_count_basis: str | None = None + metadata_columns: tuple[str, ...] = () + name: str = "orders" @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/adapters/helpers.py b/tests/unit/src/sqlbuild/adapters/helpers.py index bc838d5b6..de65f3f30 100644 --- a/tests/unit/src/sqlbuild/adapters/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/helpers.py @@ -83,9 +83,16 @@ def execute(*, connection: Any, sql: str) -> GrantRows: class MetadataCursor: """A cursor that records metadata SQL and returns fixed rows.""" - def __init__(self, *, rows: tuple[tuple[object, ...], ...], executed: list[str]) -> None: + def __init__( + self, + *, + rows: tuple[tuple[object, ...], ...], + executed: list[str], + columns: tuple[str, ...] = (), + ) -> None: self._rows: tuple[tuple[object, ...], ...] = rows self._executed: list[str] = executed + self.description: tuple[tuple[str], ...] = tuple((column,) for column in columns) def execute(self, sql: str, params: tuple[object, ...] | None = None) -> MetadataCursor: self._executed.append(f"{sql} {params or ()}") @@ -104,12 +111,13 @@ def close(self) -> None: class MetadataConnection: """A connection whose cursors share one SQL log and fixed rows.""" - def __init__(self, rows: tuple[tuple[object, ...], ...]) -> None: + def __init__(self, rows: tuple[tuple[object, ...], ...], columns: tuple[str, ...] = ()) -> None: self.rows: tuple[tuple[object, ...], ...] = rows + self.columns: tuple[str, ...] = columns self.executed: list[str] = [] def cursor(self) -> MetadataCursor: - return MetadataCursor(rows=self.rows, executed=self.executed) + return MetadataCursor(rows=self.rows, executed=self.executed, columns=self.columns) @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/adapters/test_relation_row_count_estimate.py b/tests/unit/src/sqlbuild/adapters/test_relation_row_count_estimate.py index dd28c03fb..f2c2bacc0 100644 --- a/tests/unit/src/sqlbuild/adapters/test_relation_row_count_estimate.py +++ b/tests/unit/src/sqlbuild/adapters/test_relation_row_count_estimate.py @@ -28,6 +28,7 @@ execute_through_cursor, ) +_SNOWFLAKE_SHOW_OBJECTS_COLUMNS: tuple[str, ...] = ("name", "kind", "rows", "bytes") _DATABRICKS_DETAIL_HEADER: tuple[tuple[object, ...], ...] = ( ("order_id", "bigint", None), ("", "", ""), @@ -79,28 +80,49 @@ expected_sql_fragment="pg_class", ), RelationRowCountEstimateTestCase( - description="snowflake table reports information schema row count", + description="snowflake table reports its row count from SHOW OBJECTS", adapter=SnowflakeAdapter(), - metadata_rows=((5000,),), + metadata_rows=(("ORDERS", "TABLE", 5000, 65536),), + metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS, expected_row_count=5000, expected_detail=None, - expected_sql_fragment="SELECT row_count FROM", + expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA", ), RelationRowCountEstimateTestCase( - description="snowflake view has no row count", + description="snowflake view has no row count although SHOW lists zero rows", adapter=SnowflakeAdapter(), - metadata_rows=((None,),), + metadata_rows=(("ORDERS", "VIEW", 0, 0),), + metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS, expected_row_count=None, expected_detail="no row count (view)", - expected_sql_fragment="('ORDERS', 'DEV', 'ANALYTICS')", + expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA", ), RelationRowCountEstimateTestCase( - description="snowflake missing relation is unknown", + description="snowflake neighbour matched by the LIKE wildcard is not the relation", adapter=SnowflakeAdapter(), - metadata_rows=(), + metadata_rows=(("ORDERS_ARCHIVE", "TABLE", 9, 512),), + metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS, expected_row_count=None, expected_detail="relation not found", - expected_sql_fragment="information_schema.tables", + expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA", + ), + RelationRowCountEstimateTestCase( + description="snowflake name with a backslash reads the information schema", + adapter=SnowflakeAdapter(), + metadata_rows=((5000, 65536),), + name="order\\s", + expected_row_count=5000, + expected_detail=None, + expected_sql_fragment="SELECT row_count,bytes FROM", + ), + RelationRowCountEstimateTestCase( + description="snowflake relation without a schema reads the information schema", + adapter=SnowflakeAdapter(), + metadata_rows=((5000, 65536),), + schema=None, + expected_row_count=5000, + expected_detail=None, + expected_sql_fragment="SELECT row_count,bytes FROM", ), RelationRowCountEstimateTestCase( description="databricks analyzed table reports statistics rows", @@ -185,22 +207,36 @@ RelationRowCountEstimateTestCase( description="snowflake bytes and rows are distinct", adapter=SnowflakeAdapter(), - metadata_rows=((17, 4096),), + metadata_rows=(("ORDERS", "TABLE", 17, 4096),), + metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS, expected_row_count=17, expected_detail=None, - expected_sql_fragment="SELECT row_count,bytes FROM", + expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA", metadata_only=True, expected_byte_count=4096, - expected_byte_count_basis="information_schema_bytes", + expected_byte_count_basis="show_objects_bytes", ), RelationRowCountEstimateTestCase( description="snowflake missing view statistics remain unavailable", adapter=SnowflakeAdapter(), - metadata_rows=((None, None),), + metadata_rows=(("ORDERS", "VIEW", 0, 0),), + metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS, expected_row_count=None, expected_detail="no row count (view)", + expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA", + metadata_only=True, + ), + RelationRowCountEstimateTestCase( + description="snowflake bytes without a schema keep the information schema basis", + adapter=SnowflakeAdapter(), + metadata_rows=((17, 4096),), + schema=None, + expected_row_count=17, + expected_detail=None, expected_sql_fragment="SELECT row_count,bytes FROM", metadata_only=True, + expected_byte_count=4096, + expected_byte_count_basis="information_schema_bytes", ), RelationRowCountEstimateTestCase( description="databricks bytes do not imply an analyzed row count", @@ -240,7 +276,9 @@ def test_given_catalog_metadata_when_estimating_row_count_then_reads_only_metadata( test_case: RelationRowCountEstimateTestCase, monkeypatch: pytest.MonkeyPatch ) -> None: - connection: MetadataConnection = MetadataConnection(test_case.metadata_rows) + connection: MetadataConnection = MetadataConnection( + test_case.metadata_rows, test_case.metadata_columns + ) monkeypatch.setattr(test_case.adapter, "execute", execute_through_cursor) @@ -248,7 +286,7 @@ def test_given_catalog_metadata_when_estimating_row_count_then_reads_only_metada connection=connection, database=test_case.database, schema=test_case.schema, - name="orders", + name=test_case.name, metadata_only=test_case.metadata_only, )