From 534ab31297100645fb57fd718e4a0170345c542e Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:40:13 +0000 Subject: [PATCH] feat(flow): add row inserts to frontend client (#9006) * feat(flow): add row inserts to frontend client Signed-off-by: fys * feat(flow): support hints for frontend row inserts Signed-off-by: fys * fix: handle poisoned frontend handler lock in row inserts Signed-off-by: fys --------- Signed-off-by: fys --- src/flow/src/batching_mode/frontend_client.rs | 92 +++++++++++++++++-- 1 file changed, 82 insertions(+), 10 deletions(-) diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 48314dfa21..10ce3b6820 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -19,7 +19,7 @@ use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; use api::v1::query_request::Query; -use api::v1::{CreateTableExpr, QueryRequest}; +use api::v1::{CreateTableExpr, QueryRequest, RowInsertRequests}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; use common_grpc::channel_manager::{ChannelConfig, ChannelManager, load_client_tls_config}; @@ -364,7 +364,12 @@ impl FrontendClient { database_client .handler .lock() - .unwrap() + .map_err(|e| { + UnexpectedSnafu { + reason: format!("Failed to lock database client: {e}"), + } + .build() + })? .as_ref() .context(UnexpectedSnafu { reason: "Standalone's frontend instance is not set", @@ -387,6 +392,75 @@ impl FrontendClient { } } + /// Execute row inserts on the frontend. + pub async fn row_inserts( + &self, + catalog: &str, + schema: &str, + requests: RowInsertRequests, + hints: &[(&str, &str)], + ) -> Result { + match self { + FrontendClient::Distributed { .. } => { + let db = self.get_random_active_frontend(catalog, schema).await?; + db.database + .row_inserts_with_hints(requests, hints) + .await + .with_context(|_| InvalidRequestSnafu { + context: format!("Failed to handle row inserts at {:?}", db.peer), + }) + } + FrontendClient::Standalone { + database_client, .. + } => { + let extensions = HashMap::from_iter( + hints + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())), + ); + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .extensions(extensions) + .build(); + let ctx = Arc::new(ctx); + { + let database_client = { + database_client + .handler + .lock() + .unwrap() + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + })? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + })? + }; + let resp: common_query::Output = database_client + .do_query(Request::RowInserts(requests), ctx) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + match resp.data { + OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| { + UnexpectedSnafu { + reason: format!("Failed to convert rows to u32: {}", rows), + } + .build() + })?), + _ => UnexpectedSnafu { + reason: "Unexpected output data", + } + .fail(), + } + } + } + } + } + /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( @@ -539,14 +613,12 @@ impl FrontendClient { .map_err(BoxedError::new) .context(ExternalSnafu)?; match resp.data { - common_query::OutputData::AffectedRows(rows) => { - Ok(rows.try_into().map_err(|_| { - UnexpectedSnafu { - reason: format!("Failed to convert rows to u32: {}", rows), - } - .build() - })?) - } + OutputData::AffectedRows(rows) => Ok(rows.try_into().map_err(|_| { + UnexpectedSnafu { + reason: format!("Failed to convert rows to u32: {}", rows), + } + .build() + })?), _ => UnexpectedSnafu { reason: "Unexpected output data", }