From 2237e5a8f5e33d5823d6f6e43ff2d207c00a2606 Mon Sep 17 00:00:00 2001
From: BowenXiao1999 <931759898@qq.com>
Date: Tue, 31 Jan 2023 12:08:40 +0800
Subject: [PATCH] fix: do not panic when there is None in dml ret stream

---
 src/frontend/src/handler/query.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/frontend/src/handler/query.rs b/src/frontend/src/handler/query.rs
index 832681ea6bf6b..12f3ffa89ed81 100644
--- a/src/frontend/src/handler/query.rs
+++ b/src/frontend/src/handler/query.rs
@@ -171,11 +171,17 @@ pub async fn handle_query(
         | StatementType::UPDATE_RETURNING => None,
 
         StatementType::INSERT | StatementType::DELETE | StatementType::UPDATE => {
-            let first_row_set = row_stream
-                .next()
-                .await
-                .expect("compute node should return affected rows in output")
-                .map_err(|err| RwError::from(ErrorCode::InternalError(format!("{}", err))))?;
+            let first_row_set = row_stream.next().await;
+            let first_row_set = match first_row_set {
+                None => {
+                    return Err(RwError::from(ErrorCode::InternalError(
+                        "no affected rows in output".to_string(),
+                    )))
+                }
+                Some(row) => {
+                    row.map_err(|err| RwError::from(ErrorCode::InternalError(format!("{}", err))))?
+                }
+            };
             let affected_rows_str = first_row_set[0].values()[0]
                 .as_ref()
                 .expect("compute node should return affected rows in output");