[Introspection] Support for partial results for internal queries

When distributed fan-out queries (e.g. `select * from loglet_workers`) encounter
unreachable nodes, return results from the available nodes instead of failing the
entire query. Per-node errors are surfaced as structured warnings alongside the
partial result set.

Implementation:
- ErrorCatchingStream wraps each per-node stream in NodeFanOutExecutionPlan,
  catching errors into a shared NodeWarnings (Arc<parking_lot::Mutex<Vec>>)
  and terminating the individual stream gracefully.
- QueryContext::execute() returns a QueryResult that bundles the record batch
  stream with collected warning handles, walking the physical plan tree to
  extract them from any NodeFanOutExecutionPlan nodes.
- QueryWarningStream in the gRPC handler buffers the last response and attaches
  accumulated warnings to it before yielding, piggybacking on existing data
  rather than sending a trailing empty message.
- Proto: added QueryWarning message and repeated warnings field to QueryResponse.
- CLI: restatectl collects warnings and displays them on stderr after results.

Author: AhmedSoliman

crates/admin/src/cluster_controller/grpc_svc_handler.rs

@@ -8,12 +8,15 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
 use bytes::{Bytes, BytesMut};
 use bytestring::ByteString;
 use datafusion::arrow::ipc::writer::StreamWriter;
 use datafusion::error::DataFusionError;
-use futures::StreamExt;
 use futures::stream::BoxStream;
+use futures::{Stream, StreamExt};
 use tonic::codec::CompressionEncoding;
 use tonic::{Request, Response, Status, async_trait};
 use tracing::info;
@@ -25,14 +28,15 @@ use restate_core::protobuf::cluster_ctrl_svc::{
     CreatePartitionSnapshotResponse, DescribeLogRequest, DescribeLogResponse, FindTailRequest,
     FindTailResponse, GetClusterConfigurationRequest, GetClusterConfigurationResponse,
     ListLogsRequest, ListLogsResponse, MigrateMetadataRequest, MigrateMetadataResponse,
-    QueryRequest, QueryResponse, SealAndExtendChainRequest, SealAndExtendChainResponse,
-    SealChainRequest, SealChainResponse, SealedSegment, SetClusterConfigurationRequest,
-    SetClusterConfigurationResponse, TailState, TrimLogRequest,
+    QueryRequest, QueryResponse, QueryWarning, SealAndExtendChainRequest,
+    SealAndExtendChainResponse, SealChainRequest, SealChainResponse, SealedSegment,
+    SetClusterConfigurationRequest, SetClusterConfigurationResponse, TailState, TrimLogRequest,
     cluster_ctrl_svc_server::{ClusterCtrlSvc, ClusterCtrlSvcServer},
 };
 use restate_core::{Metadata, MetadataWriter};
 use restate_metadata_store::WriteError;
 use restate_storage_query_datafusion::context::QueryContext;
+use restate_storage_query_datafusion::node_fan_out::NodeWarnings;
 use restate_types::config::{MetadataClientKind, MetadataClientOptions, NetworkingOptions};
 use restate_types::identifiers::PartitionId;
 use restate_types::logs::metadata::{Logs, SegmentIndex};
@@ -419,21 +423,37 @@ impl ClusterCtrlSvc for ClusterCtrlSvcHandler {
         request: Request<QueryRequest>,
     ) -> std::result::Result<Response<Self::QueryStream>, tonic::Status> {
         let request = request.into_inner();
-        let stream = self
+        let query_result = self
             .query_context
             .execute(&request.query)
             .await
             .map_err(datafusion_error_to_status)?;
 
-        Ok(Response::new(
-            WriteRecordBatchStream::<StreamWriter<Vec<u8>>>::new(stream, request.query)
-                .map_err(datafusion_error_to_status)?
-                .map(|item| {
-                    item.map(|encoded| QueryResponse { encoded })
-                        .map_err(datafusion_error_to_status)
-                })
-                .boxed(),
-        ))
+        let node_warnings = query_result.node_warnings;
+
+        let data_stream = WriteRecordBatchStream::<StreamWriter<Vec<u8>>>::new(
+            query_result.stream,
+            request.query,
+        )
+        .map_err(datafusion_error_to_status)?
+        .map(|item| {
+            item.map(|encoded| QueryResponse {
+                encoded,
+                ..Default::default()
+            })
+            .map_err(datafusion_error_to_status)
+        });
+
+        // Wrap the data stream to attach per-node warnings to the final
+        // response message, avoiding an extra trailing empty-data message.
+        let stream = QueryWarningStream {
+            inner: data_stream.boxed(),
+            node_warnings,
+            last_response: None,
+            done: false,
+        };
+
+        Ok(Response::new(stream.boxed()))
     }
 
     /// Migrate metadata from the current metadata store to a target store
@@ -539,6 +559,72 @@ fn serialize_value<T: StorageEncode>(value: &T) -> Bytes {
     buf.freeze()
 }
 
+/// Stream wrapper that buffers the last response from the inner stream, and
+/// when the inner stream ends, attaches any accumulated per-node warnings
+/// to that final response before yielding it.
+struct QueryWarningStream {
+    inner: BoxStream<'static, Result<QueryResponse, Status>>,
+    node_warnings: Vec<NodeWarnings>,
+    last_response: Option<Result<QueryResponse, Status>>,
+    done: bool,
+}
+
+impl Stream for QueryWarningStream {
+    type Item = Result<QueryResponse, Status>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        if self.done {
+            return Poll::Ready(None);
+        }
+
+        loop {
+            match self.inner.poll_next_unpin(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => {
+                    self.done = true;
+                    // Inner stream ended. Yield the buffered last response
+                    // with warnings attached, or a warnings-only response.
+                    let warnings = drain_node_warnings(&self.node_warnings);
+                    return match self.last_response.take() {
+                        Some(Ok(mut resp)) => {
+                            resp.warnings = warnings;
+                            Poll::Ready(Some(Ok(resp)))
+                        }
+                        Some(Err(err)) => Poll::Ready(Some(Err(err))),
+                        None if !warnings.is_empty() => {
+                            // No data at all, but we have warnings
+                            Poll::Ready(Some(Ok(QueryResponse {
+                                encoded: Default::default(),
+                                warnings,
+                            })))
+                        }
+                        None => Poll::Ready(None),
+                    };
+                }
+                Poll::Ready(Some(item)) => {
+                    // Yield the previously buffered response, buffer this one
+                    if let Some(prev) = self.last_response.replace(item) {
+                        return Poll::Ready(Some(prev));
+                    }
+                    // First item — buffer it and poll for the next
+                    continue;
+                }
+            }
+        }
+    }
+}
+
+fn drain_node_warnings(node_warnings: &[NodeWarnings]) -> Vec<QueryWarning> {
+    let mut out = Vec::new();
+    for nw in node_warnings {
+        out.extend(nw.lock().drain(..).map(|w| QueryWarning {
+            node_id: w.node_id,
+            message: w.message,
+        }));
+    }
+    out
+}
+
 fn datafusion_error_to_status(err: DataFusionError) -> Status {
     match err {
         DataFusionError::SQL(..)

crates/admin/src/rest_api/query.rs

@@ -98,18 +98,18 @@ where
         return Err(QueryError::Unavailable);
     };
 
-    let record_batch_stream = query_context.execute(&payload.query).await?;
+    let query_result = query_context.execute(&payload.query).await?;
 
     let (result_stream, content_type) = match headers.get(http::header::ACCEPT) {
         Some(v) if v == HeaderValue::from_static("application/json") => (
-            WriteRecordBatchStream::<JsonWriter>::new(record_batch_stream, payload.query)?
+            WriteRecordBatchStream::<JsonWriter>::new(query_result.stream, payload.query)?
                 .map_ok(Frame::data)
                 .left_stream(),
             "application/json",
         ),
         _ => (
             WriteRecordBatchStream::<StreamWriter<Vec<u8>>>::new(
-                record_batch_stream,
+                query_result.stream,
                 payload.query,
             )?
             .map_ok(Frame::data)

crates/admin/src/storage_accounting.rs

@@ -87,6 +87,7 @@ impl StorageAccountingTask {
             .query_context
             .execute(STORAGE_QUERY)
             .await?
+            .stream
             .collect::<Vec<_>>()
             .await
             .into_iter()

crates/core/protobuf/cluster_ctrl_svc.proto

@@ -184,6 +184,17 @@ message QueryRequest {
 message QueryResponse {
   // arrow encoded record batch
   bytes encoded = 1;
+  // Per-node errors encountered during distributed query execution.
+  // Populated on the final response message when some nodes failed but
+  // partial results are still available.
+  repeated QueryWarning warnings = 2;
+}
+
+message QueryWarning {
+  // Node identifier (e.g. "N1", "N2") where the error originated
+  string node_id = 1;
+  // Human-readable error message
+  string message = 2;
 }
 
 message MigrateMetadataRequest {

crates/storage-query-datafusion/src/context.rs

@@ -22,7 +22,7 @@ use datafusion::execution::SessionStateBuilder;
 use datafusion::execution::TaskContext;
 use datafusion::execution::context::SQLOptions;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
-use datafusion::physical_plan::SendableRecordBatchStream;
+use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream, execute_stream};
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion::sql::TableReference;
 
@@ -40,6 +40,7 @@ use restate_types::partitions::state::PartitionReplicaSetStates;
 use restate_types::schema::deployment::DeploymentResolver;
 use restate_types::schema::service::ServiceMetadataResolver;
 
+use crate::node_fan_out::NodeWarnings;
 use crate::remote_query_scanner_manager::RemoteScannerManager;
 
 const SYS_INVOCATION_VIEW: &str = "CREATE VIEW sys_invocation as SELECT
@@ -424,16 +425,25 @@ impl QueryContext {
         })
     }
 
-    pub async fn execute(
-        &self,
-        sql: &str,
-    ) -> datafusion::common::Result<SendableRecordBatchStream> {
+    pub async fn execute(&self, sql: &str) -> datafusion::common::Result<QueryResult> {
         let state = self.datafusion_context.state();
         let statement = state.sql_to_statement(sql, &datafusion::config::Dialect::PostgreSQL)?;
         let plan = state.statement_to_plan(statement).await?;
         self.sql_options.verify_plan(&plan)?;
         let df = self.datafusion_context.execute_logical_plan(plan).await?;
-        df.execute_stream().await
+
+        let task_ctx = Arc::new(df.task_ctx());
+        let physical_plan = df.create_physical_plan().await?;
+
+        // Collect NodeWarnings handles from any NodeFanOutExecutionPlan nodes
+        // in the plan tree before execution begins.
+        let node_warnings = collect_node_warnings(&physical_plan);
+
+        let stream = execute_stream(physical_plan, task_ctx)?;
+        Ok(QueryResult {
+            stream,
+            node_warnings,
+        })
     }
 
     pub fn task_ctx(&self) -> Arc<TaskContext> {
@@ -447,6 +457,31 @@ impl AsRef<SessionContext> for QueryContext {
     }
 }
 
+/// Result of a SQL query execution, containing the record batch stream
+/// and any per-node warning collectors from fan-out execution plans.
+pub struct QueryResult {
+    pub stream: SendableRecordBatchStream,
+    pub node_warnings: Vec<NodeWarnings>,
+}
+
+/// Walks the physical plan tree and collects [`NodeWarnings`] handles from
+/// any [`NodeFanOutExecutionPlan`] nodes found.
+fn collect_node_warnings(plan: &Arc<dyn ExecutionPlan>) -> Vec<NodeWarnings> {
+    use crate::node_fan_out::NodeFanOutExecutionPlan;
+
+    let mut warnings = Vec::new();
+    let mut stack = vec![Arc::clone(plan)];
+    while let Some(node) = stack.pop() {
+        if let Some(fan_out) = node.as_any().downcast_ref::<NodeFanOutExecutionPlan>() {
+            warnings.push(fan_out.node_warnings().clone());
+        }
+        for child in node.children() {
+            stack.push(Arc::clone(child));
+        }
+    }
+    warnings
+}
+
 /// Newtype to add debug implementation which is required for [`SelectPartitions`].
 #[derive(Clone, derive_more::Debug)]
 pub struct SelectPartitionsFromMetadata;

crates/storage-query-datafusion/src/idempotency/tests.rs

@@ -61,7 +61,8 @@ async fn get_idempotency_key() {
         )
         .await
         .unwrap()
-        .collect::<Vec<Result<RecordBatch, _>>>()
+        .stream
+        .collect::<Vec<datafusion::common::Result<RecordBatch>>>()
         .await
         .remove(0)
         .unwrap();

crates/storage-query-datafusion/src/inbox/tests.rs

@@ -45,7 +45,8 @@ async fn get_inbox() {
         .execute("SELECT * FROM sys_inbox ORDER BY sequence_number")
         .await
         .unwrap()
-        .collect::<Vec<Result<RecordBatch, _>>>()
+        .stream
+        .collect::<Vec<datafusion::common::Result<RecordBatch>>>()
         .await
         .remove(0)
         .unwrap();

crates/storage-query-datafusion/src/journal/tests.rs

@@ -86,7 +86,8 @@ async fn get_entries() {
         )
         .await
         .unwrap()
-        .collect::<Vec<Result<RecordBatch, _>>>()
+        .stream
+        .collect::<Vec<datafusion::common::Result<RecordBatch>>>()
         .await
         .remove(0)
         .unwrap();
@@ -179,7 +180,8 @@ async fn select_count_star() {
         .execute("SELECT COUNT(*) AS count FROM sys_journal")
         .await
         .unwrap()
-        .collect::<Vec<Result<RecordBatch, _>>>()
+        .stream
+        .collect::<Vec<datafusion::common::Result<RecordBatch>>>()
         .await
         .remove(0)
         .unwrap();

crates/storage-query-datafusion/src/lib.rs

@@ -23,7 +23,7 @@ mod keyed_service_status;
 mod log;
 pub mod loglet_worker;
 mod node;
-mod node_fan_out;
+pub mod node_fan_out;
 mod partition;
 mod partition_replica_set;
 mod partition_state;

crates/storage-query-datafusion/src/mocks.rs

@@ -16,7 +16,7 @@ use async_trait::async_trait;
 use datafusion::arrow::array::ArrayRef;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::common::DataFusionError;
-use datafusion::execution::SendableRecordBatchStream;
+
 use googletest::matcher::{Matcher, MatcherResult};
 use serde_json::Value;
 
@@ -202,7 +202,7 @@ impl MockQueryEngine {
     pub async fn execute(
         &self,
         sql: impl AsRef<str> + Send,
-    ) -> datafusion::common::Result<SendableRecordBatchStream> {
+    ) -> datafusion::common::Result<crate::context::QueryResult> {
         self.2.execute(sql.as_ref()).await
     }
 }