diff --git a/src/lib.rs b/src/lib.rs index 718782b5..1f59933e 100644 --- a/src/lib.rs +++ b/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> { #[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::()?; let mut columns = Vec::new(); diff --git a/src/query.rs b/src/query.rs index b6ef2c31..800dd21f 100644 --- a/src/query.rs +++ b/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); + } + } } diff --git a/tests/it/insert.rs b/tests/it/insert.rs index b72e137e..117f3ee5 100644 --- a/tests/it/insert.rs +++ b/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::(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::(&client, table_name).await; + assert_eq!(rows, vec!(row)) +} + #[tokio::test] async fn overrides_client_settings() { let table_name = "insert_overrides_client_settings"; diff --git a/tests/it/inserter.rs b/tests/it/inserter.rs index 1830a8e4..459ce8af 100644 --- a/tests/it/inserter.rs +++ b/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::(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::(&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() {