Frontend peek: statement logging

Author: ggevay

src/adapter/src/client.rs

@@ -264,6 +264,7 @@ impl Client {
             transient_id_gen,
             optimizer_metrics,
             persist_client,
+            statement_logging_frontend,
         } = response;
 
         let peek_client = PeekClient::new(
@@ -272,6 +273,7 @@ impl Client {
             transient_id_gen,
             optimizer_metrics,
             persist_client,
+            statement_logging_frontend,
         );
 
         let mut client = SessionClient {
@@ -704,11 +706,19 @@ impl SessionClient {
         // Attempt peek sequencing in the session task.
         // If unsupported, fall back to the Coordinator path.
         // TODO(peek-seq): wire up cancel_future
-        if let Some(resp) = self.try_frontend_peek(&portal_name).await? {
+        let mut outer_ctx_extra = outer_ctx_extra;
+        if let Some(resp) = self
+            .try_frontend_peek(&portal_name, &mut outer_ctx_extra)
+            .await?
+        {
             debug!("frontend peek succeeded");
+            // Frontend peek handled the execution and retired outer_ctx_extra if it existed.
+            // No additional work needed here.
             return Ok((resp, execute_started));
         } else {
-            debug!("frontend peek did not happen");
+            debug!("frontend peek did not happen, falling back to `Command::Execute`");
+            // If we bailed out, outer_ctx_extra is still present (if it was originally).
+            // `Command::Execute` will handle it.
         }
 
         let response = self
@@ -1020,7 +1030,9 @@ impl SessionClient {
                 | Command::StoreTransactionReadHolds { .. }
                 | Command::ExecuteSlowPathPeek { .. }
                 | Command::ExecuteCopyTo { .. }
-                | Command::ExecuteSideEffectingFunc { .. } => {}
+                | Command::ExecuteSideEffectingFunc { .. }
+                | Command::RegisterFrontendPeek { .. }
+                | Command::FrontendStatementLogging(..) => {}
             };
             cmd
         });
@@ -1110,11 +1122,12 @@ impl SessionClient {
     pub(crate) async fn try_frontend_peek(
         &mut self,
         portal_name: &str,
+        outer_ctx_extra: &mut Option<ExecuteContextExtra>,
     ) -> Result<Option<ExecuteResponse>, AdapterError> {
         if self.enable_frontend_peek_sequencing {
             let session = self.session.as_mut().expect("SessionClient invariant");
             self.peek_client
-                .try_frontend_peek_inner(portal_name, session)
+                .try_frontend_peek(portal_name, session, outer_ctx_extra)
                 .await
         } else {
             Ok(None)

src/adapter/src/command.rs

@@ -47,10 +47,13 @@ use crate::coord::ExecuteContextExtra;
 use crate::coord::appends::BuiltinTableAppendNotify;
 use crate::coord::consistency::CoordinatorInconsistencies;
 use crate::coord::peek::{PeekDataflowPlan, PeekResponseUnary};
+use crate::coord::statement_logging::WatchSetCreation;
 use crate::coord::timestamp_selection::TimestampDetermination;
 use crate::error::AdapterError;
 use crate::session::{EndTransactionAction, RowBatchStream, Session};
-use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
+use crate::statement_logging::{
+    StatementEndedExecutionReason, StatementExecutionStrategy, StatementLoggingFrontend,
+};
 use crate::util::Transmittable;
 use crate::webhook::AppendWebhookResponse;
 use crate::{AdapterNotice, AppendWebhookError, ReadHolds};
@@ -210,6 +213,9 @@ pub enum Command {
         conn_id: ConnectionId,
         max_result_size: u64,
         max_query_result_size: Option<u64>,
+        /// If statement logging is enabled, contains all info needed for installing watch sets
+        /// and logging the statement execution.
+        watch_set: Option<WatchSetCreation>,
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     },
 
@@ -219,6 +225,9 @@ pub enum Command {
         target_replica: Option<ReplicaId>,
         source_ids: BTreeSet<GlobalId>,
         conn_id: ConnectionId,
+        /// If statement logging is enabled, contains all info needed for installing watch sets
+        /// and logging the statement execution.
+        watch_set: Option<WatchSetCreation>,
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     },
 
@@ -230,6 +239,24 @@ pub enum Command {
         current_role: RoleId,
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     },
+
+    /// Register a pending peek initiated by frontend sequencing.
+    /// The Coordinator will track this peek and log its final result when it completes.
+    RegisterFrontendPeek {
+        uuid: Uuid,
+        conn_id: ConnectionId,
+        cluster_id: mz_controller_types::ClusterId,
+        depends_on: BTreeSet<GlobalId>,
+        is_fast_path: bool,
+        /// If statement logging is enabled, contains all info needed for installing watch sets
+        /// and logging the statement execution.
+        watch_set: Option<WatchSetCreation>,
+        tx: oneshot::Sender<Result<(), AdapterError>>,
+    },
+
+    /// Statement logging event from frontend peek sequencing.
+    /// No response channel needed - this is fire-and-forget.
+    FrontendStatementLogging(crate::statement_logging::FrontendStatementLoggingEvent),
 }
 
 impl Command {
@@ -257,7 +284,9 @@ impl Command {
             | Command::StoreTransactionReadHolds { .. }
             | Command::ExecuteSlowPathPeek { .. }
             | Command::ExecuteCopyTo { .. }
-            | Command::ExecuteSideEffectingFunc { .. } => None,
+            | Command::ExecuteSideEffectingFunc { .. }
+            | Command::RegisterFrontendPeek { .. }
+            | Command::FrontendStatementLogging(..) => None,
         }
     }
 
@@ -285,7 +314,9 @@ impl Command {
             | Command::StoreTransactionReadHolds { .. }
             | Command::ExecuteSlowPathPeek { .. }
             | Command::ExecuteCopyTo { .. }
-            | Command::ExecuteSideEffectingFunc { .. } => None,
+            | Command::ExecuteSideEffectingFunc { .. }
+            | Command::RegisterFrontendPeek { .. }
+            | Command::FrontendStatementLogging(..) => None,
         }
     }
 }
@@ -318,6 +349,7 @@ pub struct StartupResponse {
     pub transient_id_gen: Arc<TransientIdGen>,
     pub optimizer_metrics: OptimizerMetrics,
     pub persist_client: PersistClient,
+    pub statement_logging_frontend: StatementLoggingFrontend,
 }
 
 /// The response to [`Client::authenticate`](crate::Client::authenticate).

src/adapter/src/coord.rs

@@ -104,7 +104,9 @@ use mz_catalog::memory::objects::{
 };
 use mz_cloud_resources::{CloudResourceController, VpcEndpointConfig, VpcEndpointEvent};
 use mz_compute_client::as_of_selection;
-use mz_compute_client::controller::error::{DataflowCreationError, InstanceMissing};
+use mz_compute_client::controller::error::{
+    CollectionLookupError, DataflowCreationError, InstanceMissing,
+};
 use mz_compute_types::ComputeInstanceId;
 use mz_compute_types::dataflows::DataflowDescription;
 use mz_compute_types::plan::Plan;
@@ -190,7 +192,7 @@ use crate::coord::cluster_scheduling::SchedulingDecision;
 use crate::coord::id_bundle::CollectionIdBundle;
 use crate::coord::introspection::IntrospectionSubscribe;
 use crate::coord::peek::PendingPeek;
-use crate::coord::statement_logging::{StatementLogging, StatementLoggingId};
+use crate::coord::statement_logging::StatementLogging;
 use crate::coord::timeline::{TimelineContext, TimelineState};
 use crate::coord::timestamp_selection::{TimestampContext, TimestampDetermination};
 use crate::coord::validity::PlanValidity;
@@ -203,7 +205,9 @@ use crate::optimize::dataflows::{
 };
 use crate::optimize::{self, Optimize, OptimizerConfig};
 use crate::session::{EndTransactionAction, Session};
-use crate::statement_logging::{StatementEndedExecutionReason, StatementLifecycleEvent};
+use crate::statement_logging::{
+    StatementEndedExecutionReason, StatementLifecycleEvent, StatementLoggingId,
+};
 use crate::util::{ClientTransmitter, ResultExt};
 use crate::webhook::{WebhookAppenderInvalidator, WebhookConcurrencyLimiter};
 use crate::{AdapterNotice, ReadHolds, flags};
@@ -374,6 +378,8 @@ impl Message {
                 Command::ExecuteSlowPathPeek { .. } => "execute-slow-path-peek",
                 Command::ExecuteCopyTo { .. } => "execute-copy-to",
                 Command::ExecuteSideEffectingFunc { .. } => "execute-side-effecting-func",
+                Command::RegisterFrontendPeek { .. } => "register-frontend-peek",
+                Command::FrontendStatementLogging(..) => "frontend-statement-logging",
             },
             Message::ControllerReady {
                 controller: ControllerReadiness::Compute,
@@ -1332,7 +1338,7 @@ impl ExecuteContextExtra {
     /// called from code that knows what to do to finish up logging
     /// based on the inner value.
     #[must_use]
-    fn retire(mut self) -> Option<StatementLoggingId> {
+    pub(crate) fn retire(mut self) -> Option<StatementLoggingId> {
         let Self { statement_uuid } = &mut self;
         statement_uuid.take()
     }
@@ -3763,13 +3769,14 @@ impl Coordinator {
         objects: BTreeSet<GlobalId>,
         t: Timestamp,
         state: WatchSetResponse,
-    ) {
-        let ws_id = self.controller.install_compute_watch_set(objects, t);
+    ) -> Result<(), CollectionLookupError> {
+        let ws_id = self.controller.install_compute_watch_set(objects, t)?;
         self.connection_watch_sets
             .entry(conn_id.clone())
             .or_default()
             .insert(ws_id);
         self.installed_watch_sets.insert(ws_id, (conn_id, state));
+        Ok(())
     }
 
     /// Install a _watch set_ in the controller that is automatically associated with the given
@@ -3781,13 +3788,14 @@ impl Coordinator {
         objects: BTreeSet<GlobalId>,
         t: Timestamp,
         state: WatchSetResponse,
-    ) {
-        let ws_id = self.controller.install_storage_watch_set(objects, t);
+    ) -> Result<(), CollectionLookupError> {
+        let ws_id = self.controller.install_storage_watch_set(objects, t)?;
         self.connection_watch_sets
             .entry(conn_id.clone())
             .or_default()
             .insert(ws_id);
         self.installed_watch_sets.insert(ws_id, (conn_id, state));
+        Ok(())
     }
 
     /// Cancels pending watchsets associated with the provided connection id.

src/adapter/src/coord/catalog_implications.rs

@@ -60,9 +60,8 @@ use crate::coord::Coordinator;
 use crate::coord::catalog_implications::parsed_state_updates::{
     ParsedStateUpdate, ParsedStateUpdateKind,
 };
-use crate::coord::statement_logging::StatementLoggingId;
 use crate::coord::timeline::TimelineState;
-use crate::statement_logging::StatementEndedExecutionReason;
+use crate::statement_logging::{StatementEndedExecutionReason, StatementLoggingId};
 use crate::{AdapterError, CollectionIdBundle, ExecuteContext, ResultExt};
 
 pub mod parsed_state_updates;