Skip to content

Commit 07387ba

Browse files
kvlongekvlonge
authored
Merge branch 'main' into fix/snowflake-metadata-show
2 parents 8c676c3 + 1428b5e commit 07387ba

5 files changed

Lines changed: 112 additions & 27 deletions

File tree

‎src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py‎

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@
166166
SECONDARY_ROLES_ALL,
167167
SECONDARY_ROLES_NONE,
168168
SHOW_LIKE_ESCAPE_CHARACTER,
169+
SHOW_OBJECTS_TABLE_KIND,
169170
STATUS_COLUMN_NAME,
170171
SUCCESS_STATUS_TOKENS,
171172
TEMPORARY_FUNCTION_PREFIX,
@@ -3561,24 +3562,59 @@ def estimate_relation_row_count(
35613562
name: str,
35623563
metadata_only: bool = False,
35633564
) -> RelationRowCountEstimate:
3564-
"""Read ``ROW_COUNT`` from the information schema; views report no count."""
3565+
"""Read rows and bytes from SHOW OBJECTS, which needs no running warehouse."""
35653566

3566-
row: tuple[Any, ...] | None = self._information_schema_table_row(
3567-
connection=connection,
3568-
database=database,
3569-
schema=schema,
3570-
name=name,
3571-
projection="row_count,bytes" if metadata_only else "row_count",
3572-
)
3567+
basis: str = "show_objects_bytes"
3568+
row: tuple[Any, ...] | None
3569+
if schema is None or SHOW_LIKE_ESCAPE_CHARACTER in name:
3570+
basis = "information_schema_bytes"
3571+
row = self._information_schema_table_row(
3572+
connection=connection,
3573+
database=database,
3574+
schema=schema,
3575+
name=name,
3576+
projection="row_count,bytes",
3577+
)
3578+
else:
3579+
row = self._show_object_row_count(
3580+
connection=connection, database=database, schema=schema, name=name
3581+
)
35733582
if row is None:
35743583
return RelationRowCountEstimate(row_count=None, detail="relation not found")
35753584
byte_count: int | None = int(row[1]) if metadata_only and row[1] is not None else None
35763585
return RelationRowCountEstimate(
35773586
row_count=int(row[0]) if row[0] is not None else None,
35783587
detail="no row count (view)" if row[0] is None else None,
35793588
byte_count=byte_count,
3580-
byte_count_basis="information_schema_bytes" if byte_count is not None else None,
3589+
byte_count_basis=basis if byte_count is not None else None,
3590+
)
3591+
3592+
def _show_object_row_count(
3593+
self, *, connection: Any, database: str | None, schema: str, name: str
3594+
) -> tuple[Any, Any] | None:
3595+
"""Rows and bytes of one table; a view or other object lists neither."""
3596+
3597+
stored_name: str = self._information_schema_identifier(name)
3598+
rows: list[dict[str, object]] = self._fetch_schema_show_rows(
3599+
connection=connection,
3600+
query=(
3601+
f"SHOW OBJECTS LIKE '{show_like_pattern(stored_name)}' IN SCHEMA "
3602+
+ self._show_schema_scope(connection=connection, database=database, schema=schema)
3603+
),
3604+
show_database=self._show_database(
3605+
connection=connection, database=database, schema=schema
3606+
),
3607+
schema=schema,
35813608
)
3609+
row: dict[str, object]
3610+
for row in rows:
3611+
if str(row.get("name")) == stored_name:
3612+
return (
3613+
(row.get("rows"), row.get("bytes"))
3614+
if str(row.get("kind")) == SHOW_OBJECTS_TABLE_KIND
3615+
else (None, None)
3616+
)
3617+
return None
35823618

35833619
def inspect_row_diff_coverage(
35843620
self,

‎src/sqlbuild/adapters/snowflake/constants.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
TEXT_TYPE_NAMES: frozenset[str] = frozenset({"TEXT", "VARCHAR", "CHAR", "CHARACTER"})
1818
TRUE_METADATA_VALUE: str = "YES"
1919
SHOW_LIKE_ESCAPE_CHARACTER: str = "\\"
20+
SHOW_OBJECTS_TABLE_KIND: str = "TABLE"
2021
CURRENT_DATABASE_ATTRIBUTE: str = "_sqlbuild_current_database"
2122
STATE_TABLE_RETENTION_DAYS: int = 90
2223
STATE_TABLE_FALLBACK_RETENTION_DAYS: int = 1

‎tests/unit/src/sqlbuild/adapters/_test_types.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ class RelationRowCountEstimateTestCase:
148148
metadata_only: bool = False
149149
expected_byte_count: int | None = None
150150
expected_byte_count_basis: str | None = None
151+
metadata_columns: tuple[str, ...] = ()
152+
name: str = "orders"
151153

152154

153155
@dataclass(frozen=True)

‎tests/unit/src/sqlbuild/adapters/helpers.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,16 @@ def execute(*, connection: Any, sql: str) -> GrantRows:
8383
class MetadataCursor:
8484
"""A cursor that records metadata SQL and returns fixed rows."""
8585

86-
def __init__(self, *, rows: tuple[tuple[object, ...], ...], executed: list[str]) -> None:
86+
def __init__(
87+
self,
88+
*,
89+
rows: tuple[tuple[object, ...], ...],
90+
executed: list[str],
91+
columns: tuple[str, ...] = (),
92+
) -> None:
8793
self._rows: tuple[tuple[object, ...], ...] = rows
8894
self._executed: list[str] = executed
95+
self.description: tuple[tuple[str], ...] = tuple((column,) for column in columns)
8996

9097
def execute(self, sql: str, params: tuple[object, ...] | None = None) -> MetadataCursor:
9198
self._executed.append(f"{sql} {params or ()}")
@@ -104,12 +111,13 @@ def close(self) -> None:
104111
class MetadataConnection:
105112
"""A connection whose cursors share one SQL log and fixed rows."""
106113

107-
def __init__(self, rows: tuple[tuple[object, ...], ...]) -> None:
114+
def __init__(self, rows: tuple[tuple[object, ...], ...], columns: tuple[str, ...] = ()) -> None:
108115
self.rows: tuple[tuple[object, ...], ...] = rows
116+
self.columns: tuple[str, ...] = columns
109117
self.executed: list[str] = []
110118

111119
def cursor(self) -> MetadataCursor:
112-
return MetadataCursor(rows=self.rows, executed=self.executed)
120+
return MetadataCursor(rows=self.rows, executed=self.executed, columns=self.columns)
113121

114122

115123
@dataclass(frozen=True)

‎tests/unit/src/sqlbuild/adapters/test_relation_row_count_estimate.py‎

Lines changed: 53 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
execute_through_cursor,
2929
)
3030

31+
_SNOWFLAKE_SHOW_OBJECTS_COLUMNS: tuple[str, ...] = ("name", "kind", "rows", "bytes")
3132
_DATABRICKS_DETAIL_HEADER: tuple[tuple[object, ...], ...] = (
3233
("order_id", "bigint", None),
3334
("", "", ""),
@@ -79,28 +80,49 @@
7980
expected_sql_fragment="pg_class",
8081
),
8182
RelationRowCountEstimateTestCase(
82-
description="snowflake table reports information schema row count",
83+
description="snowflake table reports its row count from SHOW OBJECTS",
8384
adapter=SnowflakeAdapter(),
84-
metadata_rows=((5000,),),
85+
metadata_rows=(("ORDERS", "TABLE", 5000, 65536),),
86+
metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS,
8587
expected_row_count=5000,
8688
expected_detail=None,
87-
expected_sql_fragment="SELECT row_count FROM",
89+
expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA",
8890
),
8991
RelationRowCountEstimateTestCase(
90-
description="snowflake view has no row count",
92+
description="snowflake view has no row count although SHOW lists zero rows",
9193
adapter=SnowflakeAdapter(),
92-
metadata_rows=((None,),),
94+
metadata_rows=(("ORDERS", "VIEW", 0, 0),),
95+
metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS,
9396
expected_row_count=None,
9497
expected_detail="no row count (view)",
95-
expected_sql_fragment="('ORDERS', 'DEV', 'ANALYTICS')",
98+
expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA",
9699
),
97100
RelationRowCountEstimateTestCase(
98-
description="snowflake missing relation is unknown",
101+
description="snowflake neighbour matched by the LIKE wildcard is not the relation",
99102
adapter=SnowflakeAdapter(),
100-
metadata_rows=(),
103+
metadata_rows=(("ORDERS_ARCHIVE", "TABLE", 9, 512),),
104+
metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS,
101105
expected_row_count=None,
102106
expected_detail="relation not found",
103-
expected_sql_fragment="information_schema.tables",
107+
expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA",
108+
),
109+
RelationRowCountEstimateTestCase(
110+
description="snowflake name with a backslash reads the information schema",
111+
adapter=SnowflakeAdapter(),
112+
metadata_rows=((5000, 65536),),
113+
name="order\\s",
114+
expected_row_count=5000,
115+
expected_detail=None,
116+
expected_sql_fragment="SELECT row_count,bytes FROM",
117+
),
118+
RelationRowCountEstimateTestCase(
119+
description="snowflake relation without a schema reads the information schema",
120+
adapter=SnowflakeAdapter(),
121+
metadata_rows=((5000, 65536),),
122+
schema=None,
123+
expected_row_count=5000,
124+
expected_detail=None,
125+
expected_sql_fragment="SELECT row_count,bytes FROM",
104126
),
105127
RelationRowCountEstimateTestCase(
106128
description="databricks analyzed table reports statistics rows",
@@ -185,22 +207,36 @@
185207
RelationRowCountEstimateTestCase(
186208
description="snowflake bytes and rows are distinct",
187209
adapter=SnowflakeAdapter(),
188-
metadata_rows=((17, 4096),),
210+
metadata_rows=(("ORDERS", "TABLE", 17, 4096),),
211+
metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS,
189212
expected_row_count=17,
190213
expected_detail=None,
191-
expected_sql_fragment="SELECT row_count,bytes FROM",
214+
expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA",
192215
metadata_only=True,
193216
expected_byte_count=4096,
194-
expected_byte_count_basis="information_schema_bytes",
217+
expected_byte_count_basis="show_objects_bytes",
195218
),
196219
RelationRowCountEstimateTestCase(
197220
description="snowflake missing view statistics remain unavailable",
198221
adapter=SnowflakeAdapter(),
199-
metadata_rows=((None, None),),
222+
metadata_rows=(("ORDERS", "VIEW", 0, 0),),
223+
metadata_columns=_SNOWFLAKE_SHOW_OBJECTS_COLUMNS,
200224
expected_row_count=None,
201225
expected_detail="no row count (view)",
226+
expected_sql_fragment="SHOW OBJECTS LIKE 'ORDERS' IN SCHEMA",
227+
metadata_only=True,
228+
),
229+
RelationRowCountEstimateTestCase(
230+
description="snowflake bytes without a schema keep the information schema basis",
231+
adapter=SnowflakeAdapter(),
232+
metadata_rows=((17, 4096),),
233+
schema=None,
234+
expected_row_count=17,
235+
expected_detail=None,
202236
expected_sql_fragment="SELECT row_count,bytes FROM",
203237
metadata_only=True,
238+
expected_byte_count=4096,
239+
expected_byte_count_basis="information_schema_bytes",
204240
),
205241
RelationRowCountEstimateTestCase(
206242
description="databricks bytes do not imply an analyzed row count",
@@ -240,15 +276,17 @@
240276
def test_given_catalog_metadata_when_estimating_row_count_then_reads_only_metadata(
241277
test_case: RelationRowCountEstimateTestCase, monkeypatch: pytest.MonkeyPatch
242278
) -> None:
243-
connection: MetadataConnection = MetadataConnection(test_case.metadata_rows)
279+
connection: MetadataConnection = MetadataConnection(
280+
test_case.metadata_rows, test_case.metadata_columns
281+
)
244282

245283
monkeypatch.setattr(test_case.adapter, "execute", execute_through_cursor)
246284

247285
estimate: RelationRowCountEstimate = test_case.adapter.estimate_relation_row_count(
248286
connection=connection,
249287
database=test_case.database,
250288
schema=test_case.schema,
251-
name="orders",
289+
name=test_case.name,
252290
metadata_only=test_case.metadata_only,
253291
)
254292

0 commit comments

Comments
 (0)