compute: add configurable peek row iteration limit

misc/python/materialize/mzcompose/__init__.py

@@ -119,6 +119,11 @@ def get_minimal_system_parameters(
         # End of list (ordered by name)
     }
 
+    if version >= MzVersion.parse_mz("v26.38.0-dev"):
+        # Exercise the row-limit check without constraining normal test queries.
+        config["compute_peek_row_iteration_limit"] = "1000000000"
+        config["enable_compute_peek_row_iteration_limit"] = "true"
+
     if version < MzVersion.parse_mz("v0.163.0-dev"):
         config["enable_compute_active_dataflow_cancelation"] = "true"
 

misc/python/materialize/parallel_workload/action.py

@@ -2982,6 +2982,10 @@ def __init__(
         self.flags_with_values["enable_compute_peek_response_stash"] = (
             BOOLEAN_FLAG_VALUES
         )
+        self.flags_with_values["enable_compute_peek_row_iteration_limit"] = (
+            BOOLEAN_FLAG_VALUES
+        )
+        self.flags_with_values["compute_peek_row_iteration_limit"] = ["1000000000"]
         self.flags_with_values["compute_peek_response_stash_threshold_bytes"] = [
             "0",  # "force enabled"
             "1048576",  # 1 MiB, an in-between value

src/adapter/src/active_compute_sink.rs

@@ -178,7 +178,9 @@ impl ActiveSubscribe {
                 mz_ore::iter::consolidate_update_iter(merged)
             }
             Err(s) => {
-                self.send(PeekResponseUnary::Error(s));
+                self.send(PeekResponseUnary::Error(AdapterError::Unstructured(
+                    anyhow::Error::msg(s),
+                )));
                 return true;
             }
         };

src/adapter/src/coord/catalog_implications.rs

@@ -39,7 +39,7 @@ use mz_catalog::memory::objects::{
 };
 use mz_cloud_resources::VpcEndpointConfig;
 use mz_compute_client::logging::LogVariant;
-use mz_compute_client::protocol::response::PeekResponse;
+use mz_compute_client::protocol::response::{PeekError, PeekResponse};
 use mz_controller::clusters::{ClusterRole, ReplicaConfig};
 use mz_controller_types::{ClusterId, ReplicaId};
 use mz_ore::collections::CollectionExt;
@@ -992,7 +992,9 @@ impl Coordinator {
             if !peeks_to_drop.is_empty() {
                 for (dep, uuid) in peeks_to_drop {
                     if let Some(pending_peek) = self.remove_pending_peek(&uuid) {
-                        let cancel_reason = PeekResponse::Error(dep.query_terminated_error());
+                        let cancel_reason = PeekResponse::Error(PeekError::unstructured(
+                            dep.query_terminated_error(),
+                        ));
                         self.controller
                             .compute
                             .cancel_peek(pending_peek.cluster_id, uuid, cancel_reason)

src/adapter/src/coord/peek.rs

@@ -85,13 +85,9 @@ pub(crate) struct PendingPeek {
 #[derive(Debug)]
 pub enum PeekResponseUnary {
     Rows(Box<dyn RowIterator + Send + Sync>),
-    Error(String),
+    Error(AdapterError),
     Canceled,
     /// A dependency was dropped during execution.
-    ///
-    /// N.B. This is a bit of a workaround for the fact that our Error variant
-    /// is unstructured and right now we specifically care about this error and
-    /// need to render differently based on context.
     DependencyDropped(DroppedDependency),
 }
 
@@ -1049,7 +1045,7 @@ impl crate::coord::Coordinator {
             let rows = match result {
                 Ok(rows) => rows,
                 Err(e) => {
-                    yield PeekResponseUnary::Error(e.to_string());
+                    yield PeekResponseUnary::Error(AdapterError::Unstructured(anyhow::anyhow!(e)));
                     return;
                 }
             };
@@ -1064,7 +1060,11 @@ impl crate::coord::Coordinator {
                         &duration_histogram,
                     ) {
                         Ok((rows, _size_bytes)) => yield PeekResponseUnary::Rows(Box::new(rows)),
-                        Err(e) => yield PeekResponseUnary::Error(e),
+                        Err(e) => {
+                            yield PeekResponseUnary::Error(AdapterError::Unstructured(
+                                anyhow::Error::msg(e),
+                            ))
+                        }
                     }
                 }
                 PeekResponse::Stashed(response) => {
@@ -1221,7 +1221,11 @@ impl crate::coord::Coordinator {
 
                         match result_rows {
                             Ok(result_rows) => yield PeekResponseUnary::Rows(Box::new(result_rows)),
-                            Err(e) => yield PeekResponseUnary::Error(e),
+                            Err(e) => {
+                                yield PeekResponseUnary::Error(AdapterError::Unstructured(
+                                    anyhow::Error::msg(e),
+                                ))
+                            }
                         }
                     }
 
@@ -1236,7 +1240,7 @@ impl crate::coord::Coordinator {
                     yield PeekResponseUnary::Canceled;
                 }
                 PeekResponse::Error(e) => {
-                    yield PeekResponseUnary::Error(e);
+                    yield PeekResponseUnary::Error(e.into());
                 }
             }
         })

src/adapter/src/coord/sequencer/inner.rs

@@ -2987,9 +2987,7 @@ impl Coordinator {
                                     };
                                 }
                                 PeekResponseUnary::Canceled => break Err(AdapterError::Canceled),
-                                PeekResponseUnary::Error(e) => {
-                                    break Err(AdapterError::Unstructured(anyhow!(e)));
-                                }
+                                PeekResponseUnary::Error(e) => break Err(e),
                                 PeekResponseUnary::DependencyDropped(dep) => {
                                     break Err(dep.to_concurrent_dependency_drop());
                                 }

src/adapter/src/error.rs

@@ -36,7 +36,7 @@ use mz_sql::rbac;
 use mz_sql::session::vars::VarError;
 use mz_storage_types::connections::ConnectionValidationError;
 use mz_storage_types::controller::StorageError;
-use mz_storage_types::errors::CollectionMissing;
+use mz_storage_types::errors::{CollectionMissing, DataflowError};
 use smallvec::SmallVec;
 use timely::progress::Antichain;
 use tokio::sync::oneshot;
@@ -69,6 +69,8 @@ pub enum AdapterError {
     DuplicateCursor(String),
     /// An error while evaluating an expression.
     Eval(EvalError),
+    /// An error produced while executing a dataflow.
+    Dataflow(Box<DataflowError>),
     /// An error occurred while planning the statement.
     Explain(ExplainError),
     /// The ID allocator exhausted all valid IDs.
@@ -173,6 +175,11 @@ pub enum AdapterError {
     },
     /// Result size of a query is too large.
     ResultSize(String),
+    /// A query exceeded the configured compute peek row iteration limit.
+    PeekRowIterationLimitExceeded {
+        /// The configured per-worker limit.
+        limit: usize,
+    },
     /// The specified feature is not permitted in safe mode.
     SafeModeViolation(String),
     /// The current transaction had the wrong set of write locks.
@@ -495,6 +502,15 @@ fn eval_error_code(err: &EvalError) -> SqlState {
     }
 }
 
+fn dataflow_error_code(error: &DataflowError) -> SqlState {
+    match error {
+        DataflowError::EvalError(error) => eval_error_code(error),
+        DataflowError::DecodeError(_)
+        | DataflowError::SourceError(_)
+        | DataflowError::EnvelopeError(_) => SqlState::INTERNAL_ERROR,
+    }
+}
+
 impl AdapterError {
     pub fn into_response(self, severity: Severity) -> ErrorResponse {
         ErrorResponse {
@@ -522,6 +538,10 @@ impl AdapterError {
             }
             AdapterError::Catalog(c) => c.detail(),
             AdapterError::Eval(e) => e.detail(),
+            AdapterError::Dataflow(e) => match &**e {
+                DataflowError::EvalError(e) => e.detail(),
+                _ => None,
+            },
             AdapterError::RelationOutsideTimeDomain { relations, names } => Some(format!(
                 "The following relations in the query are outside the transaction's time domain:\n{}\n{}",
                 relations
@@ -562,6 +582,11 @@ impl AdapterError {
                      objects. Reduce the number of dependencies, or raise the \
                      read_then_write_max_dependencies system parameter."
             )),
+            AdapterError::PeekRowIterationLimitExceeded { limit } => Some(format!(
+                "The query attempted to examine more than {limit} rows on a single compute \
+                 worker. This limit prevents long-running SELECT queries from delaying other \
+                 work on the cluster."
+            )),
             AdapterError::SafeModeViolation(_) => Some(
                 "The Materialize server you are connected to is running in \
                  safe mode, which limits the features that are available."
@@ -735,6 +760,10 @@ impl AdapterError {
             ),
             AdapterError::Catalog(c) => c.hint(),
             AdapterError::Eval(e) => e.hint(),
+            AdapterError::Dataflow(e) => match &**e {
+                DataflowError::EvalError(e) => e.hint(),
+                _ => None,
+            },
             AdapterError::AlterClusterUnmanagedWhileReconfiguring => Some(
                 "Cancel the reconfiguration by altering the cluster back to its current \
                 configuration, or wait for it to settle, then convert."
@@ -798,6 +827,13 @@ impl AdapterError {
                  statement_timeout = '120s'`."
                     .into(),
             ),
+            AdapterError::PeekRowIterationLimitExceeded { .. } => Some(
+                "Reduce the number of rows the query must examine, for example by querying an \
+                 indexed, more selective result. Queries with `LIMIT` and no `ORDER BY` can also \
+                 stop early. To permit this query, increase \
+                 `compute_peek_row_iteration_limit`."
+                    .into(),
+            ),
             AdapterError::PlanError(e) => e.hint(),
             AdapterError::UnallowedOnCluster { cluster, .. } => {
                 (cluster != MZ_CATALOG_SERVER_CLUSTER.name).then(||
@@ -864,6 +900,7 @@ impl AdapterError {
             // exhaustively so the catch-all `INTERNAL_ERROR` no longer applies
             // to errors that are really the user's fault. See SQL-326.
             AdapterError::Eval(e) => eval_error_code(e),
+            AdapterError::Dataflow(e) => dataflow_error_code(e),
             AdapterError::Explain(_) => SqlState::INTERNAL_ERROR,
             AdapterError::IdExhaustionError => SqlState::INTERNAL_ERROR,
             AdapterError::Internal(_) => SqlState::INTERNAL_ERROR,
@@ -923,6 +960,7 @@ impl AdapterError {
             AdapterError::RelationOutsideTimeDomain { .. } => SqlState::INVALID_TRANSACTION_STATE,
             AdapterError::ResourceExhaustion { .. } => SqlState::INSUFFICIENT_RESOURCES,
             AdapterError::ResultSize(_) => SqlState::OUT_OF_MEMORY,
+            AdapterError::PeekRowIterationLimitExceeded { .. } => SqlState::PROGRAM_LIMIT_EXCEEDED,
             AdapterError::SafeModeViolation(_) => SqlState::INTERNAL_ERROR,
             AdapterError::SubscribeOnlyTransaction => SqlState::INVALID_TRANSACTION_STATE,
             AdapterError::Optimizer(e) => match e {
@@ -1197,6 +1235,7 @@ impl fmt::Display for AdapterError {
                 write!(f, "cursor {} already exists", name.quoted())
             }
             AdapterError::Eval(e) => e.fmt(f),
+            AdapterError::Dataflow(e) => e.fmt(f),
             AdapterError::Explain(e) => e.fmt(f),
             AdapterError::IdExhaustionError => f.write_str("ID allocator exhausted all valid IDs"),
             AdapterError::Internal(e) => write!(f, "internal error: {}", e),
@@ -1237,6 +1276,12 @@ impl fmt::Display for AdapterError {
                     "selection has too many transitive dependencies to validate (limit {max_rw_dependencies})"
                 )
             }
+            AdapterError::PeekRowIterationLimitExceeded { limit } => {
+                write!(
+                    f,
+                    "query exceeded the configured row iteration limit of {limit} rows"
+                )
+            }
             AdapterError::ReplaceMaterializedViewSealed { name } => {
                 write!(
                     f,
@@ -1554,6 +1599,20 @@ impl From<EvalError> for AdapterError {
     }
 }
 
+impl From<mz_compute_client::protocol::response::PeekError> for AdapterError {
+    fn from(error: mz_compute_client::protocol::response::PeekError) -> Self {
+        use mz_compute_client::protocol::response::PeekError;
+
+        match error {
+            PeekError::Dataflow(error) => AdapterError::Dataflow(error),
+            PeekError::Unstructured(error) => AdapterError::Unstructured(anyhow::Error::msg(error)),
+            PeekError::RowIterationLimitExceeded { limit } => {
+                AdapterError::PeekRowIterationLimitExceeded { limit }
+            }
+        }
+    }
+}
+
 impl From<ExplainError> for AdapterError {
     fn from(e: ExplainError) -> AdapterError {
         match e {
@@ -1688,3 +1747,47 @@ impl From<ConnectionValidationError> for AdapterError {
 }
 
 impl Error for AdapterError {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[mz_ore::test]
+    fn peek_row_iteration_limit_error_is_user_facing() {
+        let response = AdapterError::PeekRowIterationLimitExceeded { limit: 1000 }
+            .into_response(Severity::Error);
+
+        assert_eq!(response.code, SqlState::PROGRAM_LIMIT_EXCEEDED);
+        assert_eq!(
+            response.message,
+            "query exceeded the configured row iteration limit of 1000 rows"
+        );
+        assert_eq!(
+            response.detail.as_deref(),
+            Some(
+                "The query attempted to examine more than 1000 rows on a single compute worker. \
+                 This limit prevents long-running SELECT queries from delaying other work on the \
+                 cluster."
+            )
+        );
+        assert!(
+            response
+                .hint
+                .as_deref()
+                .is_some_and(|hint| hint.contains("compute_peek_row_iteration_limit"))
+        );
+    }
+
+    #[mz_ore::test]
+    fn structured_dataflow_error_preserves_message_and_code() {
+        use mz_compute_client::protocol::response::PeekError;
+
+        let dataflow_error = DataflowError::from(EvalError::DivisionByZero);
+        let expected_message = dataflow_error.to_string();
+        let response =
+            AdapterError::from(PeekError::from(dataflow_error)).into_response(Severity::Error);
+
+        assert_eq!(response.code, SqlState::DIVISION_BY_ZERO);
+        assert_eq!(response.message, expected_message);
+    }
+}

src/compute-client/Cargo.toml

@@ -41,5 +41,8 @@ tokio.workspace = true
 tracing.workspace = true
 uuid = { workspace = true, features = ["serde", "v4"] }
 
+[dev-dependencies]
+bincode.workspace = true
+
 [features]
 default = []

src/compute-client/src/controller.rs

@@ -170,7 +170,7 @@ impl PeekNotification {
                     result_size: u64::cast_from(result_size),
                 }
             }
-            PeekResponse::Error(err) => Self::Error(err.clone()),
+            PeekResponse::Error(err) => Self::Error(err.to_string()),
             PeekResponse::Canceled => Self::Canceled,
         }
     }

src/compute-client/src/controller/instance.rs

@@ -64,8 +64,8 @@ use crate::protocol::command::{
 };
 use crate::protocol::history::ComputeCommandHistory;
 use crate::protocol::response::{
-    ComputeResponse, CopyToResponse, FrontiersResponse, PeekResponse, StatusResponse,
-    SubscribeBatch, SubscribeResponse,
+    ComputeResponse, CopyToResponse, FrontiersResponse, PeekError as ProtocolPeekError,
+    PeekResponse, StatusResponse, SubscribeBatch, SubscribeResponse,
 };
 
 #[derive(Error, Debug)]
@@ -1346,7 +1346,8 @@ impl Instance {
             self.deliver_response(response);
         }
         for uuid in to_drop {
-            let response = PeekResponse::Error(ERROR_TARGET_REPLICA_FAILED.into());
+            let response =
+                PeekResponse::Error(ProtocolPeekError::unstructured(ERROR_TARGET_REPLICA_FAILED));
             self.finish_peek(uuid, response);
         }
 

src/compute-client/src/protocol/command.rs

@@ -233,12 +233,13 @@ pub enum ComputeCommand {
     /// After receiving a `Peek` command, the replica must eventually produce a single
     /// [`PeekResponse`]:
     ///
-    ///    * For peeks that were not cancelled: either [`Rows`] or [`Error`].
-    ///    * For peeks that were cancelled: either [`Rows`], or [`Error`], or [`Canceled`].
+    ///    * For peeks that were not cancelled: [`Rows`], [`Stashed`], or [`Error`].
+    ///    * For peeks that were cancelled: any response above, or [`Canceled`].
     ///
     /// [`PeekResponse`]: super::response::PeekResponse
     /// [`PeekResponse::Error`]: super::response::PeekResponse::Error
     /// [`Rows`]: super::response::PeekResponse::Rows
+    /// [`Stashed`]: super::response::PeekResponse::Stashed
     /// [`Error`]: super::response::PeekResponse::Error
     /// [`Canceled`]: super::response::PeekResponse::Canceled
     Peek(Box<Peek>),