diff --git a/target_snowflake/connector.py b/target_snowflake/connector.py index 96b2dcc..1c81ed1 100644 --- a/target_snowflake/connector.py +++ b/target_snowflake/connector.py @@ -116,6 +116,21 @@ def inspector(self) -> sqlalchemy.Inspector: self._inspector = sqlalchemy.inspect(self._engine) return self._inspector + def invalidate_table_cache(self, full_table_name: str) -> None: + """Discard cached reflection state for a table after DDL. + + Dropping the `table_cache` entry alone is not enough. The Inspector + memoises reflection for its lifetime, and snowflake-sqlalchemy caches + columns per *schema*, so a table created after the first reflection of + that schema stays invisible and reflecting it raises NoSuchTableError. + The Inspector is dropped so the next reflection queries Snowflake. + + Args: + full_table_name: the table whose cached state is now stale. + """ + self.table_cache.pop(full_table_name, None) + self._inspector = None + def get_table_columns( self, full_table_name: str | FullyQualifiedName, diff --git a/target_snowflake/sinks.py b/target_snowflake/sinks.py index 6f2719a..d71377e 100644 --- a/target_snowflake/sinks.py +++ b/target_snowflake/sinks.py @@ -88,7 +88,7 @@ def setup(self) -> None: self.conform_name(self.schema_name, object_type="schema"), ) - self.connector.table_cache.pop(self.full_table_name, None) + self.connector.invalidate_table_cache(self.full_table_name) try: self.connector.prepare_table( @@ -105,7 +105,7 @@ def setup(self) -> None: ) raise - self.connector.table_cache.pop(self.full_table_name, None) + self.connector.invalidate_table_cache(self.full_table_name) if self.config.get("load_method", "upsert") == "overwrite": self.logger.info("load_method=overwrite: truncating %s", self.full_table_name) diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index 215e0e6..15305e2 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -193,3 +193,37 @@ def test_format_identifier(connector: SnowflakeConnector, config: dict, identifi connector.config.update(config) formatted = connector.format_identifier(identifier) assert formatted == expected_formatted + + +def test_invalidate_table_cache_drops_inspector(connector: SnowflakeConnector): + """Invalidating a table must drop the cached Inspector, not just table_cache. + + snowflake-sqlalchemy memoises columns per schema on the Inspector, so + reusing it after a CREATE TABLE reflects the schema as it was before the + table existed and raises NoSuchTableError. + """ + connector.table_cache["DB.SCHEMA.USERS"] = {"id": mock.Mock()} + connector.table_cache["DB.SCHEMA.POSTS"] = {"id": mock.Mock()} + connector._inspector = mock.Mock(spec=sa.Inspector) + + connector.invalidate_table_cache("DB.SCHEMA.USERS") + + assert "DB.SCHEMA.USERS" not in connector.table_cache + assert "DB.SCHEMA.POSTS" in connector.table_cache + assert connector._inspector is None + + +def test_inspector_rebuilt_after_invalidation(connector: SnowflakeConnector): + """A fresh Inspector is built on next access, so reflection re-queries.""" + connector._cached_engine = mock.Mock() + + with mock.patch("sqlalchemy.inspect") as inspect_mock: + inspect_mock.side_effect = [mock.sentinel.stale, mock.sentinel.fresh] + + assert connector.inspector is mock.sentinel.stale + assert connector.inspector is mock.sentinel.stale # cached + + connector.invalidate_table_cache("DB.SCHEMA.USERS") + + assert connector.inspector is mock.sentinel.fresh + assert inspect_mock.call_count == 2 diff --git a/tests/unit/test_sinks.py b/tests/unit/test_sinks.py index b7d4936..4d8b9c6 100644 --- a/tests/unit/test_sinks.py +++ b/tests/unit/test_sinks.py @@ -7,9 +7,11 @@ from unittest import mock import pytest +import sqlalchemy as sa from singer_sdk.helpers._batch import JSONLinesEncoding from target_snowflake.arrow_batch import ArrowEncoding +from target_snowflake.connector import SnowflakeConnector from target_snowflake.sinks import SnowflakeSink @@ -148,3 +150,42 @@ def fake_convert(manifest, output_dir, *, clean_up_source_files): # noqa: ARG00 SnowflakeSink._process_arrow_batch_files(fake_sink, ["file:///tmp/a.arrow"]) # noqa: SLF001 assert not os.path.exists(captured_output_dir["path"]) + + +def test_setup_invalidates_cached_inspector(): + """setup() must drop the connector's cached Inspector after creating the table. + + snowflake-sqlalchemy memoises columns per schema on the Inspector, so a table + created here stays invisible to later reflection unless the Inspector is + dropped. That surfaced as NoSuchTableError in activate_version for every + stream after the first. + """ + connector = SnowflakeConnector( + config={ + "user": "test_user", + "password": "test_password", + "account": "test_account", + "database": "TEST_DATABASE", + }, + ) + connector._inspector = mock.Mock(spec=sa.Inspector) + connector.table_cache["DB.SCHEMA.TABLE"] = {"id": mock.Mock()} + connector.prepare_schema = mock.MagicMock() + connector.prepare_table = mock.MagicMock() + + fake_sink = SimpleNamespace( + schema_name="SCHEMA", + full_table_name="DB.SCHEMA.TABLE", + schema={}, + key_properties=[], + config={}, + logger=mock.MagicMock(), + connector=connector, + conform_name=lambda name, object_type=None: name, # noqa: ARG005 + conform_schema=lambda schema: schema, + ) + + SnowflakeSink.setup(fake_sink) + + assert connector._inspector is None + assert "DB.SCHEMA.TABLE" not in connector.table_cache