fix: use a derived query ID for insert schema lookups

src/lib.rs

@@ -718,6 +718,24 @@ impl Client {
         self
     }
 
+    /// Builds the internal `DESCRIBE TABLE` query used to fetch the schema before an insert.
+    fn insert_metadata_query(&self, raw_table_name: &str) -> query::Query {
+        let mut query = self
+            .query(&_priv::row_insert_metadata_query(raw_table_name))
+            .with_setting("describe_include_subcolumns", "0");
+
+        // Link the schema lookup to the INSERT without reusing its query id.
+        // An empty id still asks the server to generate an id for each query.
+        if let Some(query_id) = self
+            .get_setting(settings::QUERY_ID)
+            .filter(|id| !id.is_empty())
+        {
+            query = query.with_setting(settings::QUERY_ID, format!("{query_id}-describe"));
+        }
+
+        query
+    }
+
     async fn get_insert_metadata(&self, raw_table_name: &str) -> Result<Arc<InsertMetadata>> {
         #[derive(::serde::Deserialize, clickhouse_macros::Row)]
         #[clickhouse(crate = "self")]
@@ -745,8 +763,7 @@ impl Client {
         let mut write_lock = self.insert_metadata_cache.0.write().await;
 
         let mut columns_cursor = self
-            .query(&_priv::row_insert_metadata_query(raw_table_name))
-            .with_setting("describe_include_subcolumns", "0")
+            .insert_metadata_query(raw_table_name)
             .fetch::<DescribeColumn>()?;
 
         let mut columns = Vec::new();

src/query.rs

@@ -433,4 +433,44 @@ mod tests {
             "unexpected error: {err_str:?}"
         );
     }
+
+    #[test]
+    fn insert_metadata_query_derives_query_id() {
+        let client = Client::default()
+            .with_url("http://localhost:8123")
+            .with_setting(settings::QUERY_ID, "my-insert")
+            .with_setting(settings::SESSION_ID, "my-session")
+            .with_setting("max_block_size", "1000");
+
+        let query = client.insert_metadata_query("foo");
+
+        assert_eq!(
+            query.client.get_setting(settings::QUERY_ID),
+            Some("my-insert-describe")
+        );
+        assert_eq!(
+            query.client.get_setting(settings::SESSION_ID),
+            Some("my-session")
+        );
+        assert_eq!(query.client.get_setting("max_block_size"), Some("1000"));
+        assert_eq!(
+            query.client.get_setting("describe_include_subcolumns"),
+            Some("0")
+        );
+
+        // The client itself (and so the `INSERT`) keeps the `query_id`.
+        assert_eq!(client.get_setting(settings::QUERY_ID), Some("my-insert"));
+    }
+
+    #[test]
+    fn insert_metadata_query_preserves_generated_query_ids() {
+        for query_id in [None, Some("")] {
+            let mut client = Client::default();
+            if let Some(query_id) = query_id {
+                client = client.with_setting(settings::QUERY_ID, query_id);
+            }
+            let query = client.insert_metadata_query("foo");
+            assert_eq!(query.client.get_setting(settings::QUERY_ID), query_id);
+        }
+    }
 }

tests/it/insert.rs

@@ -73,6 +73,59 @@ async fn keeps_client_settings() {
     assert_eq!(rows, vec!(row))
 }
 
+// https://github.com/ClickHouse/clickhouse-rs/issues/477
+#[tokio::test]
+async fn describe_query_id_is_linked_to_insert() {
+    let table_name = "insert_describe_query_id_is_linked_to_insert";
+    let query_id = uuid::Uuid::new_v4().to_string();
+
+    let client = prepare_database!();
+    create_simple_table(&client, table_name).await;
+
+    let row = SimpleRow::new(42, "foo");
+
+    // Triggers the internal `DESCRIBE TABLE` with `query_id` set on the client.
+    let mut insert = client
+        .clone()
+        .with_setting("query_id", &query_id)
+        .insert::<SimpleRow>(table_name)
+        .await
+        .unwrap();
+
+    insert.write(&row).await.unwrap();
+    insert.end().await.unwrap();
+
+    flush_query_log(&client).await;
+
+    let describe_query_id = format!("{query_id}-describe");
+    let queries = client
+        .query(
+            "
+            SELECT query_id, query_kind
+            FROM system.query_log
+            WHERE query_id IN (?, ?)
+            AND type = 'QueryFinish'
+            ORDER BY event_time_microseconds
+            ",
+        )
+        .bind(&query_id)
+        .bind(&describe_query_id)
+        .fetch_all::<(String, String)>()
+        .await
+        .unwrap();
+
+    assert_eq!(
+        queries,
+        vec![
+            (describe_query_id, "Describe".into()),
+            (query_id, "Insert".into()),
+        ]
+    );
+
+    let rows = fetch_rows::<SimpleRow>(&client, table_name).await;
+    assert_eq!(rows, vec!(row))
+}
+
 #[tokio::test]
 async fn overrides_client_settings() {
     let table_name = "insert_overrides_client_settings";

tests/it/inserter.rs

@@ -249,6 +249,55 @@ async fn keeps_client_settings() {
     assert_eq!(rows, vec!(row))
 }
 
+/// Similar to [`crate::insert::describe_query_id_is_linked_to_insert`].
+#[tokio::test]
+async fn describe_query_id_is_linked_to_insert() {
+    let table_name = "inserter_describe_query_id_is_linked_to_insert";
+    let query_id = uuid::Uuid::new_v4().to_string();
+
+    let client = prepare_database!();
+    create_simple_table(&client, table_name).await;
+
+    let row = SimpleRow::new(42, "foo");
+
+    let mut inserter = client
+        .inserter::<SimpleRow>(table_name)
+        .with_setting("query_id", &query_id);
+
+    inserter.write(&row).await.unwrap();
+    inserter.end().await.unwrap();
+
+    flush_query_log(&client).await;
+
+    let describe_query_id = format!("{query_id}-describe");
+    let queries = client
+        .query(
+            "
+            SELECT query_id, query_kind
+            FROM system.query_log
+            WHERE query_id IN (?, ?)
+            AND type = 'QueryFinish'
+            ORDER BY event_time_microseconds
+            ",
+        )
+        .bind(&query_id)
+        .bind(&describe_query_id)
+        .fetch_all::<(String, String)>()
+        .await
+        .unwrap();
+
+    assert_eq!(
+        queries,
+        vec![
+            (describe_query_id, "Describe".into()),
+            (query_id, "Insert".into()),
+        ]
+    );
+
+    let rows = fetch_rows::<SimpleRow>(&client, table_name).await;
+    assert_eq!(rows, vec!(row))
+}
+
 /// Similar to [`crate::insert::overrides_client_settings`] with minor differences.
 #[tokio::test]
 async fn overrides_client_settings() {