fix: dedupe copy_into(qualify=True) via temp table and merge

README.md

@@ -68,7 +68,7 @@ Available (public) methods:
   - whether the table should be fully refreshed
   - a list of target columns, if not all the columns are present in the file to be loaded (or not all need to be written)
   - whether to sync tags (provided in the table structure) to the table/columns
-  - whether to perform a qualify on the table after loading the data. If this is used, a list of primary keys (and optionally of replication keys) should also be provided. If only the primary keys are supplied, those will also be used to determine which records are kept (non deterministic).
+  - whether to dedupe the loaded data (qualify). If this is used, a list of primary keys (and optionally of replication keys) should also be provided. If only the primary keys are supplied, those will also be used to determine which records are kept (non deterministic). Unless `full_refresh` is set, the data is copied to a `<table>_temp` table, deduped there and merged into the destination, so the destination never holds duplicates visible to readers. With `full_refresh` the table is copied into directly and then rebuilt with the qualify.
   - a stage parameter to use an existing stage instead of creating a temporary one
 - *create_table*: runs the create table statement, with optional full refresh to recreate an existing table.
 - *setup_file_format*: given a file format object, creates the corresponding resource in Snowflake

snowflake_utils/models/table.py

@@ -253,39 +253,52 @@ def copy_into(
                 {files_clause}
                 {self._include_metadata()}
                 """
-        if qualify:
-            self._copy(
-                copy_query,
-                path,
-                file_format,
-                storage_integration,
-                full_refresh,
-                sync_tags,
-                stage,
-                create_table,
-                copy_grants,
-            )
-            with connect() as connection:
-                cursor = connection.cursor()
-                self.qualify(
-                    cursor=cursor,
-                    primary_keys=primary_keys,
-                    replication_keys=replication_keys,
+        if qualify and not full_refresh:
+            # dedupe in a temp table so the live table never holds duplicates
+            def copy_callable(table: Table, sync_tags: bool) -> None:
+                return table.copy_into(
+                    path=path,
+                    file_format=file_format,
+                    storage_integration=storage_integration,
+                    match_by_column_name=match_by_column_name,
+                    target_columns=target_columns,
+                    sync_tags=sync_tags,
+                    stage=stage,
+                    files=files,
+                    create_table=create_table or table is not self,
+                    copy_grants=copy_grants,
                 )
-                if sync_tags and self.table_structure:
-                    self.sync_tags(cursor)
-        else:
-            return self._copy(
-                copy_query,
-                path,
-                file_format,
-                storage_integration,
-                full_refresh,
-                sync_tags,
-                stage,
-                create_table,
-                copy_grants,
+
+            return self._merge(
+                copy_callable,
+                primary_keys,
+                replication_keys,
+                qualify=True,
+                sync_tags=sync_tags,
+            )
+
+        result = self._copy(
+            copy_query,
+            path,
+            file_format,
+            storage_integration,
+            full_refresh,
+            sync_tags,
+            stage,
+            create_table,
+            copy_grants,
+        )
+        if not qualify:
+            return result
+        with connect() as connection:
+            cursor = connection.cursor()
+            self.qualify(
+                cursor=cursor,
+                primary_keys=primary_keys,
+                replication_keys=replication_keys,
             )
+            if sync_tags and self.table_structure:
+                self.sync_tags(cursor)
 
     def create_table(
         self, full_refresh: bool, execute_statement: callable, copy_grants: bool = True
@@ -331,16 +344,21 @@ def _merge(
         primary_keys: list[str] = ["id"],
         replication_keys: list[str] | None = None,
         qualify: bool = False,
+        sync_tags: bool = True,
     ) -> None:
         with connect() as connection:
             cursor = connection.cursor()
             if not self.exists(cursor):
-                copy_callable(self, sync_tags=True)
+                copy_callable(self, sync_tags=sync_tags)
                 if qualify:
                     self.qualify(cursor, primary_keys, replication_keys)
+                    if sync_tags and self.table_structure:
+                        self.sync_tags(cursor)
                 return None
 
         temp_table = self.model_copy(update={"name": f"{self.name}_temp"})
+        with connect() as connection:
+            connection.cursor().execute(f"drop table if exists {temp_table.fqn}")
         copy_callable(temp_table, sync_tags=False)
         if qualify:
             with connect() as connection:
@@ -349,9 +367,6 @@ def _merge(
 
         with connect() as connection:
             cursor = connection.cursor()
-            cursor.execute(
-                self.get_create_table_statement(full_refresh=False, copy_grants=True)
-            )
             old_columns = {x.name: x.data_type for x in self.get_columns(cursor)}
             new_columns = temp_table.get_columns(cursor)
 
@@ -364,7 +379,7 @@ def _merge(
                     temp_table, new_columns, old_columns, primary_keys
                 )
             )
-            if self.table_structure:
+            if sync_tags and self.table_structure:
                 self.sync_tags(cursor)
             temp_table.drop(cursor)
 

tests/test_models.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import os
 from datetime import datetime
@@ -407,6 +408,147 @@ def test_merge(mock_merge, mock_copy):
             }
 
 
+@patch.object(Table, "_copy")
+@patch.object(Table, "_merge")
+def test_copy_into_qualify_merges_instead_of_copying_into_live_table(
+    mock_merge, mock_copy
+):
+    test_table.copy_into(
+        path=path,
+        file_format=parquet_file_format,
+        storage_integration=storage_integration,
+        primary_keys=["id"],
+        qualify=True,
+        sync_tags=True,
+    )
+
+    mock_copy.assert_not_called()
+    mock_merge.assert_called_once()
+    _, kwargs = mock_merge.call_args
+    assert kwargs == {"qualify": True, "sync_tags": True}
+
+
+@patch.object(Table, "qualify")
+@patch.object(Table, "_copy")
+@patch.object(Table, "_merge")
+def test_copy_into_qualify_full_refresh_keeps_copy_then_qualify(
+    mock_merge, mock_copy, mock_qualify
+):
+    with patch("snowflake_utils.models.table.connect") as mock_connect:
+        mock_connect.return_value = make_mock_conn()
+        test_table.copy_into(
+            path=path,
+            file_format=parquet_file_format,
+            storage_integration=storage_integration,
+            primary_keys=["id"],
+            qualify=True,
+            full_refresh=True,
+        )
+
+    mock_merge.assert_not_called()
+    mock_copy.assert_called_once()
+    mock_qualify.assert_called_once()
+
+
+def copy_args(call) -> dict:
+    return inspect.signature(Table._copy).bind(*call.args, **call.kwargs).arguments
+
+
+@pytest.mark.parametrize("create_table", [True, False])
+@patch.object(Table, "_copy", autospec=True)
+@patch.object(Table, "_merge")
+def test_copy_into_qualify_always_creates_temp_table(
+    mock_merge, mock_copy, create_table
+):
+    test_table.copy_into(
+        path=path,
+        file_format=parquet_file_format,
+        storage_integration=storage_integration,
+        target_columns=["id"],
+        qualify=True,
+        create_table=create_table,
+    )
+    copy_callable = mock_merge.call_args.args[0]
+
+    copy_callable(
+        test_table.model_copy(update={"name": "PYTEST_temp"}), sync_tags=False
+    )
+    copy_callable(test_table, sync_tags=True)
+
+    temp_call, live_call = (copy_args(c) for c in mock_copy.call_args_list)
+    assert temp_call["sync_tags"] is False and temp_call["create_table"] is True
+    assert live_call["sync_tags"] is True and live_call["create_table"] is create_table
+    assert "COPY INTO PUBLIC.PYTEST_temp (id)" in temp_call["query"].replace("\n", " ")
+
+
+@patch.object(Table, "drop")
+@patch.object(Table, "get_columns")
+@patch.object(Table, "exists", return_value=True)
+@patch.object(Table, "_copy")
+def test_copy_into_qualify_existing_table_without_structure(
+    mock_copy, mock_exists, mock_get_columns, mock_drop
+):
+    mock_get_columns.return_value = [Column(name="id", data_type="integer")]
+    mock_cursor = make_mock_cursor()
+    inferred_table = Table(name="PYTEST_INFERRED", schema_name="PUBLIC")
+    with patch("snowflake_utils.models.table.connect") as mock_connect:
+        mock_connect.return_value = make_mock_conn(cursor=mock_cursor)
+        inferred_table.copy_into(
+            path=path,
+            file_format=parquet_file_format,
+            storage_integration=storage_integration,
+            primary_keys=["id"],
+            qualify=True,
+        )
+
+    statements = [
+        " ".join(c.args[0].split()) for c in mock_cursor.execute.call_args_list
+    ]
+    assert any(
+        s.lower().startswith("merge into public.pytest_inferred") for s in statements
+    )
+
+
+@patch.object(Table, "sync_tags")
+@patch.object(Table, "drop")
+@patch.object(Table, "get_columns")
+@patch.object(Table, "exists", return_value=True)
+@patch.object(Table, "_copy")
+def test_copy_into_qualify_never_rebuilds_live_table(
+    mock_copy, mock_exists, mock_get_columns, mock_drop, mock_sync_tags
+):
+    mock_get_columns.return_value = [Column(name="id", data_type="integer")]
+    mock_cursor = make_mock_cursor()
+    with patch("snowflake_utils.models.table.connect") as mock_connect:
+        mock_connect.return_value = make_mock_conn(cursor=mock_cursor)
+        test_table.copy_into(
+            path=path,
+            file_format=parquet_file_format,
+            storage_integration=storage_integration,
+            primary_keys=["id"],
+            qualify=True,
+        )
+
+    statements = [
+        " ".join(c.args[0].split()) for c in mock_cursor.execute.call_args_list
+    ]
+    assert not any(
+        s.lower().startswith("create or replace table public.pytest ")
+        for s in statements
+    )
+    assert any(
+        s.lower().startswith("create or replace table public.pytest_temp")
+        for s in statements
+    )
+    assert "drop table if exists PUBLIC.PYTEST_temp" in statements
+    assert any(
+        s.lower().startswith("merge into public.pytest as dest") for s in statements
+    )
+    # the live table is only ever written by the MERGE: COPY targets the temp table
+    assert "PUBLIC.PYTEST_temp" in mock_copy.call_args.args[0]
+    mock_sync_tags.assert_not_called()
+
+
 @patch("snowflake_utils.settings.connect")
 def test_single_column_update(mock_connect):
     mock_cursor = make_mock_cursor()