Revert query_tags changes for CI baseline test

Author: jiabin-hu

examples/query_tags_example.py

@@ -16,7 +16,7 @@
 
 Special cases:
 - If a value is None, only the key is included (no colon or value)
-- Special characters (comma, colon and backslash) in values are automatically escaped
+- Special characters (:, ,, \\) in values are automatically escaped
 - Keys are not escaped (should be controlled identifiers)
 """
 
@@ -74,11 +74,6 @@
         result = cursor.fetchone()
         print(f"  Special Chars Query Result: {result[0]}")
 
-        # Query 4: No tags (demonstrates tags don't persist from previous queries)
-        cursor.execute("SELECT 4")
-        result = cursor.fetchone()
-        print(f"  No Tags Query Result: {result[0]}")
-
 print()
 
 # Example 3: Async execution with query tags
@@ -91,31 +86,11 @@
 
     with connection.cursor() as cursor:
         cursor.execute_async(
-            "SELECT 5",
+            "SELECT 4",
             query_tags={"team": "data-eng", "mode": "async"}
         )
         cursor.get_async_execution_result()
         result = cursor.fetchone()
         print(f"  Async Query Result: {result[0]}")
 
-print()
-
-# Example 4: executemany with query tags
-print("Example 4: executemany with query tags")
-with sql.connect(
-    server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
-    http_path=os.getenv("DATABRICKS_HTTP_PATH"),
-    access_token=os.getenv("DATABRICKS_TOKEN"),
-) as connection:
-
-    with connection.cursor() as cursor:
-        # Execute multiple queries with the same tags
-        cursor.executemany(
-            "SELECT ?",
-            [[6], [7], [8]],
-            query_tags={"team": "data-eng", "batch": "executemany"}
-        )
-        result = cursor.fetchone()
-        print(f"  Executemany Query Result (last): {result[0]}")
-
-print("\n=== Query Tags Example Complete ===")
+print("\n=== Query Tags Example Complete ===")

src/databricks/sql/backend/sea/backend.py

@@ -463,9 +463,6 @@ def execute_command(
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
         row_limit: Optional[int] = None,
-        query_tags: Optional[
-            Dict[str, Optional[str]]
-        ] = None,  # TODO: implement query_tags for SEA backend
     ) -> Union[SeaResultSet, None]:
         """
         Execute a SQL command using the SEA backend.

src/databricks/sql/client.py

@@ -1459,12 +1459,7 @@ def get_async_execution_result(self):
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-    def executemany(
-        self,
-        operation,
-        seq_of_parameters,
-        query_tags: Optional[Dict[str, Optional[str]]] = None,
-    ):
+    def executemany(self, operation, seq_of_parameters):
         """
         Execute the operation once for every set of passed in parameters.
 
@@ -1473,14 +1468,10 @@ def executemany(
 
         Only the final result set is retained.
 
-        :param query_tags: Optional dictionary of query tags to apply for all queries in this batch.
-            Tags are key-value pairs that can be used to identify and categorize queries.
-            Example: {"team": "data-eng", "application": "etl"}
-
         :returns self
         """
         for parameters in seq_of_parameters:
-            self.execute(operation, parameters, query_tags=query_tags)
+            self.execute(operation, parameters)
         return self
 
     @log_latency(StatementType.METADATA)

src/databricks/sql/utils.py

@@ -898,17 +898,15 @@ def concat_table_chunks(
         return pyarrow.concat_tables(table_chunks)
 
 
-def serialize_query_tags(
-    query_tags: Optional[Dict[str, Optional[str]]]
-) -> Optional[str]:
+def serialize_query_tags(query_tags: Optional[Dict[str, Optional[str]]]) -> Optional[str]:
     """
     Serialize query_tags dictionary to a string format.
 
     Format: "key1:value1,key2:value2"
     Special cases:
     - If value is None, omit the colon and value (e.g., "key1:value1,key2,key3:value3")
     - Escape special characters (:, ,, \\) in values with a leading backslash
-    - Backslashes in keys are escaped; other special characters in keys are not escaped
+    - Keys are not escaped (assumed to be controlled identifiers)
 
     Args:
         query_tags: Dictionary of query tags where keys are strings and values are optional strings
@@ -922,21 +920,20 @@ def serialize_query_tags(
     def escape_value(value: str) -> str:
         """Escape special characters in tag values."""
         # Escape backslash first to avoid double-escaping
-        value = value.replace("\\", r"\\")
+        value = value.replace("\\", "\\\\")
         # Escape colon and comma
-        value = value.replace(":", r"\:")
-        value = value.replace(",", r"\,")
+        value = value.replace(":", "\\:")
+        value = value.replace(",", "\\,")
         return value
 
     serialized_parts = []
     for key, value in query_tags.items():
-        escaped_key = key.replace("\\", r"\\")
         if value is None:
             # No colon or value when value is None
-            serialized_parts.append(escaped_key)
+            serialized_parts.append(key)
         else:
             escaped_value = escape_value(value)
-            serialized_parts.append(f"{escaped_key}:{escaped_value}")
+            serialized_parts.append(f"{key}:{escaped_value}")
 
     return ",".join(serialized_parts)
 

tests/unit/test_util.py

@@ -180,19 +180,19 @@ def test_serialize_query_tags_with_special_chars(self):
         query_tags = {
             "key1": "value:with:colons",
             "key2": "value,with,commas",
-            "key3": r"value\with\backslashes",
+            "key3": "value\\with\\backslashes",
         }
         result = serialize_query_tags(query_tags)
         assert (
             result
-            == r"key1:value\:with\:colons,key2:value\,with\,commas,key3:value\\with\\backslashes"
+            == "key1:value\\:with\\:colons,key2:value\\,with\\,commas,key3:value\\\\with\\\\backslashes"
         )
 
     def test_serialize_query_tags_with_mixed_special_chars(self):
         """Test query tags with mixed special characters"""
-        query_tags = {"key1": r"a:b,c\d"}
+        query_tags = {"key1": "a:b,c\\d"}
         result = serialize_query_tags(query_tags)
-        assert result == r"key1:a\:b\,c\\d"
+        assert result == "key1:a\\:b\\,c\\\\d"
 
     def test_serialize_query_tags_empty_dict(self):
         """Test serialization with empty dictionary"""
@@ -206,17 +206,17 @@ def test_serialize_query_tags_none(self):
         assert result is None
 
     def test_serialize_query_tags_with_special_chars_in_key(self):
-        """Test query tags with special characters in keys (only backslashes are escaped in keys)"""
+        """Test query tags with special characters in keys (keys are not escaped)"""
         query_tags = {
             "key:with:colons": "value1",
             "key,with,commas": "value2",
-            r"key\with\backslashes": "value3",
+            "key\\with\\backslashes": "value3",
         }
         result = serialize_query_tags(query_tags)
-        # Only backslashes are escaped in keys; colons and commas in keys are not escaped
+        # Keys are not escaped, only values are
         assert (
             result
-            == r"key:with:colons:value1,key,with,commas:value2,key\\with\\backslashes:value3"
+            == "key:with:colons:value1,key,with,commas:value2,key\\with\\backslashes:value3"
         )
 
     def test_serialize_query_tags_all_none_values(self):