Merge pull request #5412 from zhang2014/fix/clickhouse_dead_loop

fixes(handler): fix clickhouse handler dead loop when error
databendlabs · May 17, 2022 · 353ac9f · 353ac9f
2 parents b97c58f + 83da2d5
commit 353ac9f
Showing 1 changed file with 27 additions and 2 deletions.
diff --git a/query/src/servers/clickhouse/interactive_worker_base.rs b/query/src/servers/clickhouse/interactive_worker_base.rs
@@ -223,9 +223,34 @@ impl InteractiveWorkerBase {
                     Err(e)
                 }
                 Ok(mut data_stream) => {
-                    while let Some(block) = data_stream.next().await {
-                        data_tx.send(BlockItem::Block(block)).await.ok();
+                    'worker: loop {
+                        match data_stream.next().await {
+                            None => {
+                                break 'worker;
+                            }
+                            Some(Ok(data)) => {
+                                if let Err(cause) = data_tx.send(BlockItem::Block(Ok(data))).await {
+                                    tracing::warn!(
+                                        "Cannot send data to channel, cause: {:?}",
+                                        cause
+                                    );
+                                    break 'worker;
+                                }
+                            }
+                            Some(Err(error_code)) => {
+                                if let Err(cause) =
+                                    data_tx.send(BlockItem::Block(Err(error_code))).await
+                                {
+                                    tracing::warn!(
+                                        "Cannot send data to channel, cause: {:?}",
+                                        cause
+                                    );
+                                }
+                                break 'worker;
+                            }
+                        }
                     }
+
                     let _ = interpreter
                         .finish()
                         .await