check response

risingwavelabs · Jan 19, 2025 · 0e058d1 · 0e058d1
1 parent 787a9f8
commit 0e058d1
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 25 deletions.
diff --git a/proto/task_service.proto b/proto/task_service.proto
@@ -50,8 +50,13 @@ message FastInsertRequest {
 }
 
 message FastInsertResponse {
-  // Optional error message for failed task.
-  string error_message = 1;
+  enum Status {
+    UNSPECIFIED = 0;
+    SUCCEEDED = 1;
+    DML_FAILED = 2;
+  }
+  Status status = 1;
+  string error_message = 2;
 }
 
 message CancelTaskRequest {

diff --git a/src/batch/src/executor/fast_insert.rs b/src/batch/src/executor/fast_insert.rs
@@ -209,7 +209,6 @@ mod tests {
 
         let json_value = json!(header_map);
         let jsonb_val = JsonbVal::from(json_value);
-        // Add 4 ListValues to ArrayBuilder
         builder.append(Some(jsonb_val.as_scalar_ref()));
 
         // Use builder to obtain a single (List) column DataChunk

diff --git a/src/batch/src/rpc/service/task_service.rs b/src/batch/src/rpc/service/task_service.rs
@@ -18,8 +18,8 @@ use risingwave_common::util::tracing::TracingContext;
 use risingwave_pb::batch_plan::{FastInsertNode, TaskOutputId};
 use risingwave_pb::task_service::task_service_server::TaskService;
 use risingwave_pb::task_service::{
-    CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, ExecuteRequest, FastInsertRequest,
-    FastInsertResponse, GetDataResponse, TaskInfoResponse,
+    fast_insert_response, CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, ExecuteRequest,
+    FastInsertRequest, FastInsertResponse, GetDataResponse, TaskInfoResponse,
 };
 use risingwave_storage::dispatch_state_store;
 use thiserror_ext::AsReport;
@@ -132,9 +132,19 @@ impl TaskService for BatchServiceImpl {
         let res = self.do_fast_insert(insert_node, req.wait_epoch).await;
         match res {
             Ok(_) => Ok(Response::new(FastInsertResponse {
-                error_message: String::from("success"),
+                status: fast_insert_response::Status::Succeeded.into(),
+                error_message: "".to_string(),
             })),
-            Err(e) => Err(e.into()),
+            Err(e) => match e {
+                BatchError::Dml(e) => Ok(Response::new(FastInsertResponse {
+                    status: fast_insert_response::Status::DmlFailed.into(),
+                    error_message: format!("{}", e.as_report()),
+                })),
+                _ => {
+                    error!(error = %e.as_report(), "failed to fast insert");
+                    Err(e.into())
+                }
+            },
         }
     }
 }

diff --git a/src/frontend/src/scheduler/distributed/query.rs b/src/frontend/src/scheduler/distributed/query.rs
@@ -231,11 +231,6 @@ impl QueryExecution {
         compute_client_pool: ComputeClientPoolRef,
         catalog_reader: CatalogReader,
     ) -> HashMap<StageId, Arc<StageExecution>> {
-        // println!(
-        //     "WKXLOG self.query.stage_graph.stages.len(): {}, self.query.stage_graph.stages: {:?}",
-        //     self.query.stage_graph.stages.len(),
-        //     self.query.stage_graph.stages
-        // );
         let mut stage_executions: HashMap<StageId, Arc<StageExecution>> =
             HashMap::with_capacity(self.query.stage_graph.stages.len());
 

diff --git a/src/frontend/src/scheduler/fast_insert.rs b/src/frontend/src/scheduler/fast_insert.rs
@@ -22,7 +22,7 @@ use risingwave_batch::worker_manager::worker_node_manager::WorkerNodeSelector;
 use risingwave_common::hash::WorkerSlotMapping;
 use risingwave_pb::batch_plan::FastInsertNode;
 use risingwave_pb::common::WorkerNode;
-use risingwave_pb::task_service::FastInsertRequest;
+use risingwave_pb::task_service::{FastInsertRequest, FastInsertResponse};
 
 use crate::catalog::TableId;
 use crate::scheduler::{SchedulerError, SchedulerResult};
@@ -55,7 +55,7 @@ impl FastInsertExecution {
         }
     }
 
-    pub async fn my_execute(self) -> SchedulerResult<()> {
+    pub async fn my_execute(self) -> SchedulerResult<FastInsertResponse> {
         let workers = self.choose_worker(
             &TableId::new(self.fast_insert_node.table_id),
             self.fast_insert_node.session_id,
@@ -67,8 +67,7 @@ impl FastInsertExecution {
             wait_epoch: self.wait_for_persistence,
         };
         let response = client.fast_insert(request).await?;
-        println!("WKXLOG response: {:?}", response);
-        Ok(())
+        Ok(response)
     }
 
     #[inline(always)]

diff --git a/src/frontend/src/scheduler/local.rs b/src/frontend/src/scheduler/local.rs
@@ -111,8 +111,6 @@ impl LocalQueryExecution {
         let plan_fragment = self.create_plan_fragment()?;
         let plan_node = plan_fragment.root.unwrap();
 
-        // println!("WKXLOG plan_node: {:?}", plan_node);
-
         let executor = ExecutorBuilder::new(
             &plan_node,
             &task_id,
@@ -123,8 +121,6 @@ impl LocalQueryExecution {
 
         let executor = executor.build().await?;
 
-        // println!("WKXLOG executor: {:?}", executor);
-
         #[for_await]
         for chunk in executor.execute() {
             yield chunk?;

diff --git a/src/frontend/src/webhook/mod.rs b/src/frontend/src/webhook/mod.rs
@@ -64,6 +64,7 @@ pub(super) mod handlers {
     use risingwave_pb::batch_plan::FastInsertNode;
     use risingwave_pb::catalog::WebhookSourceInfo;
     use risingwave_pb::plan_common::DefaultColumns;
+    use risingwave_pb::task_service::fast_insert_response;
     // use risingwave_sqlparser::ast::{Query, SetExpr, Statement, Value, Values};
     use utils::{header_map_to_json, verify_signature};
 
@@ -133,10 +134,13 @@ pub(super) mod handlers {
         }
 
         let mut builder = JsonbArrayBuilder::with_type(1, DataType::Jsonb);
-        // TODO(kexiang): handle errors
-        let json_value = Value::from_text(&body.to_vec()).unwrap();
+        let json_value = Value::from_text(&body.to_vec()).map_err(|e| {
+            err(
+                anyhow!(e).context("Failed to parse body"),
+                StatusCode::UNPROCESSABLE_ENTITY,
+            )
+        })?;
         let jsonb_val = JsonbVal::from(json_value);
-        // Add 4 ListValues to ArrayBuilder
         builder.append(Some(jsonb_val.as_scalar_ref()));
 
         // Use builder to obtain a single (List) column DataChunk
@@ -150,8 +154,14 @@ pub(super) mod handlers {
             session.clone(),
         );
         let res = execution.my_execute().await.unwrap();
-
-        Ok(())
+        if res.status == fast_insert_response::Status::Succeeded as i32 {
+            Ok(())
+        } else {
+            Err(err(
+                anyhow!("Failed to fast insert: {}", res.error_message),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            ))
+        }
     }
 
     async fn acquire_table_info(