From 40ada94c3e9346462a7918aa8a8ffb5db1acf682 Mon Sep 17 00:00:00 2001
From: Oliver Mannion <125105+tekumara@users.noreply.github.com>
Date: Sat, 29 Jun 2024 11:04:38 +1000
Subject: [PATCH] add test_merge using test case from snowflake docs

---
 tests/test_fakes.py | 49 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/tests/test_fakes.py b/tests/test_fakes.py
index 2778ad0..cce5c6b 100644
--- a/tests/test_fakes.py
+++ b/tests/test_fakes.py
@@ -945,6 +945,55 @@ def test_identifier(cur: snowflake.connector.cursor.SnowflakeCursor):
     assert cur.fetchall() == [(1,)]
 
 
+def test_merge(conn: snowflake.connector.SnowflakeConnection):
+    *_, dcur = conn.execute_string(
+        """
+        CREATE OR REPLACE TABLE t1 (
+            t1Key INT PRIMARY KEY,
+            val VARCHAR(50),
+            status VARCHAR(20)
+        );
+
+        CREATE OR REPLACE TABLE t2 (
+            t2Key INT PRIMARY KEY,
+            newVal VARCHAR(50),
+            newStatus VARCHAR(20),
+            isNewStatus INT,
+            marked INT
+        );
+
+        INSERT INTO t1 (t1Key, val, status) VALUES
+        (1, 'Old Value 1', 'Old Status 1'),
+        (2, 'Old Value 2', 'Old Status 2'),
+        (3, 'Old Value 3', 'Old Status 3'),
+        (4, 'Old Value 4', 'Old Status 4');
+
+        INSERT INTO t2 (t2Key, newVal, newStatus, isNewStatus, marked) VALUES
+        (1, 'New Value 1', 'New Status 1', 1, 0),  -- Case: WHEN MATCHED AND t2.isNewStatus = 1 THEN UPDATE
+        (2, 'New Value 2', 'New Status 2', 0, 1),  -- Case: WHEN MATCHED AND t2.marked = 1 THEN DELETE
+        (3, 'New Value 3', 'New Status 3', 0, 0),  -- Case: WHEN MATCHED THEN UPDATE
+        (5, 'New Value 5', 'New Status 5', 0, 0);  -- Case: WHEN NOT MATCHED THEN INSERT
+
+        MERGE INTO t1 USING t2 ON t1.t1Key = t2.t2Key
+            WHEN MATCHED AND t2.marked = 1 THEN DELETE
+            WHEN MATCHED AND t2.isNewStatus = 1 THEN UPDATE SET val = t2.newVal, status = t2.newStatus
+            WHEN MATCHED THEN UPDATE SET val = t2.newVal
+            WHEN NOT MATCHED THEN INSERT (t1Key, val, status) VALUES (t2.t2Key, t2.newVal, t2.newStatus);
+        """,
+        cursor_class=snowflake.connector.cursor.DictCursor, # type: ignore see https://github.com/snowflakedb/snowflake-connector-python/issues/1984
+    )
+
+    assert dcur.fetchall() == [{"number of rows inserted": 1, "number of rows updated": 2, "number of rows deleted": 1}]
+
+    dcur.execute("select * from t1 order by t1key")
+    assert dcur.fetchall() == [
+        {"T1KEY": 1, "VAL": "New Value 1", "STATUS": "New Status 1"},
+        {"T1KEY": 3, "VAL": "New Value 3", "STATUS": "Old Status 3"},
+        {"T1KEY": 4, "VAL": "Old Value 4", "STATUS": "Old Status 4"},
+        {"T1KEY": 5, "VAL": "New Value 5", "STATUS": "New Status 5"},
+    ]
+
+
 def test_nop_regexes():
     with fakesnow.patch(nop_regexes=["^CALL.*"]), snowflake.connector.connect() as conn, conn.cursor() as cur:
         cur.execute("call this_procedure_does_not_exist('foo', 'bar);")