address review comments

Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>

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 (:, ,, \\) in values are automatically escaped
+- Special characters (comma, colon and backslash) in values are automatically escaped
 - Keys are not escaped (should be controlled identifiers)
 """
 
@@ -118,4 +118,4 @@
         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/utils.py

@@ -908,7 +908,7 @@ def serialize_query_tags(
     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
-    - Keys are not escaped (assumed to be controlled identifiers)
+    - Backslashes in keys are escaped; other special characters in keys are not escaped
 
     Args:
         query_tags: Dictionary of query tags where keys are strings and values are optional strings
@@ -922,20 +922,21 @@ 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("\\", "\\\\")
+        value = value.replace("\\", r"\\")
         # Escape colon and comma
-        value = value.replace(":", "\\:")
-        value = value.replace(",", "\\,")
+        value = value.replace(":", r"\:")
+        value = value.replace(",", r"\,")
         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(key)
+            serialized_parts.append(escaped_key)
         else:
             escaped_value = escape_value(value)
-            serialized_parts.append(f"{key}:{escaped_value}")
+            serialized_parts.append(f"{escaped_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": "value\\with\\backslashes",
+            "key3": r"value\with\backslashes",
         }
         result = serialize_query_tags(query_tags)
         assert (
             result
-            == "key1:value\\:with\\:colons,key2:value\\,with\\,commas,key3:value\\\\with\\\\backslashes"
+            == r"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": "a:b,c\\d"}
+        query_tags = {"key1": r"a:b,c\d"}
         result = serialize_query_tags(query_tags)
-        assert result == "key1:a\\:b\\,c\\\\d"
+        assert result == r"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 (keys are not escaped)"""
+        """Test query tags with special characters in keys (only backslashes are escaped in keys)"""
         query_tags = {
             "key:with:colons": "value1",
             "key,with,commas": "value2",
-            "key\\with\\backslashes": "value3",
+            r"key\with\backslashes": "value3",
         }
         result = serialize_query_tags(query_tags)
-        # Keys are not escaped, only values are
+        # Only backslashes are escaped in keys; colons and commas in keys are not escaped
         assert (
             result
-            == "key:with:colons:value1,key,with,commas:value2,key\\with\\backslashes:value3"
+            == r"key:with:colons:value1,key,with,commas:value2,key\\with\\backslashes:value3"
         )
 
     def test_serialize_query_tags_all_none_values(self):