feat(flow): add row inserts to frontend client (#9006)

* feat(flow): add row inserts to frontend client

Signed-off-by: fys <fengys1996@gmail.com>

* feat(flow): support hints for frontend row inserts

Signed-off-by: fys <fengys1996@gmail.com>

* fix: handle poisoned frontend handler lock in row inserts

Signed-off-by: fys <fengys1996@gmail.com>

---------

Signed-off-by: fys <fengys1996@gmail.com>
This commit is contained in:
fys
2026-09-03 03:40:13 +00:00
committed by GitHub
parent 9c135ebcb3
commit 534ab31297
+82 -10
View File
@@ -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<u32, Error> {
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",
}