Fix flaky QCC integration test by mocking server response (#2866)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: sfc-gh-pczajka

test/integ/test_connection.py

@@ -1280,55 +1280,6 @@ def test_disable_query_context_cache(conn_cnx) -> None:
         assert conn.query_context_cache is None
 
 
-@pytest.mark.aws
-@pytest.mark.skipolddriver
-def test_qcc_updated_on_failed_query_with_hybrid_table(conn_cnx) -> None:
-    """SNOW-3010877: QCC must be merged even when a query fails.
-
-    Reproduces the exact hybrid table scenario from the bug report:
-    a duplicate-key INSERT inside a transaction fails, but the server
-    still sends queryContext in the error response.  The client must
-    merge it so that subsequent queries on different GS nodes see the
-    updated sessionDPO / read snapshot watermark.
-
-    Note: full behavioural verification depends on the server returning
-    queryContext in error responses, which is gated behind a feature flag.
-    On accounts where the flag is off the QCC simply won't change, so the
-    >= assertion still passes.  The unit tests with mocked responses are
-    the authoritative behavioural tests for this feature.
-    """
-    with conn_cnx() as conn:
-        cur = conn.cursor()
-        table_name = f"test_qcc_hybrid_{random_string(5)}"
-        try:
-            cur.execute(
-                f"create or replace hybrid table {table_name} "
-                f"(pk text primary key, val text)"
-            )
-
-            # Q1 — successful insert
-            cur.execute(f"insert into {table_name} values ('k1', 'v1')")
-            qcc_after_success = len(conn.query_context_cache)
-            assert (
-                qcc_after_success > 0
-            ), "QCC should have entries after successful query"
-
-            # Q2 — duplicate key insert, expected to fail
-            with pytest.raises(ProgrammingError):
-                cur.execute(f"insert into {table_name} values ('k1', 'v1')")
-
-            # The critical assertion: QCC must still have been updated
-            # despite the query failure.  The entry count must not have
-            # decreased — it should stay the same or grow.
-            qcc_after_failure = len(conn.query_context_cache)
-            assert qcc_after_failure >= qcc_after_success, (
-                f"QCC should not shrink after a failed query: "
-                f"before={qcc_after_success}, after={qcc_after_failure}"
-            )
-        finally:
-            cur.execute(f"drop table if exists {table_name}")
-
-
 @pytest.mark.aws
 @pytest.mark.skipolddriver
 def test_qcc_tracks_multi_database_hybrid_tables(conn_cnx) -> None:

test/unit/test_query_context_cache.py

@@ -516,3 +516,81 @@ def test_qcc_updated_regardless_of_query_success(success):
     entry_ids = {e.id for e in conn.query_context_cache._get_elements()}
     assert 1 in entry_ids, "Original entry should still be present"
     assert 2 in entry_ids, "New entry from server response should have been merged"
+
+
+def test_qcc_not_shrunk_after_failed_query_with_hybrid_table():
+    """SNOW-3010877: QCC must not shrink when a query on a hybrid table fails.
+
+    Simulates the hybrid table scenario: a successful INSERT populates
+    the QCC with 2 entries (main + DPO), then a duplicate-key INSERT
+    fails but the server still sends queryContext in the error response.
+    The connector must merge it so the cache does not shrink.
+    """
+    from snowflake.connector.connection import SnowflakeConnection
+
+    conn = MagicMock()
+    conn._disable_query_context_cache = False
+    conn.is_query_context_cache_disabled = False
+    conn.query_context_cache = QueryContextCache(5)
+    conn.query_context_cache_size = 5
+
+    conn.get_query_context = types.MethodType(
+        SnowflakeConnection.get_query_context, conn
+    )
+    conn.set_query_context = types.MethodType(
+        SnowflakeConnection.set_query_context, conn
+    )
+
+    # Q1 — successful insert: server returns queryContext with 2 entries
+    conn.rest.request.return_value = {
+        "success": True,
+        "data": {
+            "queryId": "success-query-id",
+            "queryContext": {
+                "entries": [
+                    {"id": 0, "timestamp": 100, "priority": 0, "context": "main_ctx"},
+                    {"id": 1, "timestamp": 100, "priority": 1, "context": "ht_dpo_ctx"},
+                ]
+            },
+        },
+    }
+    SnowflakeConnection.cmd_query(
+        conn, "INSERT INTO ht VALUES ('k1','v1')", 1, uuid.uuid4()
+    )
+    qcc_after_success = len(conn.query_context_cache)
+    assert qcc_after_success == 2, "QCC should have 2 entries after successful query"
+
+    # Q2 — duplicate key insert fails, but server still sends queryContext
+    conn.rest.request.return_value = {
+        "success": False,
+        "code": "100072",
+        "message": "Duplicate key value violates unique constraint",
+        "data": {
+            "queryId": "failed-query-id",
+            "queryContext": {
+                "entries": [
+                    {
+                        "id": 0,
+                        "timestamp": 200,
+                        "priority": 0,
+                        "context": "main_ctx_v2",
+                    },
+                    {
+                        "id": 1,
+                        "timestamp": 200,
+                        "priority": 1,
+                        "context": "ht_dpo_ctx_v2",
+                    },
+                ]
+            },
+        },
+    }
+    SnowflakeConnection.cmd_query(
+        conn, "INSERT INTO ht VALUES ('k1','v1')", 2, uuid.uuid4()
+    )
+
+    qcc_after_failure = len(conn.query_context_cache)
+    assert qcc_after_failure >= qcc_after_success, (
+        f"QCC should not shrink after a failed query: "
+        f"before={qcc_after_success}, after={qcc_after_failure}"
+    )