improve tests and refactor

Author: joe-clickhouse

clickhouse_connect/driver/httpclient.py

@@ -197,18 +197,6 @@ def __init__(self,
                 self._setting_status('http_headers_progress_interval_ms').is_writable:
             self._progress_interval = str(min(120000, max(10000, (send_receive_timeout - 5) * 1000)))
 
-    def _ensure_query_id(self, settings: Optional[dict]) -> Optional[dict]:
-        if not self._autogenerate_query_id:
-            return settings
-
-        if settings is None:
-            settings = {}
-
-        if "query_id" not in settings:
-            settings["query_id"] = str(uuid.uuid4())
-
-        return settings
-
     def set_client_setting(self, key, value):
         str_value = self._validate_setting(key, value, common.get_setting('invalid_setting_action'))
         if str_value is not None:
@@ -233,8 +221,6 @@ def _prep_query(self, context: QueryContext):
         return final_query + fmt
 
     def _query_with_context(self, context: QueryContext) -> QueryResult:
-        context.settings = self._ensure_query_id(context.settings)
-
         headers = {}
         params = {}
         if self.database:
@@ -327,8 +313,6 @@ def data_insert(self, context: InsertContext) -> QuerySummary:
             logger.debug('No data included in insert, skipping')
             return QuerySummary()
 
-        context.settings = self._ensure_query_id(context.settings)
-
         def error_handler(resp: HTTPResponse):
             # If we actually had a local exception when building the insert, throw that instead
             if context.insert_exception:
@@ -364,8 +348,6 @@ def raw_insert(self, table: str = None,
         """
         See BaseClient doc_string for this method
         """
-        settings = self._ensure_query_id(settings)
-
         params = {}
         headers = {'Content-Type': 'application/octet-stream'}
         if compression:
@@ -409,8 +391,6 @@ def command(self,
         """
         See BaseClient doc_string for this method
         """
-        settings = self._ensure_query_id(settings)
-
         cmd, params = bind_query(cmd, parameters, self.server_tz)
         headers = {}
         payload = None
@@ -521,6 +501,10 @@ def _raw_request(self,
             final_params['http_headers_progress_interval_ms'] = self._progress_interval
         final_params = dict_copy(self.params, final_params)
         final_params = dict_copy(final_params, params)
+
+        if self._autogenerate_query_id and "query_id" not in final_params:
+            final_params["query_id"] = str(uuid.uuid4())
+
         url = f'{self.url}?{urlencode(final_params)}'
         kwargs = {
             'headers': headers,

tests/integration_tests/test_client.py

@@ -20,6 +20,15 @@
 klm,1,"""
 
 
+def _is_valid_uuid_v4(id_string: str) -> bool:
+    """Helper function to validate that a string is a valid UUID v4"""
+    try:
+        parsed_uuid = uuid.UUID(id_string)
+        return parsed_uuid.version == 4
+    except (ValueError, AttributeError):
+        return False
+
+
 def test_ping(test_client: Client):
     assert test_client.ping() is True
 
@@ -388,41 +397,75 @@ def test_role_setting_works(test_client: Client, test_config: TestConfig):
     assert res.result_rows == [([role_limited],)]
 
 
-# pylint: disable=protected-access
-def test_autogenerate_query_id(test_client: Client, test_table_engine: str, test_config: TestConfig):
-    def _is_valid_uuid_v4_string(id_string: str) -> bool:
-        parsed_uuid = uuid.UUID(id_string)
-        return parsed_uuid.version == 4
-
+def test_query_id_autogeneration(test_client: Client, test_table_engine: str):
+    """Test that query_id is auto-generated for query(), command(), and insert() methods"""
     result = test_client.query("SELECT 1")
-    assert _is_valid_uuid_v4_string(result.query_id)
+    assert _is_valid_uuid_v4(result.query_id)
 
-    manual_query_id = "test_manual_query_id_12345"
-    result = test_client.query("SELECT 2", settings={"query_id": manual_query_id})
-    assert result.query_id == manual_query_id
+    summary = test_client.command("DROP TABLE IF EXISTS test_query_id_nonexistent")
+    assert _is_valid_uuid_v4(summary.query_id())
 
-    summary = test_client.command("DROP TABLE IF EXISTS does_not_exist")
-    assert _is_valid_uuid_v4_string(summary.query_id())
+    test_client.command("DROP TABLE IF EXISTS test_query_id_insert")
+    test_client.command(f"CREATE TABLE test_query_id_insert (id UInt32) ENGINE {test_table_engine} ORDER BY id")
+    summary = test_client.insert("test_query_id_insert", [[1], [2], [3]], column_names=["id"])
+    assert _is_valid_uuid_v4(summary.query_id())
+    test_client.command("DROP TABLE test_query_id_insert")
+
+
+def test_query_id_manual_override(test_client: Client):
+    """Test that manually specified query_id is respected and not overwritten"""
+    manual_query_id = "test_manual_query_id_override"
+    result = test_client.query("SELECT 1", settings={"query_id": manual_query_id})
+    assert result.query_id == manual_query_id
 
-    test_client.command("DROP TABLE IF EXISTS test_autogen_query_id")
-    test_client.command(f"CREATE TABLE test_autogen_query_id (id UInt32) ENGINE {test_table_engine} ORDER BY id")
-    summary = test_client.insert("test_autogen_query_id", [[1], [2], [3]], column_names=["id"])
-    assert _is_valid_uuid_v4_string(summary.query_id())
-    test_client.command("DROP TABLE test_autogen_query_id")
 
-    # Create client with autogenerate_query_id disabled
+# pylint: disable=protected-access
+def test_query_id_disabled(test_config: TestConfig):
+    """Test that autogenerate_query_id=False works correctly"""
     client_no_autogen = create_client(
         host=test_config.host,
         port=test_config.port,
         username=test_config.username,
         password=test_config.password,
         autogenerate_query_id=False,
     )
-    result = client_no_autogen.query("SELECT 4")
-    # Even with autogen disabled the server generates a query_id
-    #  so we still expect a query_id in the result, it's just generated by the server...
-    #  not sure how to verify that specifically though.
-    assert _is_valid_uuid_v4_string(result.query_id)
 
-    # We can verify that the setting on the client is False though
     assert client_no_autogen._autogenerate_query_id is False
+
+    # Even with autogen disabled, server generates a query_id
+    result = client_no_autogen.query("SELECT 1")
+    assert _is_valid_uuid_v4(result.query_id)
+
+    client_no_autogen.close()
+
+
+def test_query_id_in_query_logs(test_client: Client, test_config: TestConfig):
+    """Test that query_id appears in ClickHouse's system.query_log for observability"""
+    if test_config.cloud:
+        pytest.skip("Skipping query_log test in cloud environment")
+
+    def check_in_logs(test_query_id):
+        max_retries = 30
+        for _ in range(max_retries):
+            log_result = test_client.query(
+                f"SELECT query_id FROM system.query_log WHERE query_id = '{test_query_id}' AND " "event_time > now() - 30 LIMIT 1"
+            )
+
+            if len(log_result.result_set) > 0:
+                assert log_result.result_set[0][0] == test_query_id
+                return
+
+            sleep(0.1)
+
+        # If we get here, query_id never appeared in logs
+        pytest.fail(f"query_id '{test_query_id}' did not appear in system.query_log after {max_retries * 0.1}s")
+
+    # Manual override check
+    test_query_id_manual = "test_query_id_in_logs"
+    test_client.query("SELECT 1 as num", settings={"query_id": test_query_id_manual})
+    check_in_logs(test_query_id_manual)
+
+    # Autogen check
+    result = test_client.query("SELECT 2 as num")
+    test_query_id_auto = result.query_id
+    check_in_logs(test_query_id_auto)