Merge pull request #30844 from benesch/copy-to-s3-empty

Reapply "storage/copy-to-s3: emit empty file even if input is empty"

Author: benesch

Cargo.lock

@@ -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.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,44 @@ 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()
+                    .expect("known to be exactly one copy to s3 sink");
+                match &sink_desc.connection {
+                    ComputeSinkConnection::CopyToS3Oneshot(conn) => {
+                        mz_storage_types::sinks::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/coord/sequencer/inner/peek.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");

src/adapter/src/optimize/copy_to.rs

@@ -51,10 +51,6 @@ pub struct S3MultiPartUploader {
     upload_handles: Vec<JoinHandle<Result<(Option<String>, i32), S3MultiPartUploadError>>>,
 }
 
-/// The smallest allowable part number (inclusive).
-///
-/// From <https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html>
-const AWS_S3_MIN_PART_COUNT: i32 = 1;
 /// The largest allowable part number (inclusive).
 ///
 /// From <https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html>
@@ -218,18 +214,12 @@ impl S3MultiPartUploader {
         }
     }
 
-    /// Method to finish the multi part upload. If the buffer is not empty,
-    /// it flushes the buffer first and then makes a call to `complete_multipart_upload`.
+    /// Finishes the multi part upload.
+    ///
     /// Returns the number of parts and number of bytes uploaded.
     pub async fn finish(mut self) -> Result<CompletedUpload, S3MultiPartUploadError> {
-        if self.buffer.len() > 0 {
-            let remaining = self.buffer.split();
-            self.upload_part_internal(remaining.freeze())?;
-        }
-
-        if self.part_count < AWS_S3_MIN_PART_COUNT {
-            return Err(S3MultiPartUploadError::AtLeastMinPartNumber);
-        }
+        let remaining = self.buffer.split();
+        self.upload_part_internal(remaining.freeze())?;
 
         let mut parts: Vec<CompletedPart> = Vec::with_capacity(self.upload_handles.len());
         for handle in self.upload_handles {
@@ -336,11 +326,6 @@ pub enum S3MultiPartUploadError {
         AWS_S3_MAX_PART_COUNT
     )]
     ExceedsMaxPartNumber,
-    #[error(
-        "multi-part upload should have at least {} part",
-        AWS_S3_MIN_PART_COUNT
-    )]
-    AtLeastMinPartNumber,
     #[error("multi-part upload will exceed configured file_size_limit: {} bytes", .0)]
     UploadExceedsMaxFileLimit(u64),
     #[error("{}", .0.display_with_causes())]
@@ -503,7 +488,7 @@ mod tests {
     #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
     #[cfg_attr(coverage, ignore)] // https://github.com/MaterializeInc/database-issues/issues/5586
     #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_method` on OS `linux`
-    async fn multi_part_upload_error() -> Result<(), S3MultiPartUploadError> {
+    async fn multi_part_upload_no_data() -> Result<(), S3MultiPartUploadError> {
         let sdk_config = defaults().load().await;
         let (bucket, key) = match s3_bucket_key_for_test() {
             Some(tuple) => tuple,
@@ -514,12 +499,20 @@ mod tests {
         let uploader =
             S3MultiPartUploader::try_new(&sdk_config, bucket.clone(), key.clone(), config).await?;
 
-        // Calling finish without adding any data should error
-        let err = uploader.finish().await.unwrap_err();
-        assert_eq!(
-            err.to_string(),
-            "multi-part upload should have at least 1 part"
-        );
+        // Calling finish without adding any data should succeed.
+        uploader.finish().await.unwrap();
+
+        // The file should exist but have no content.
+        let s3_client = s3::new_client(&sdk_config);
+        let uploaded_object = s3_client
+            .get_object()
+            .bucket(bucket)
+            .key(key)
+            .send()
+            .await
+            .unwrap();
+
+        assert_eq!(uploaded_object.content_length(), Some(0));
 
         Ok(())
     }

src/aws-util/src/s3_uploader.rs

@@ -124,7 +124,7 @@ where
                 .get(&compute_state.worker_config),
         };
 
-        mz_storage_operators::s3_oneshot_sink::copy_to(
+        let token = mz_storage_operators::s3_oneshot_sink::copy_to(
             input,
             error_stream,
             sink.up_to.clone(),
@@ -138,6 +138,7 @@ where
         );
 
         Some(Rc::new(scopeguard::guard((), move |_| {
+            let _token = token;
             if let Some(protocol_handle) = response_protocol_weak.upgrade() {
                 std::mem::drop(protocol_handle.borrow_mut().take())
             }