storage: move copy to preflight to adapter

A copy to operation involves some preflight checks: validating that the
requested prefix of S3 is empty, and writing an INCOMPLETE sentinel
file.

Previously these checks were run as the first operation on every replica
participating in the copy to operation. This was racy though: a lagging
replica could write the INCOMPLETE sentinel file *after* the leading
replica finished the copy and had thus *deleted* the INCOMPLETE sentinel
file. Since the lagging replica's work is canceled once the leading
replica's work is finished, this could result in a INCOMPLETE sentinel
file that never got removed.

This commit moves these preflight checks to the adapter, where they are
performed exactly once before any replica is notified of the copy to
operation. This ensures that the INCOMPLETE file is written at most once
in an copy to operation, and in particular that lagging replicas can't
cause the file to reappear afer the leading replica finishes the
operation.

src/adapter/BUILD.bazel

@@ -65,6 +65,7 @@ rust_library(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",
@@ -135,6 +136,7 @@ rust_test(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",
@@ -185,6 +187,7 @@ rust_doc_test(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",
@@ -255,6 +258,7 @@ rust_test(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",
@@ -325,6 +329,7 @@ rust_test(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",
@@ -395,6 +400,7 @@ rust_test(
         "//src/sql-parser:mz_sql_parser",
         "//src/ssh-util:mz_ssh_util",
         "//src/storage-client:mz_storage_client",
+        "//src/storage-operators:mz_storage_operators",
         "//src/storage-types:mz_storage_types",
         "//src/timestamp-oracle:mz_timestamp_oracle",
         "//src/tls-util:mz_tls_util",

src/adapter/Cargo.toml

@@ -67,6 +67,7 @@ mz-sql = { path = "../sql" }
 mz-sql-parser = { path = "../sql-parser" }
 mz-ssh-util = { path = "../ssh-util" }
 mz-storage-client = { path = "../storage-client" }
+mz-storage-operators = { path = "../storage-operators" }
 mz-storage-types = { path = "../storage-types" }
 mz-tls-util = { path = "../tls-util" }
 mz-tracing = { path = "../tracing" }

src/adapter/src/coord.rs

@@ -426,8 +426,10 @@ pub enum PeekStage {
     /// Final stage for an explain.
     ExplainPlan(PeekStageExplainPlan),
     ExplainPushdown(PeekStageExplainPushdown),
-    /// Final stage for a copy to.
-    CopyTo(PeekStageCopyTo),
+    /// Preflight checks for a copy to operation.
+    CopyToPreflight(PeekStageCopyTo),
+    /// Final stage for a copy to which involves shipping the dataflow.
+    CopyToDataflow(PeekStageCopyTo),
 }
 
 #[derive(Debug)]

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

@@ -14,6 +14,7 @@ use std::sync::Arc;
 use http::Uri;
 use itertools::Either;
 use maplit::btreemap;
+use mz_compute_types::sinks::ComputeSinkConnection;
 use mz_controller_types::ClusterId;
 use mz_expr::{CollectionPlan, ResultSpec};
 use mz_ore::cast::CastFrom;
@@ -69,7 +70,8 @@ impl Staged for PeekStage {
             PeekStage::Finish(stage) => &mut stage.validity,
             PeekStage::ExplainPlan(stage) => &mut stage.validity,
             PeekStage::ExplainPushdown(stage) => &mut stage.validity,
-            PeekStage::CopyTo(stage) => &mut stage.validity,
+            PeekStage::CopyToPreflight(stage) => &mut stage.validity,
+            PeekStage::CopyToDataflow(stage) => &mut stage.validity,
         }
     }
 
@@ -94,7 +96,8 @@ impl Staged for PeekStage {
             PeekStage::ExplainPushdown(stage) => {
                 coord.peek_explain_pushdown(ctx.session(), stage).await
             }
-            PeekStage::CopyTo(stage) => coord.peek_copy_to_dataflow(ctx, stage).await,
+            PeekStage::CopyToPreflight(stage) => coord.peek_copy_to_preflight(stage).await,
+            PeekStage::CopyToDataflow(stage) => coord.peek_copy_to_dataflow(ctx, stage).await,
         }
     }
 
@@ -687,7 +690,7 @@ impl Coordinator {
                         }
                         Ok(Either::Right(global_lir_plan)) => {
                             let optimizer = optimizer.unwrap_right();
-                            PeekStage::CopyTo(PeekStageCopyTo {
+                            PeekStage::CopyToPreflight(PeekStageCopyTo {
                                 validity,
                                 optimizer,
                                 global_lir_plan,
@@ -947,6 +950,42 @@ impl Coordinator {
         Ok(StageResult::Response(resp))
     }
 
+    #[instrument]
+    async fn peek_copy_to_preflight(
+        &mut self,
+        copy_to: PeekStageCopyTo,
+    ) -> Result<StageResult<Box<PeekStage>>, AdapterError> {
+        let connection_context = self.connection_context().clone();
+        Ok(StageResult::Handle(mz_ore::task::spawn(
+            || "peek copy to preflight",
+            async {
+                let sinks = &copy_to.global_lir_plan.df_desc().sink_exports;
+                if sinks.len() != 1 {
+                    return Err(AdapterError::Internal(
+                        "expected exactly one copy to s3 sink".into(),
+                    ));
+                }
+                let (sink_id, sink_desc) = sinks.first_key_value().unwrap();
+                match &sink_desc.connection {
+                    ComputeSinkConnection::CopyToS3Oneshot(conn) => {
+                        mz_storage_operators::s3_oneshot_sink::preflight(
+                            connection_context,
+                            &conn.aws_connection,
+                            &conn.upload_info,
+                            conn.connection_id,
+                            *sink_id,
+                        )
+                        .await?;
+                        Ok(Box::new(PeekStage::CopyToDataflow(copy_to)))
+                    }
+                    _ => Err(AdapterError::Internal(
+                        "expected copy to s3 oneshot sink".into(),
+                    )),
+                }
+            },
+        )))
+    }
+
     #[instrument]
     async fn peek_copy_to_dataflow(
         &mut self,

src/adapter/src/optimize/copy_to.rs

@@ -139,6 +139,10 @@ pub struct GlobalLirPlan {
 }
 
 impl GlobalLirPlan {
+    pub fn df_desc(&self) -> &LirDataflowDescription {
+        &self.df_desc
+    }
+
     pub fn sink_id(&self) -> GlobalId {
         let sink_exports = &self.df_desc.sink_exports;
         let sink_id = sink_exports.keys().next().expect("valid sink");