SNOW-862821 [Python Connector] Update query context cache response bo…

…dy to align with JDBC (#1646)

* fix

* clean up

* fix merge gate

* fix comments

* fix comments

src/snowflake/connector/_query_context_cache.py

@@ -3,7 +3,6 @@
 #
 from __future__ import annotations
 
-import json
 from functools import total_ordering
 from hashlib import md5
 from logging import getLogger
@@ -151,13 +150,13 @@ def _get_elements(self) -> Iterable[QueryContextElement]:
     def _last(self) -> QueryContextElement:
         return self._tree_set[-1]
 
-    def serialize_to_json(self) -> str:
+    def serialize_to_dict(self) -> dict:
         with self._lock:
-            logger.debug("serialize_to_json() called")
+            logger.debug("serialize_to_dict() called")
             self.log_cache_entries()
 
             if len(self._tree_set) == 0:
-                return ""
+                return {}  # we should return an empty dict
 
             try:
                 data = {
@@ -166,22 +165,24 @@ def serialize_to_json(self) -> str:
                             "id": qce.id,
                             "timestamp": qce.read_timestamp,
                             "priority": qce.priority,
-                            "context": qce.context,
+                            "context": {"base64Data": qce.context}
+                            if qce.context is not None
+                            else {},
                         }
                         for qce in self._tree_set
                     ]
                 }
-                # Serialize the data to JSON
-                serialized_data = json.dumps(data)
+                # Because on GS side, `context` field is an object with `base64Data`  string member variable,
+                # we should serialize `context` field to an object instead of string directly to stay consistent with GS side.
 
-                logger.debug(
-                    f"serialize_to_json(): data to send to server {serialized_data}"
-                )
+                logger.debug(f"serialize_to_dict(): data to send to server {data}")
 
-                return serialized_data
+                # query context shoule be an object field of the HTTP request body JSON and on GS side. here we should only return a dict
+                # and let the outer HTTP request body to convert the entire big dict to a single JSON.
+                return data
             except Exception as e:
-                logger.debug(f"serialize_to_json(): Exception {e}")
-                return ""
+                logger.debug(f"serialize_to_dict(): Exception {e}")
+                return {}
 
     def deserialize_json_dict(self, data: dict) -> None:
         with self._lock:

src/snowflake/connector/connection.py

@@ -1107,8 +1107,10 @@ def cmd_query(
             # binding parameters. This is for qmarks paramstyle.
             data["bindings"] = binding_params
         if not _no_results:
-            # not an async query
-            data["queryContext"] = self.get_query_context()
+            # not an async query.
+            queryContext = self.get_query_context()
+            #  Here queryContextDTO should be a dict object field, same with `parameters` field
+            data["queryContextDTO"] = queryContext
         client = "sfsql_file_transfer" if is_file_transfer else "sfsql"
 
         if logger.getEffectiveLevel() <= logging.DEBUG:
@@ -1647,10 +1649,10 @@ def initialize_query_context_cache(self) -> None:
         if not self.is_query_context_cache_disabled:
             self.query_context_cache = QueryContextCache(self.query_context_cache_size)
 
-    def get_query_context(self) -> str | None:
+    def get_query_context(self) -> dict | None:
         if self.is_query_context_cache_disabled:
             return None
-        return self.query_context_cache.serialize_to_json()
+        return self.query_context_cache.serialize_to_dict()
 
     def set_query_context(self, data: dict) -> None:
         if not self.is_query_context_cache_disabled:

test/unit/test_query_context_cache.py

@@ -51,6 +51,37 @@ def reset_data(self):
         self.priorities = [BASE_PRIORITY + i for i in range(self.capacity)]
 
 
+def serialize_to_json(qcc: QueryContextCache) -> str:
+    """
+    This function is used for testing the deserialize_to_dict function is correct. Note that `QueryContextCache::serialize_to_dict` function is used in sending the query context back to GS.
+    It has slight difference with this function on `context` field.
+    """
+    with qcc._lock:
+        qcc.log_cache_entries()
+
+        if len(qcc._tree_set) == 0:
+            return ""
+
+        try:
+            data = {
+                "entries": [
+                    {
+                        "id": qce.id,
+                        "timestamp": qce.read_timestamp,
+                        "priority": qce.priority,
+                        "context": qce.context,
+                    }
+                    for qce in qcc._tree_set
+                ]
+            }
+            # Serialize the data to JSON
+            serialized_data = json.dumps(data)
+
+            return serialized_data
+        except Exception:
+            return ""
+
+
 @pytest.fixture()
 def expected_data() -> ExpectedQCCData:
     data = ExpectedQCCData(MAX_CAPACITY)
@@ -376,7 +407,8 @@ def test_serialization_deserialization_with_null_context(
 ):
     assert_cache_with_data(qcc_with_data_null_context, expected_data_with_null_context)
 
-    data = qcc_with_data_null_context.serialize_to_json()
+    data = serialize_to_json(qcc_with_data_null_context)
+
     qcc_with_data_null_context.clear_cache()
     assert len(qcc_with_data_null_context) == 0
 
@@ -390,7 +422,7 @@ def test_serialization_deserialization(
 ):
     assert_cache_with_data(qcc_with_data, expected_data)
 
-    data = qcc_with_data.serialize_to_json()
+    data = serialize_to_json(qcc_with_data)
     qcc_with_data.clear_cache()
     assert len(qcc_with_data) == 0