fix: finalize S3 uploads for CSV/TBL writes

Author: prantogg

spatialbench-cli/src/generate.rs

@@ -45,12 +45,13 @@ pub trait Source: Send {
 /// Something that can write the contents of a buffer somewhere
 ///
 /// For example, this is implemented for a file writer.
-pub trait Sink: Send {
+pub trait Sink: Send + Sized {
     /// Write all data from the buffer to the sink
     fn sink(&mut self, buffer: &[u8]) -> Result<(), io::Error>;
 
-    /// Complete and flush any remaining data from the sink
-    fn flush(self) -> Result<(), io::Error>;
+    /// Complete and flush any remaining data from the sink, returning it
+    /// so the caller can perform additional finalization (e.g. async S3 upload).
+    fn flush(self) -> Result<Self, io::Error>;
 }
 
 /// Generates data in parallel from a series of [`Source`] and writes to a [`Sink`]
@@ -69,7 +70,7 @@ pub async fn generate_in_chunks<G, I, S>(
     mut sink: S,
     sources: I,
     num_threads: usize,
-) -> Result<(), io::Error>
+) -> Result<S, io::Error>
 where
     G: Source + 'static,
     I: Iterator<Item = G>,
@@ -86,7 +87,7 @@ where
 
     // write the header
     let Some(first) = sources.peek() else {
-        return Ok(()); // no sources
+        return Ok(sink); // no sources
     };
     let header = first.header(Vec::new());
     tx.send(header)
@@ -131,7 +132,8 @@ where
             sink.sink(&buffer)?;
             captured_recycler.return_buffer(buffer);
         }
-        // No more input, flush the sink and return
+        // No more input, flush the sink and return it so the caller can
+        // perform additional finalization (e.g. async S3 upload).
         sink.flush()
     });
 

spatialbench-cli/src/main.rs

@@ -421,6 +421,11 @@ impl<W: Write> WriterSink<W> {
             statistics: WriteStatistics::new("buffers"),
         }
     }
+
+    /// Consume the sink and return the inner writer for further finalization.
+    fn into_inner(self) -> W {
+        self.inner
+    }
 }
 
 impl<W: Write + Send> Sink for WriterSink<W> {
@@ -430,7 +435,8 @@ impl<W: Write + Send> Sink for WriterSink<W> {
         self.inner.write_all(buffer)
     }
 
-    fn flush(mut self) -> Result<(), io::Error> {
-        self.inner.flush()
+    fn flush(mut self) -> Result<Self, io::Error> {
+        self.inner.flush()?;
+        Ok(self)
     }
 }

spatialbench-cli/src/runner.rs

@@ -197,7 +197,8 @@ where
     match plan.output_location() {
         OutputLocation::Stdout => {
             let sink = WriterSink::new(io::stdout());
-            generate_in_chunks(sink, sources, num_threads).await
+            generate_in_chunks(sink, sources, num_threads).await?;
+            Ok(())
         }
         OutputLocation::File(path) => {
             // if the output already exists, skip running
@@ -224,7 +225,9 @@ where
             info!("Writing to S3: {}", uri);
             let s3_writer = S3Writer::with_client(Arc::clone(client), path);
             let sink = WriterSink::new(s3_writer);
-            generate_in_chunks(sink, sources, num_threads).await
+            let sink = generate_in_chunks(sink, sources, num_threads).await?;
+            sink.into_inner().finish().await?;
+            Ok(())
         }
     }
 }

spatialbench-cli/src/s3_writer.rs

@@ -533,4 +533,36 @@ mod tests {
         writer.write_all(&[4, 5]).unwrap();
         assert_eq!(writer.total_bytes(), 5);
     }
+
+    /// Verify that `std::io::Write::flush()` does NOT upload data to S3.
+    /// Data is only uploaded when `finish()` is called. This test guards
+    /// against the bug where CSV/TBL writes were silently lost because
+    /// the `WriterSink` called `flush()` (a no-op) but never `finish()`.
+    #[tokio::test]
+    async fn flush_does_not_upload_without_finish() {
+        let store = Arc::new(InMemory::new());
+        let mut writer = S3Writer::with_client(store.clone(), "output/flush_test.csv");
+
+        let data = b"col1,col2\nfoo,bar\n";
+        writer.write_all(data).unwrap();
+        writer.flush().unwrap();
+
+        // Data should NOT be in the store yet — flush is a no-op
+        let result = store.get(&ObjectPath::from("output/flush_test.csv")).await;
+        assert!(
+            result.is_err(),
+            "data should not be uploaded before finish()"
+        );
+
+        // Now call finish — data should appear
+        let total = writer.finish().await.unwrap();
+        assert_eq!(total, data.len());
+
+        let result = store
+            .get(&ObjectPath::from("output/flush_test.csv"))
+            .await
+            .unwrap();
+        let stored = result.bytes().await.unwrap();
+        assert_eq!(stored.as_ref(), data);
+    }
 }