diff --git a/crates/sqlbuild-python/src/bindings/_helpers/native_host/warehouse_catalog.rs b/crates/sqlbuild-python/src/bindings/_helpers/native_host/warehouse_catalog.rs index 112caea0a2..1fd2b10c0b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/native_host/warehouse_catalog.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/native_host/warehouse_catalog.rs @@ -3,7 +3,7 @@ use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyValueError; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; -use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::types::{PyDict, PyDictMethods, PyFrozenSet, PyTuple}; use sqlbuild_orchestration::warehouse_metadata::constants::MAX_HISTORY_RECORDS; use sqlbuild_orchestration::warehouse_metadata::models::{ CatalogRead, WarehouseCatalogColumn, WarehouseCatalogRecord, WarehouseHistoryPage, @@ -49,23 +49,43 @@ pub(crate) fn read_catalog( "catalog listing is not authoritative for this namespace; scan incomplete", )); } - let relations = call_with( - &adapter.getattr("list_relations")?, - &[ - ("connection", connection.clone()), - ( - "database", - request.location.database.clone().into_bound_py_any(py)?, - ), - ( - "schemas", - (request.location.schema.clone(),).into_bound_py_any(py)?, - ), - ], + let schema_scoped = adapter.is_instance( + &py.import("sqlbuild.adapter.relations.types")? + .getattr("SchemaScopedMetadataReader")?, )?; - let mut selected: Vec<(String, Bound<'_, PyAny>)> = Vec::new(); - for relation in relations.try_iter()? { - let relation = relation?; + let mut selected: Vec<(String, String, Bound<'_, PyAny>)> = Vec::new(); + let listed = if schema_scoped { + field( + &call_with( + &adapter.getattr("read_schema_relation_listing")?, + &[ + ("connection", connection.clone()), + namespace[0].clone(), + namespace[1].clone(), + ], + )?, + "entries", + )? + } else { + call_with( + &adapter.getattr("list_relations")?, + &[ + ("connection", connection.clone()), + namespace[0].clone(), + ( + "schemas", + (request.location.schema.clone(),).into_bound_py_any(py)?, + ), + ], + )? + }; + for item in listed.try_iter()? { + let item = item?; + let (relation, key) = if schema_scoped { + (field(&item, "relation")?, field(&item, "stored_name")?) + } else { + (item.clone(), field(&item, "name")?) + }; let name: String = field(&relation, "name")?.extract()?; if request .after @@ -73,7 +93,8 @@ pub(crate) fn read_catalog( .and_then(|cursor| cursor[0].as_deref()) .is_none_or(|after| name.as_str() > after) { - selected.push((name, relation)); + let relation_type: String = field(&relation, "relation_type")?.extract()?; + selected.push((name, relation_type, key)); } } selected.sort_by(|left, right| left.0.cmp(&right.0)); @@ -85,28 +106,43 @@ pub(crate) fn read_catalog( next_after: None, }); } - let names = PyTuple::new(py, selected.iter().map(|(name, _)| name))?; - let columns = call_with( - &adapter.getattr("get_all_columns")?, - &[ - ("connection", connection.clone()), - ("database", namespace[0].1.clone()), - ( - "schemas", - (request.location.schema.clone(),).into_bound_py_any(py)?, - ), - ("names", names.into_any()), - ], - )?; + let keys = selected.iter().map(|(_, _, key)| key); + let columns = if schema_scoped { + field( + &call_with( + &adapter.getattr("read_schema_column_listing")?, + &[ + ("connection", connection.clone()), + namespace[0].clone(), + namespace[1].clone(), + ("stored_names", PyFrozenSet::new(py, keys)?.into_any()), + ], + )?, + "columns_by_stored_name", + )? + } else { + call_with( + &adapter.getattr("get_all_columns")?, + &[ + ("connection", connection.clone()), + namespace[0].clone(), + ( + "schemas", + (request.location.schema.clone(),).into_bound_py_any(py)?, + ), + ("names", PyTuple::new(py, keys)?.into_any()), + ], + )? + }; let columns = columns.cast::()?; let next_after = if more { - selected.last().map(|(name, _)| vec![Some(name.clone())]) + selected.last().map(|(name, _, _)| vec![Some(name.clone())]) } else { None }; let mut records: Vec = Vec::new(); - for (name, relation) in selected { - let relation_columns = columns.get_item(name.as_str())?.ok_or_else(|| { + for (name, relation_type, key) in selected { + let relation_columns = columns.get_item(key)?.ok_or_else(|| { PyValueError::new_err("catalog relation columns are unavailable; scan incomplete") })?; let mut decoded: Vec = Vec::new(); @@ -121,7 +157,7 @@ pub(crate) fn read_catalog( state_database: request.location.database.clone(), state_schema: request.location.schema.clone(), name, - relation_type: field(&relation, "relation_type")?.extract()?, + relation_type, columns: decoded, })); } diff --git a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py index afe6a25c9d..729942f5d8 100644 --- a/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py +++ b/src/sqlbuild/adapters/snowflake/classes/snowflake_adapter.py @@ -4201,6 +4201,25 @@ def schema_exists( database: str | None, schema: str, ) -> bool: + """Check with SHOW SCHEMAS, which needs no running warehouse, where the name allows.""" + + show_database: str | None = ( + database if database is not None else self._current_database(connection=connection) + ) + if show_database is not None and SHOW_LIKE_ESCAPE_CHARACTER not in schema: + try: + rows: list[dict[str, object]] = self._fetch_show_rows( + connection=connection, + query=( + f"SHOW SCHEMAS LIKE '{show_like_pattern(schema.upper())}' " + f"IN DATABASE {self.render_identifier(show_database)}" + ), + ) + except Exception as error: + if not is_missing_object_error(error): + raise + return False + return any(str(row.get("name")).upper() == schema.upper() for row in rows) query: str = ( "SELECT 1 FROM information_schema.schemata WHERE UPPER(schema_name) = UPPER(%s)" ) diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py b/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py index 30ed59be03..e06ff737f7 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/_test_types.py @@ -417,3 +417,12 @@ class SnowflakePlatformDetectionTestCase: description: str environment: dict[str, str] expected_value: str + + +@dataclass(frozen=True) +class SnowflakeSchemaExistsTestCase: + description: str + schema: str + cursor: object + expected_exists: bool + expected_sql_fragment: str diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/helpers.py b/tests/unit/src/sqlbuild/adapters/snowflake/helpers.py index 4882cc8bc5..0c6d24b4f4 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/helpers.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/helpers.py @@ -78,6 +78,14 @@ def fetchmany(self, size: int) -> list[tuple[object, ...]]: return self._rows[:size] +class FakeSnowflakeMissingObjectCursor(FakeSnowflakeQueryCursor): + """Cursor double whose statement fails as Snowflake reports a missing object.""" + + def execute(self, sql: str) -> None: + super().execute(sql) + raise RuntimeError("Database 'GONE' does not exist or not authorized.") + + class FakeSnowflakeDescribeConnection: """Connection double returning a preconfigured describe cursor.""" diff --git a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py index 6a9dcf2cdd..5c64c55bcf 100644 --- a/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py +++ b/tests/unit/src/sqlbuild/adapters/snowflake/test_client.py @@ -47,6 +47,7 @@ SnowflakeRenderPythonFunctionTestCase, SnowflakeRenderTableFunctionTestCase, SnowflakeSchemaDiffTestCase, + SnowflakeSchemaExistsTestCase, SnowflakeTableFreshnessBatchTestCase, SnowflakeTableFreshnessMetadataErrorTestCase, SnowflakeTableFreshnessMetadataTestCase, @@ -59,6 +60,7 @@ FakeSnowflakeMetadataConnection, FakeSnowflakeMetadataCursor, FakeSnowflakeMetadataSequenceConnection, + FakeSnowflakeMissingObjectCursor, FakeSnowflakeQueryCursor, FakeSnowflakeRawConnection, build_bulk_columns_rows, @@ -1357,3 +1359,52 @@ def test_given_snowflake_relation_timestamps_when_listing_then_exposes_utc_age_m str(test_case.expected_created_at), str(test_case.expected_last_altered_at), ) + + +@pytest.mark.parametrize( + "test_case", + [ + SnowflakeSchemaExistsTestCase( + description="SHOW SCHEMAS finds a schema stored in another case", + schema="staging", + cursor=FakeSnowflakeQueryCursor((("name",),), [("Staging",)]), + expected_exists=True, + expected_sql_fragment="SHOW SCHEMAS LIKE 'STAGING' IN DATABASE", + ), + SnowflakeSchemaExistsTestCase( + description="a neighbour matched by the LIKE wildcard is not the schema", + schema="stg_a", + cursor=FakeSnowflakeQueryCursor((("name",),), [("STGXA",)]), + expected_exists=False, + expected_sql_fragment="SHOW SCHEMAS LIKE 'STG_A' IN DATABASE", + ), + SnowflakeSchemaExistsTestCase( + description="a missing database holds no schema", + schema="staging", + cursor=FakeSnowflakeMissingObjectCursor((("name",),), []), + expected_exists=False, + expected_sql_fragment="SHOW SCHEMAS LIKE 'STAGING' IN DATABASE", + ), + SnowflakeSchemaExistsTestCase( + description="a name with a backslash reads the information schema", + schema="stag\\ing", + cursor=FakeSnowflakeMetadataCursor(row=(1,)), + expected_exists=True, + expected_sql_fragment="information_schema.schemata", + ), + ], + ids=lambda case: case.description, +) +def test_given_schema_when_checking_existence_then_show_needs_no_running_warehouse( + test_case: SnowflakeSchemaExistsTestCase, +) -> None: + connection: FakeSnowflakeDescribeConnection = FakeSnowflakeDescribeConnection( + cast(Any, test_case.cursor) + ) + + exists: bool = SnowflakeAdapter().schema_exists( + connection=cast(Any, connection), database="ANALYTICS", schema=test_case.schema + ) + + assert exists is test_case.expected_exists + assert test_case.expected_sql_fragment in str(cast(Any, test_case.cursor).executed_sql)