Merge pull request #34261 from ggevay/frontend-peek-copy-to

Frontend peek sequencing -- `COPY TO`

Author: ggevay

src/adapter/src/client.rs

@@ -1016,7 +1016,8 @@ impl SessionClient {
                 | Command::GetComputeInstanceClient { .. }
                 | Command::GetOracle { .. }
                 | Command::DetermineRealTimeRecentTimestamp { .. }
-                | Command::ExecuteSlowPathPeek { .. } => {}
+                | Command::ExecuteSlowPathPeek { .. }
+                | Command::ExecuteCopyTo { .. } => {}
             };
             cmd
         });

src/adapter/src/command.rs

@@ -52,7 +52,7 @@ use crate::session::{EndTransactionAction, RowBatchStream, Session};
 use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
 use crate::util::Transmittable;
 use crate::webhook::AppendWebhookResponse;
-use crate::{AdapterNotice, AppendWebhookError};
+use crate::{AdapterNotice, AppendWebhookError, optimize};
 
 #[derive(Debug)]
 pub struct CatalogSnapshot {
@@ -199,6 +199,15 @@ pub enum Command {
         max_query_result_size: Option<u64>,
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     },
+
+    ExecuteCopyTo {
+        global_lir_plan: Box<optimize::copy_to::GlobalLirPlan>,
+        compute_instance: ComputeInstanceId,
+        target_replica: Option<ReplicaId>,
+        source_ids: BTreeSet<GlobalId>,
+        conn_id: ConnectionId,
+        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
+    },
 }
 
 impl Command {
@@ -222,7 +231,8 @@ impl Command {
             | Command::GetComputeInstanceClient { .. }
             | Command::GetOracle { .. }
             | Command::DetermineRealTimeRecentTimestamp { .. }
-            | Command::ExecuteSlowPathPeek { .. } => None,
+            | Command::ExecuteSlowPathPeek { .. }
+            | Command::ExecuteCopyTo { .. } => None,
         }
     }
 
@@ -246,7 +256,8 @@ impl Command {
             | Command::GetComputeInstanceClient { .. }
             | Command::GetOracle { .. }
             | Command::DetermineRealTimeRecentTimestamp { .. }
-            | Command::ExecuteSlowPathPeek { .. } => None,
+            | Command::ExecuteSlowPathPeek { .. }
+            | Command::ExecuteCopyTo { .. } => None,
         }
     }
 }

src/adapter/src/coord.rs

@@ -367,6 +367,7 @@ impl Message {
                     "determine-real-time-recent-timestamp"
                 }
                 Command::ExecuteSlowPathPeek { .. } => "execute-slow-path-peek",
+                Command::ExecuteCopyTo { .. } => "execute-copy-to",
             },
             Message::ControllerReady {
                 controller: ControllerReadiness::Compute,

src/adapter/src/coord/command_handler.rs

@@ -344,6 +344,28 @@ impl Coordinator {
 
                     let _ = tx.send(result);
                 }
+
+                Command::ExecuteCopyTo {
+                    global_lir_plan,
+                    compute_instance,
+                    target_replica,
+                    source_ids,
+                    conn_id,
+                    tx,
+                } => {
+                    // implement_copy_to spawns a background task that sends the response
+                    // through tx when the COPY TO completes (or immediately if setup fails).
+                    // We just call it and let it handle all response sending.
+                    self.implement_copy_to(
+                        *global_lir_plan,
+                        compute_instance,
+                        target_replica,
+                        source_ids,
+                        conn_id,
+                        tx,
+                    )
+                    .await;
+                }
             }
         }
         .instrument(debug_span!("handle_command"))

src/adapter/src/coord/peek.rs

@@ -28,6 +28,7 @@ use mz_compute_client::protocol::command::PeekTarget;
 use mz_compute_client::protocol::response::PeekResponse;
 use mz_compute_types::ComputeInstanceId;
 use mz_compute_types::dataflows::{DataflowDescription, IndexImport};
+use mz_compute_types::sinks::ComputeSinkConnection;
 use mz_controller_types::ClusterId;
 use mz_expr::explain::{HumanizedExplain, HumanizerMode, fmt_text_constant_rows};
 use mz_expr::row::RowCollection;
@@ -37,6 +38,7 @@ use mz_expr::{
 };
 use mz_ore::cast::CastFrom;
 use mz_ore::str::{StrExt, separated};
+use mz_ore::task;
 use mz_ore::tracing::OpenTelemetryContext;
 use mz_persist_client::Schemas;
 use mz_persist_types::codec_impls::UnitSchema;
@@ -49,10 +51,13 @@ use mz_repr::{
 use mz_storage_types::sources::SourceData;
 use serde::{Deserialize, Serialize};
 use timely::progress::{Antichain, Timestamp};
+use tokio::sync::oneshot;
+use tracing::{Instrument, Span};
 use uuid::Uuid;
 
+use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyTo};
 use crate::coord::timestamp_selection::TimestampDetermination;
-use crate::optimize::OptimizerError;
+use crate::optimize::{self, OptimizerError};
 use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
 use crate::util::ResultExt;
 use crate::{AdapterError, ExecuteContextExtra, ExecuteResponse};
@@ -1238,6 +1243,117 @@ impl crate::coord::Coordinator {
         .await
     }
 
+    /// Implements a COPY TO command by validating S3 connection, shipping the dataflow,
+    /// and spawning a background task to wait for completion.
+    /// This is called from the command handler for ExecuteCopyTo.
+    ///
+    /// This method inlines the logic from peek_copy_to_preflight and peek_copy_to_dataflow
+    /// to avoid the complexity of the staging mechanism for the frontend peek sequencing path.
+    ///
+    /// This method does NOT block waiting for completion. Instead, it spawns a background task that
+    /// will send the response through the provided tx channel when the COPY TO completes.
+    /// All errors (setup or execution) are sent through tx.
+    pub(crate) async fn implement_copy_to(
+        &mut self,
+        global_lir_plan: optimize::copy_to::GlobalLirPlan,
+        compute_instance: ComputeInstanceId,
+        target_replica: Option<ReplicaId>,
+        source_ids: BTreeSet<GlobalId>,
+        conn_id: ConnectionId,
+        tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
+    ) {
+        // Helper to send error and return early
+        let send_err = |tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
+                        e: AdapterError| {
+            let _ = tx.send(Err(e));
+        };
+
+        let sink_id = global_lir_plan.sink_id();
+
+        // # Inlined from peek_copy_to_preflight
+
+        let connection_context = self.connection_context().clone();
+        let sinks = &global_lir_plan.df_desc().sink_exports;
+
+        if sinks.len() != 1 {
+            send_err(
+                tx,
+                AdapterError::Internal("expected exactly one copy to s3 sink".into()),
+            );
+            return;
+        }
+        let (sink_id_check, sink_desc) = sinks
+            .first_key_value()
+            .expect("known to be exactly one copy to s3 sink");
+        assert_eq!(sink_id, *sink_id_check);
+
+        match &sink_desc.connection {
+            ComputeSinkConnection::CopyToS3Oneshot(conn) => {
+                if let Err(e) = mz_storage_types::sinks::s3_oneshot_sink::preflight(
+                    connection_context,
+                    &conn.aws_connection,
+                    &conn.upload_info,
+                    conn.connection_id,
+                    sink_id,
+                )
+                .await
+                {
+                    send_err(tx, e.into());
+                    return;
+                }
+            }
+            _ => {
+                send_err(
+                    tx,
+                    AdapterError::Internal("expected copy to s3 oneshot sink".into()),
+                );
+                return;
+            }
+        }
+
+        // # Inlined from peek_copy_to_dataflow
+
+        let (df_desc, _df_meta) = global_lir_plan.unapply();
+
+        // Create and register ActiveCopyTo.
+        // Note: sink_tx/sink_rx is the channel for the compute sink to notify completion
+        // This is different from the command's tx which sends the response to the client
+        let (sink_tx, sink_rx) = oneshot::channel();
+        let active_copy_to = ActiveCopyTo {
+            conn_id: conn_id.clone(),
+            tx: sink_tx,
+            cluster_id: compute_instance,
+            depends_on: source_ids,
+        };
+
+        // Add metadata for the new COPY TO. CopyTo returns a `ready` future, so it is safe to drop.
+        drop(
+            self.add_active_compute_sink(sink_id, ActiveComputeSink::CopyTo(active_copy_to))
+                .await,
+        );
+
+        self.ship_dataflow(df_desc, compute_instance, target_replica)
+            .await;
+
+        // Spawn background task to wait for completion
+        // We must NOT await sink_rx here directly, as that would block the coordinator's main task
+        // from processing the completion message. Instead, we spawn a background task that will
+        // send the result through tx when the COPY TO completes.
+        let span = Span::current();
+        task::spawn(
+            || "copy to completion",
+            async move {
+                let res = sink_rx.await;
+                let result = match res {
+                    Ok(res) => res,
+                    Err(_) => Err(AdapterError::Internal("copy to sender dropped".into())),
+                };
+                let _ = tx.send(result);
+            }
+            .instrument(span),
+        );
+    }
+
     /// Constructs an [`ExecuteResponse`] that that will send some rows to the
     /// client immediately, as opposed to asking the dataflow layer to send along
     /// the rows after some computation.

src/adapter/src/coord/sequencer.rs

@@ -13,11 +13,13 @@
 //! Logic for executing a planned SQL query.
 
 use std::collections::BTreeSet;
+use std::str::FromStr;
 use std::sync::Arc;
 
 use futures::FutureExt;
 use futures::future::LocalBoxFuture;
 use futures::stream::FuturesOrdered;
+use http::Uri;
 use inner::return_if_err;
 use maplit::btreemap;
 use mz_catalog::memory::objects::Cluster;
@@ -28,12 +30,13 @@ use mz_ore::cast::CastFrom;
 use mz_ore::tracing::OpenTelemetryContext;
 use mz_persist_client::stats::SnapshotPartStats;
 use mz_repr::explain::{ExprHumanizerExt, TransientItem};
-use mz_repr::{CatalogItemId, Diff, GlobalId, IntoRowIterator, Row, Timestamp};
+use mz_repr::{CatalogItemId, Datum, Diff, GlobalId, IntoRowIterator, Row, RowArena, Timestamp};
 use mz_sql::catalog::{CatalogError, SessionCatalog};
 use mz_sql::names::ResolvedIds;
 use mz_sql::plan::{
     self, AbortTransactionPlan, CommitTransactionPlan, CopyFromSource, CreateRolePlan,
-    CreateSourcePlanBundle, FetchPlan, MutationKind, Params, Plan, PlanKind, RaisePlan,
+    CreateSourcePlanBundle, FetchPlan, HirScalarExpr, MutationKind, Params, Plan, PlanKind,
+    RaisePlan,
 };
 use mz_sql::rbac;
 use mz_sql::session::metadata::SessionMetadata;
@@ -50,7 +53,7 @@ use tokio::sync::oneshot;
 use tracing::{Instrument, Level, Span, event};
 
 use crate::ExecuteContext;
-use crate::catalog::Catalog;
+use crate::catalog::{Catalog, CatalogState};
 use crate::command::{Command, ExecuteResponse, Response};
 use crate::coord::appends::{DeferredOp, DeferredPlan};
 use crate::coord::validity::PlanValidity;
@@ -61,6 +64,7 @@ use crate::coord::{
 use crate::error::AdapterError;
 use crate::explain::insights::PlanInsightsContext;
 use crate::notice::AdapterNotice;
+use crate::optimize::dataflows::{EvalTime, ExprPrepStyle, prep_scalar_expr};
 use crate::optimize::peek;
 use crate::session::{
     EndTransactionAction, Session, StateRevision, TransactionOps, TransactionStatus, WriteOp,
@@ -1010,6 +1014,39 @@ pub(crate) fn emit_optimizer_notices(
     }
 }
 
+/// Evaluates a COPY TO target URI expression and validates it.
+///
+/// This function is shared between the old peek sequencing (sequence_copy_to)
+/// and the new frontend peek sequencing to avoid code duplication.
+pub fn eval_copy_to_uri(
+    to: HirScalarExpr,
+    session: &Session,
+    catalog_state: &CatalogState,
+) -> Result<Uri, AdapterError> {
+    let style = ExprPrepStyle::OneShot {
+        logical_time: EvalTime::NotAvailable,
+        session,
+        catalog_state,
+    };
+    let mut to = to.lower_uncorrelated()?;
+    prep_scalar_expr(&mut to, style)?;
+    let temp_storage = RowArena::new();
+    let evaled = to.eval(&[], &temp_storage)?;
+    if evaled == Datum::Null {
+        coord_bail!("COPY TO target value can not be null");
+    }
+    let to_url = match Uri::from_str(evaled.unwrap_str()) {
+        Ok(url) => {
+            if url.scheme_str() != Some("s3") {
+                coord_bail!("only 's3://...' urls are supported as COPY TO target");
+            }
+            url
+        }
+        Err(e) => coord_bail!("could not parse COPY TO target url: {}", e),
+    };
+    Ok(to_url)
+}
+
 /// Returns a future that will execute EXPLAIN FILTER PUSHDOWN, i.e., compute the filter pushdown
 /// statistics for the given collections with the given MFPs.
 ///