From 0917d6fa31838a1fd811b41069d7799e36160234 Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 1 Oct 2026 12:21:25 +0200 Subject: [PATCH 1/4] fix: dedupe copy_into(qualify=True) via temp table and merge copy_into(qualify=True) appended into the live table and then ran `create or replace table ... qualify row_number() = 1`, so readers saw duplicate keys for the minutes between the two steps. Route it through _merge instead: COPY into _temp, dedupe there, MERGE into the live table. full_refresh keeps the old path. - _merge takes sync_tags so the caller's flag is honoured - re-sync tags after the first-load qualify, which drops them Claude-Session: https://claude.ai/code/session_011PFLPuv77fp5M3JXdZvW1v --- README.md | 2 +- snowflake_utils/models/table.py | 33 +++++++++- tests/test_models.py | 109 ++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index cb5331d..aa2a4e5 100644 --- a/README.md +++ b/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 `
_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 diff --git a/snowflake_utils/models/table.py b/snowflake_utils/models/table.py index 2644019..2ceaf78 100644 --- a/snowflake_utils/models/table.py +++ b/snowflake_utils/models/table.py @@ -253,7 +253,31 @@ def copy_into( {files_clause} {self._include_metadata()} """ - if qualify: + if qualify and not full_refresh: + # Dedup in a temp table and merge, so readers never see the live + # table holding duplicates between the COPY and the QUALIFY. + 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, + ) + + return self._merge( + copy_callable, + primary_keys, + replication_keys, + qualify=True, + sync_tags=sync_tags, + ) + elif qualify: self._copy( copy_query, path, @@ -331,13 +355,16 @@ 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"}) @@ -364,7 +391,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) diff --git a/tests/test_models.py b/tests/test_models.py index 41abe09..d17d269 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -407,6 +407,115 @@ 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() + + +@pytest.mark.parametrize("create_table", [True, False]) +@patch.object(Table, "_copy") +@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 = mock_copy.call_args_list + # positional args of _copy: ..., sync_tags, stage, create_table, copy_grants + assert temp_call.args[5] is False and temp_call.args[7] is True + assert live_call.args[5] is True and live_call.args[7] is create_table + assert "COPY INTO PUBLIC.PYTEST_temp (id)" in temp_call.args[0].replace("\n", " ") + + +@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 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() From 9b128d16a3005fcec8dbe41008349de98b34be40 Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 1 Oct 2026 12:26:11 +0200 Subject: [PATCH 2/4] refactor: simplify copy_into qualify path and harden tests - single _copy call for full_refresh+qualify - drop stale
_temp before the temp COPY - bind _copy args by name in tests Claude-Session: https://claude.ai/code/session_011PFLPuv77fp5M3JXdZvW1v --- snowflake_utils/models/table.py | 59 ++++++++++++++------------------- tests/test_models.py | 17 ++++++---- 2 files changed, 36 insertions(+), 40 deletions(-) diff --git a/snowflake_utils/models/table.py b/snowflake_utils/models/table.py index 2ceaf78..51f988b 100644 --- a/snowflake_utils/models/table.py +++ b/snowflake_utils/models/table.py @@ -254,8 +254,7 @@ def copy_into( {self._include_metadata()} """ if qualify and not full_refresh: - # Dedup in a temp table and merge, so readers never see the live - # table holding duplicates between the COPY and the QUALIFY. + # 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, @@ -277,39 +276,29 @@ def copy_callable(table: Table, sync_tags: bool) -> None: qualify=True, sync_tags=sync_tags, ) - elif 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 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, + + 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 @@ -368,6 +357,8 @@ def _merge( 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: diff --git a/tests/test_models.py b/tests/test_models.py index d17d269..40fc2c0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,3 +1,4 @@ +import inspect import logging import os from datetime import datetime @@ -449,8 +450,12 @@ def test_copy_into_qualify_full_refresh_keeps_copy_then_qualify( 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") +@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 @@ -470,11 +475,10 @@ def test_copy_into_qualify_always_creates_temp_table( ) copy_callable(test_table, sync_tags=True) - temp_call, live_call = mock_copy.call_args_list - # positional args of _copy: ..., sync_tags, stage, create_table, copy_grants - assert temp_call.args[5] is False and temp_call.args[7] is True - assert live_call.args[5] is True and live_call.args[7] is create_table - assert "COPY INTO PUBLIC.PYTEST_temp (id)" in temp_call.args[0].replace("\n", " ") + 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, "sync_tags") @@ -508,6 +512,7 @@ def test_copy_into_qualify_never_rebuilds_live_table( 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 ) From 1b87fdbca05361a6be83b3d5adb20df74a20283d Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 1 Oct 2026 12:28:57 +0200 Subject: [PATCH 3/4] fix: skip redundant create-table in _merge for existing tables With no table_structure the statement needs the stage/file format that are only set on the temp copy, so merging into an existing table raised "Call setup_stage to set the stage". The branch already knows the table exists. Claude-Session: https://claude.ai/code/session_011PFLPuv77fp5M3JXdZvW1v --- snowflake_utils/models/table.py | 3 --- tests/test_models.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/snowflake_utils/models/table.py b/snowflake_utils/models/table.py index 51f988b..7d7eccb 100644 --- a/snowflake_utils/models/table.py +++ b/snowflake_utils/models/table.py @@ -367,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) diff --git a/tests/test_models.py b/tests/test_models.py index 40fc2c0..c0ce976 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -481,6 +481,34 @@ def test_copy_into_qualify_always_creates_temp_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") From 358d827c614cd0cbe1ee60b84f039534f15b300f Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Thu, 1 Oct 2026 14:28:43 +0200 Subject: [PATCH 4/4] fix: address review on qualify merge (loaded columns, replication keys, temp name) - update only the loaded columns when target_columns is given, so unloaded columns are not overwritten with NULL - with replication keys, only update live rows that are not newer than the incoming row (matches the old qualify ordering, NULL-safe) - use a unique
_temp_ per run so concurrent COPYs don't collide - drop the temp table when the load or merge fails Claude-Session: https://claude.ai/code/session_011PFLPuv77fp5M3JXdZvW1v --- README.md | 2 +- snowflake_utils/models/column.py | 12 ++++ snowflake_utils/models/table.py | 86 +++++++++++++++++------- tests/test_models.py | 111 ++++++++++++++++++++++++++++++- 4 files changed, 182 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index aa2a4e5..0266f5f 100644 --- a/README.md +++ b/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 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 `
_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. + - 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 uniquely named `
_temp_` table, deduped there and merged into the destination (updating only the loaded columns, and only rows that are not older by the replication keys), 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 diff --git a/snowflake_utils/models/column.py b/snowflake_utils/models/column.py index fa9a19f..0b1002e 100644 --- a/snowflake_utils/models/column.py +++ b/snowflake_utils/models/column.py @@ -25,6 +25,18 @@ def tmp(x: str) -> str: ) +def _is_newer_or_equal(keys: list[str]) -> str: + *head, last = (k.upper() for k in keys) + expr = f'(dest."{last}" is null or tmp."{last}" >= dest."{last}")' + for key in reversed(head): + dest, tmp = f'dest."{key}"', f'tmp."{key}"' + expr = ( + f"(({dest} is null and {tmp} is not null) or {tmp} > {dest}" + f" or (equal_null({tmp}, {dest}) and {expr}))" + ) + return expr + + def _inserts(columns: list[Column], old_columns: dict[str, str]) -> str: return ",".join( _possibly_cast(f'tmp."{c.name}"', old_columns.get(c.name), c.data_type) diff --git a/snowflake_utils/models/table.py b/snowflake_utils/models/table.py index 7d7eccb..582f420 100644 --- a/snowflake_utils/models/table.py +++ b/snowflake_utils/models/table.py @@ -2,13 +2,21 @@ from collections import defaultdict from functools import partial from typing import ClassVar +from uuid import uuid4 from pydantic import BaseModel, Field from snowflake.connector.cursor import SnowflakeCursor from ..queries import execute_statement from ..settings import SnowflakeSettings, connect, governance_settings -from .column import Column, MetadataColumn, _inserts, _matched, _type_cast +from .column import ( + Column, + MetadataColumn, + _inserts, + _is_newer_or_equal, + _matched, + _type_cast, +) from .enums import MatchByColumnName, TagLevel from .file_format import FileFormat, InlineFileFormat from .table_structure import TableStructure @@ -275,6 +283,7 @@ def copy_callable(table: Table, sync_tags: bool) -> None: replication_keys, qualify=True, sync_tags=sync_tags, + target_columns=target_columns, ) result = self._copy( @@ -345,6 +354,7 @@ def _merge( replication_keys: list[str] | None = None, qualify: bool = False, sync_tags: bool = True, + target_columns: list[str] | None = None, ) -> None: with connect() as connection: cursor = connection.cursor() @@ -356,32 +366,54 @@ def _merge( 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: + temp_table = self.model_copy( + update={"name": f"{self.name}_temp_{uuid4().hex[:8]}"} + ) + try: + copy_callable(temp_table, sync_tags=False) + if qualify: + with connect() as connection: + cursor = connection.cursor() + temp_table.qualify(cursor, primary_keys, replication_keys) + with connect() as connection: cursor = connection.cursor() - temp_table.qualify(cursor, primary_keys, replication_keys) - - with connect() as connection: - cursor = connection.cursor() - old_columns = {x.name: x.data_type for x in self.get_columns(cursor)} - new_columns = temp_table.get_columns(cursor) - - for column in new_columns: - if column.name not in old_columns: - self.add_column(cursor, column) - - cursor.execute( - self._merge_statement( - temp_table, new_columns, old_columns, primary_keys + old_columns = {x.name: x.data_type for x in self.get_columns(cursor)} + new_columns = temp_table.get_columns(cursor) + + for column in new_columns: + if column.name not in old_columns: + self.add_column(cursor, column) + + loaded = { + c.upper() + for c in [ + *(target_columns or []), + *primary_keys, + *(m.name for m in self.include_metadata), + ] + } + merge_columns = [ + c + for c in new_columns + if not target_columns or c.name.upper() in loaded + ] + cursor.execute( + self._merge_statement( + temp_table, + merge_columns, + old_columns, + primary_keys, + replication_keys if qualify else None, + ) ) - ) - if sync_tags and self.table_structure: - self.sync_tags(cursor) - temp_table.drop(cursor) + if sync_tags and self.table_structure: + self.sync_tags(cursor) + temp_table.drop(cursor) + except BaseException: + with connect() as connection: + connection.cursor().execute(f"drop table if exists {temp_table.fqn}") + raise def merge( self, @@ -482,10 +514,14 @@ def _merge_statement( columns: list[Column], old_columns: dict[str, str], primary_keys: list[str], + replication_keys: list[str] | None = None, ) -> str: pkes = " and ".join( f'dest."{c.upper()}" = tmp."{c.upper()}"' for c in primary_keys ) + update_condition = ( + f" and {_is_newer_or_equal(replication_keys)}" if replication_keys else "" + ) matched = _matched(columns, old_columns) column_names = ",".join(f'"{c.name}"' for c in columns) inserts = _inserts(columns, old_columns) @@ -498,7 +534,7 @@ def _merge_statement( merge into {self.fqn} as dest using {temp_table.fqn} tmp ON {pkes} - when matched then update set {matched} + when matched{update_condition} then update set {matched} when not matched then insert ({column_names}) VALUES ({inserts}) """ diff --git a/tests/test_models.py b/tests/test_models.py index c0ce976..d0576df 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,6 +1,7 @@ import inspect import logging import os +import re from datetime import datetime from unittest.mock import MagicMock, patch @@ -425,7 +426,7 @@ def test_copy_into_qualify_merges_instead_of_copying_into_live_table( mock_copy.assert_not_called() mock_merge.assert_called_once() _, kwargs = mock_merge.call_args - assert kwargs == {"qualify": True, "sync_tags": True} + assert kwargs == {"qualify": True, "sync_tags": True, "target_columns": None} @patch.object(Table, "qualify") @@ -540,15 +541,119 @@ def test_copy_into_qualify_never_rebuilds_live_table( 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] + assert "PUBLIC.PYTEST_temp_" in mock_copy.call_args.args[0] + mock_drop.assert_called_once() mock_sync_tags.assert_not_called() +def run_qualified_copy_into(table: Table, **kwargs) -> MagicMock: + 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) + table.copy_into( + path=path, + file_format=parquet_file_format, + storage_integration=storage_integration, + primary_keys=["id"], + qualify=True, + **kwargs, + ) + return mock_cursor + + +def executed_statements(mock_cursor: MagicMock) -> list[str]: + return [" ".join(c.args[0].split()) for c in mock_cursor.execute.call_args_list] + + +@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_only_updates_loaded_columns( + mock_copy, mock_exists, mock_get_columns, mock_drop +): + mock_get_columns.return_value = [ + Column(name=n, data_type="text") for n in ("id", "name", "last_name") + ] + + statements = executed_statements( + run_qualified_copy_into(test_table, target_columns=["name"]) + ) + + merge = next(s for s in statements if s.lower().startswith("merge into")) + update_clause = merge.split("when matched")[1].split("when not matched")[0] + assert 'dest."name"' in update_clause and 'dest."id"' in update_clause + assert "last_name" not in update_clause + + +def test_merge_statement_updates_only_newer_rows_with_replication_keys(): + columns = [Column(name="id", data_type="text")] + temp = test_table.model_copy(update={"name": "PYTEST_temp"}) + + with_keys = " ".join( + test_table._merge_statement( + temp, columns, {}, ["id"], ["updated_at", "etl_file_ingested_at"] + ).split() + ) + without_keys = " ".join( + test_table._merge_statement(temp, columns, {}, ["id"]).split() + ) + + condition = with_keys.split("when matched")[1].split("then update")[0] + assert 'tmp."UPDATED_AT" > dest."UPDATED_AT"' in condition + assert 'tmp."ETL_FILE_INGESTED_AT" >= dest."ETL_FILE_INGESTED_AT"' in condition + assert "when matched then update" in without_keys + + +@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_uses_a_new_temp_table_per_run( + mock_copy, mock_exists, mock_get_columns, mock_drop +): + mock_get_columns.return_value = [Column(name="id", data_type="text")] + + run_qualified_copy_into(test_table) + run_qualified_copy_into(test_table) + + temp_names = { + re.search(r"PYTEST_temp_[0-9a-f]{8}", c.args[0]).group() + for c in mock_copy.call_args_list + } + assert len(temp_names) == 2 + + +@patch.object(Table, "get_columns") +@patch.object(Table, "exists", return_value=True) +@patch.object(Table, "_copy", side_effect=RuntimeError("copy failed")) +def test_copy_into_qualify_drops_temp_table_when_the_load_fails( + mock_copy, mock_exists, mock_get_columns +): + 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) + with pytest.raises(RuntimeError, match="copy failed"): + test_table.copy_into( + path=path, + file_format=parquet_file_format, + storage_integration=storage_integration, + primary_keys=["id"], + qualify=True, + ) + + statements = executed_statements(mock_cursor) + assert any( + re.fullmatch(r"drop table if exists PUBLIC\.PYTEST_temp_[0-9a-f]{8}", s) + for s in statements + ) + assert not any(s.lower().startswith("merge into") for s in statements) + + @patch("snowflake_utils.settings.connect") def test_single_column_update(mock_connect): mock_cursor = make_mock_cursor()