cache-control to no-store for metadata uploads to S3

src/storage/metrics_layer.rs

@@ -26,8 +26,8 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use futures_util::{Stream, StreamExt, stream::BoxStream};
 use object_store::{
-    CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
-    PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions,
+    Attribute, CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
+    ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions,
     Result as ObjectStoreResult, path::Path,
 };
 
@@ -93,13 +93,28 @@ pub fn error_to_status_code(err: &object_store::Error) -> &'static str {
 pub struct MetricLayer<T: ObjectStore> {
     inner: T,
     provider: String,
+    cache_control_no_store: bool,
 }
 
 impl<T: ObjectStore> MetricLayer<T> {
     pub fn new(inner: T, provider: &str) -> Self {
         Self {
             inner,
             provider: provider.to_string(),
+            cache_control_no_store: false,
+        }
+    }
+
+    pub fn with_cache_control_no_store(mut self, enabled: bool) -> Self {
+        self.cache_control_no_store = enabled;
+        self
+    }
+
+    fn set_cache_control(&self, location: &Path, attributes: &mut object_store::Attributes) {
+        // Parquet data is immutable and may be cached, while mutable metadata
+        // (manifests, schemas, stream metadata, etc.) must not be cached.
+        if self.cache_control_no_store && !location.as_ref().ends_with(".parquet") {
+            attributes.insert(Attribute::CacheControl, "no-store".into());
         }
     }
 }
@@ -116,8 +131,9 @@ impl<T: ObjectStore> ObjectStore for MetricLayer<T> {
         &self,
         location: &Path,
         payload: PutPayload,
-        opts: PutOptions,
+        mut opts: PutOptions,
     ) -> ObjectStoreResult<PutResult> {
+        self.set_cache_control(location, &mut opts.attributes);
         let _guard = InflightGuard::new(&self.provider, "PUT");
         let time = time::Instant::now();
         let put_result = self.inner.put_opts(location, payload, opts).await;
@@ -137,8 +153,9 @@ impl<T: ObjectStore> ObjectStore for MetricLayer<T> {
     async fn put_multipart_opts(
         &self,
         location: &Path,
-        opts: PutMultipartOptions,
+        mut opts: PutMultipartOptions,
     ) -> ObjectStoreResult<Box<dyn MultipartUpload>> {
+        self.set_cache_control(location, &mut opts.attributes);
         let _guard = InflightGuard::new(&self.provider, "PUT_MULTIPART");
         let time = time::Instant::now();
         let result = self.inner.put_multipart_opts(location, opts).await;
@@ -324,3 +341,43 @@ impl<T> Stream for StreamMetricWrapper<'_, T> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use object_store::{Attribute, Attributes, memory::InMemory, path::Path};
+
+    use super::MetricLayer;
+
+    #[test]
+    fn s3_metadata_uploads_disable_caching() {
+        let layer = MetricLayer::new(InMemory::new(), "s3").with_cache_control_no_store(true);
+        let mut attributes = Attributes::new();
+
+        layer.set_cache_control(&Path::from("stream/stream.json"), &mut attributes);
+
+        assert_eq!(
+            attributes.get(&Attribute::CacheControl).map(AsRef::as_ref),
+            Some("no-store")
+        );
+    }
+
+    #[test]
+    fn s3_parquet_uploads_remain_cacheable() {
+        let layer = MetricLayer::new(InMemory::new(), "s3").with_cache_control_no_store(true);
+        let mut attributes = Attributes::new();
+
+        layer.set_cache_control(&Path::from("stream/events.parquet"), &mut attributes);
+
+        assert_eq!(attributes.get(&Attribute::CacheControl), None);
+    }
+
+    #[test]
+    fn cache_control_is_disabled_by_default() {
+        let layer = MetricLayer::new(InMemory::new(), "s3");
+        let mut attributes = Attributes::new();
+
+        layer.set_cache_control(&Path::from("stream/stream.json"), &mut attributes);
+
+        assert_eq!(attributes.get(&Attribute::CacheControl), None);
+    }
+}

src/storage/s3.rs

@@ -122,6 +122,15 @@ pub struct S3Config {
     )]
     pub set_checksum: bool,
 
+    /// Add `Cache-Control: no-store` to non-Parquet object uploads.
+    #[arg(
+        long,
+        env = "P_S3_CACHE_CONTROL_NO_STORE",
+        value_name = "bool",
+        default_value = "false"
+    )]
+    pub cache_control_no_store: bool,
+
     /// Set client to use virtual hosted style acess
     #[arg(
         long,
@@ -335,7 +344,8 @@ impl ObjectStorageProvider for S3Config {
 
         // limit objectstore to a concurrent request limit
         let s3 = LimitStore::new(s3, super::MAX_OBJECT_STORE_REQUESTS);
-        let s3 = MetricLayer::new(s3, "s3");
+        let s3 =
+            MetricLayer::new(s3, "s3").with_cache_control_no_store(self.cache_control_no_store);
 
         let object_store_registry = DefaultObjectStoreRegistry::new();
         let url = ObjectStoreUrl::parse(format!("s3://{}", self.bucket_name)).unwrap();
@@ -348,7 +358,8 @@ impl ObjectStorageProvider for S3Config {
         let s3 = self.get_default_builder().build().unwrap();
         // limit objectstore to a concurrent request limit
         let s3 = LimitStore::new(s3, super::MAX_OBJECT_STORE_REQUESTS);
-        let s3 = MetricLayer::new(s3, "s3");
+        let s3 =
+            MetricLayer::new(s3, "s3").with_cache_control_no_store(self.cache_control_no_store);
         Arc::new(S3 {
             client: Arc::new(s3),
             bucket: self.bucket_name.clone(),