pass bucket to every storage call (#786)

* pass bucket to every storage call

* pass bucket in queue

Author: dinmukhamedm

app-server/src/api/v1/datasets.rs

@@ -1,4 +1,4 @@
-use std::{env, sync::Arc};
+use std::sync::Arc;
 
 use actix_web::{HttpResponse, get, post, web};
 use futures_util::StreamExt;
@@ -32,7 +32,6 @@ async fn get_datapoints(
     let clickhouse = clickhouse.into_inner().as_ref().clone();
     let query = params.into_inner();
 
-    // Still get dataset metadata from PostgreSQL
     let dataset_id =
         db::datasets::get_dataset_id_by_name(&db.pool, &query.name, project_id).await?;
 
@@ -52,10 +51,8 @@ async fn get_datapoints(
     )
     .await?;
 
-    // Get total count from ClickHouse
     let total_count = ch_datapoints::count_datapoints(clickhouse, project_id, dataset_id).await?;
 
-    // Convert CHDatapoints to Datapoints
     let datapoints: Vec<Datapoint> = ch_datapoints
         .into_iter()
         .map(|ch_dp| ch_dp.into())
@@ -108,7 +105,6 @@ async fn create_datapoints(
         })));
     }
 
-    // Get dataset metadata from PostgreSQL
     let dataset_id =
         db::datasets::get_dataset_id_by_name(&db.pool, &request.dataset_name, project_id).await?;
 
@@ -131,13 +127,11 @@ async fn create_datapoints(
         })
         .collect();
 
-    // Convert to ClickHouse datapoints
     let ch_datapoints: Vec<ch_datapoints::CHDatapoint> = datapoints
         .iter()
         .map(|dp| ch_datapoints::CHDatapoint::from_datapoint(dp, project_id))
         .collect();
 
-    // Insert into ClickHouse
     ch_datapoints::insert_datapoints(clickhouse, ch_datapoints).await?;
 
     Ok(HttpResponse::Created().json(serde_json::json!({
@@ -160,7 +154,6 @@ async fn get_parquet(
     let project_id = project_api_key.project_id;
     let db = db.into_inner();
 
-    // Get parquet paths from database
     let parquet_path =
         db::datasets::get_parquet_path(&db.pool, project_id, dataset_id, &name).await?;
 
@@ -170,15 +163,14 @@ async fn get_parquet(
         })));
     };
 
-    // Get object metadata to determine file size
-    let content_length = storage
-        .get_size(&parquet_path, &env::var("S3_EXPORTS_BUCKET").ok())
-        .await?;
+    let Ok(bucket) = std::env::var("S3_EXPORTS_BUCKET") else {
+        return Ok(HttpResponse::InternalServerError().json(serde_json::json!({
+            "error": "exports storage is not configured"
+        })));
+    };
+    let content_length = storage.get_size(&bucket, &parquet_path).await?;
 
-    // Stream the file from S3
-    let get_response = storage
-        .get_stream(&parquet_path, &env::var("S3_EXPORTS_BUCKET").ok())
-        .await?;
+    let get_response = storage.get_stream(&bucket, &parquet_path).await?;
 
     let filename = parquet_path.split('/').last().unwrap_or(&name);
 

app-server/src/api/v1/payloads.rs

@@ -31,7 +31,12 @@ pub async fn get_payload(
     let key = format!("project/{}/{}", project_id, payload_id);
 
     // Get the payload stream from storage
-    let mut stream = match storage.as_ref().get_stream(&key, &None).await {
+    let Ok(bucket) = std::env::var("S3_TRACE_PAYLOADS_BUCKET") else {
+        return Ok(HttpResponse::InternalServerError().json(serde_json::json!({
+            "error": "payloads storage is not configured"
+        })));
+    };
+    let mut stream = match storage.as_ref().get_stream(&bucket, &key).await {
         Ok(stream) => stream,
         Err(e) => {
             log::error!("Failed to retrieve payload from storage: {:?}", e);

app-server/src/language_model/chat_message.rs

@@ -436,6 +436,7 @@ impl ChatMessageContentPart {
         &self,
         project_id: &Uuid,
         storage: Arc<Storage>,
+        bucket: &str,
     ) -> Result<ChatMessageContentPart> {
         match self {
             ChatMessageContentPart::Image(image) => {
@@ -460,7 +461,7 @@ impl ChatMessageContentPart {
                         // Leave intact in case of error
                         return Ok(self.clone());
                     }
-                    let url = storage.store(data, &key).await?;
+                    let url = storage.store(&bucket, &key, data).await?;
                     Ok(ChatMessageContentPart::ImageUrl(ChatMessageImageUrl {
                         url,
                         detail: Some(format!("media_type:{};base64", media_type)),
@@ -484,7 +485,7 @@ impl ChatMessageContentPart {
                     // Leave intact in case of error
                     return Ok(self.clone());
                 }
-                let url = storage.store(data, &key).await?;
+                let url = storage.store(&bucket, &key, data).await?;
                 Ok(ChatMessageContentPart::DocumentUrl(
                     ChatMessageDocumentUrl {
                         media_type: document.source.media_type.clone(),
@@ -505,7 +506,7 @@ impl ChatMessageContentPart {
                         // Leave intact in case of error
                         return Ok(self.clone());
                     }
-                    let url = storage.store(data, &key).await?;
+                    let url = storage.store(&bucket, &key, data).await?;
                     Ok(ChatMessageContentPart::ImageUrl(ChatMessageImageUrl {
                         url,
                         detail: image_url.detail.clone(),
@@ -526,7 +527,7 @@ impl ChatMessageContentPart {
                     // Leave intact in case of error
                     return Ok(self.clone());
                 }
-                let url = storage.store(image.image.clone(), &key).await?;
+                let url = storage.store(&bucket, &key, image.image.clone()).await?;
                 Ok(ChatMessageContentPart::ImageUrl(ChatMessageImageUrl {
                     url,
                     detail: image

app-server/src/main.rs

@@ -428,8 +428,6 @@ fn main() -> anyhow::Result<()> {
                     let s3_client = aws_sdk_s3::Client::new(&aws_sdk_config);
                     let s3_storage = storage::s3::S3Storage::new(
                         s3_client,
-                        env::var("S3_TRACE_PAYLOADS_BUCKET")
-                            .expect("S3_TRACE_PAYLOADS_BUCKET must be set"),
                         mq_for_http.clone(),
                     );
                     Arc::new(s3_storage.into())

app-server/src/storage/mock.rs

@@ -9,25 +9,21 @@ pub struct MockStorage;
 impl super::StorageTrait for MockStorage {
     type StorageBytesStream =
         Pin<Box<dyn futures_util::stream::Stream<Item = bytes::Bytes> + Send + 'static>>;
-    async fn store(&self, _data: Vec<u8>, _key: &str) -> Result<String> {
+    async fn store(&self, _bucket: &str, _key: &str, _data: Vec<u8>) -> Result<String> {
         Ok("mock".to_string())
     }
 
-    async fn store_direct(&self, _data: Vec<u8>, _key: &str) -> Result<String> {
+    async fn store_direct(&self, _bucket: &str, _key: &str, _data: Vec<u8>) -> Result<String> {
         Ok("mock".to_string())
     }
 
-    async fn get_stream(
-        &self,
-        _key: &str,
-        _bucket: &Option<String>,
-    ) -> Result<Self::StorageBytesStream> {
+    async fn get_stream(&self, _bucket: &str, _key: &str) -> Result<Self::StorageBytesStream> {
         Ok(Box::pin(futures_util::stream::once(async move {
             bytes::Bytes::new()
         })))
     }
 
-    async fn get_size(&self, _key: &str, _bucket: &Option<String>) -> Result<u64> {
+    async fn get_size(&self, _bucket: &str, _key: &str) -> Result<u64> {
         Ok(0)
     }
 }

app-server/src/storage/mod.rs

@@ -22,6 +22,7 @@ pub const PAYLOADS_ROUTING_KEY: &str = "payloads_routing_key";
 pub struct QueuePayloadMessage {
     pub key: String,
     pub data: Vec<u8>,
+    pub bucket: String,
 }
 
 use mock::MockStorage;
@@ -37,14 +38,10 @@ pub enum Storage {
 #[enum_delegate::register]
 pub trait StorageTrait {
     type StorageBytesStream: futures_util::stream::Stream<Item = bytes::Bytes>;
-    async fn store(&self, data: Vec<u8>, key: &str) -> Result<String>;
-    async fn store_direct(&self, data: Vec<u8>, key: &str) -> Result<String>;
-    async fn get_stream(
-        &self,
-        key: &str,
-        bucket: &Option<String>,
-    ) -> Result<Self::StorageBytesStream>;
-    async fn get_size(&self, key: &str, bucket: &Option<String>) -> Result<u64>;
+    async fn store(&self, bucket: &str, key: &str, data: Vec<u8>) -> Result<String>;
+    async fn store_direct(&self, bucket: &str, key: &str, data: Vec<u8>) -> Result<String>;
+    async fn get_stream(&self, bucket: &str, key: &str) -> Result<Self::StorageBytesStream>;
+    async fn get_size(&self, bucket: &str, key: &str) -> Result<u64>;
 }
 
 pub fn create_key(project_id: &Uuid, file_extension: &Option<String>) -> String {
@@ -118,7 +115,7 @@ async fn inner_process_payloads(storage: Arc<Storage>, queue: Arc<MessageQueue>)
         };
 
         let store_payload = || async {
-            storage.store_direct(message.data.clone(), &message.key).await.map_err(|e| {
+            storage.store_direct(&message.bucket, &message.key, message.data.clone()).await.map_err(|e| {
                 log::error!("Failed attempt to store payload. Will retry according to backoff policy. Error: {:?}", e);
                 backoff::Error::transient(e)
             })

app-server/src/storage/s3.rs

@@ -11,17 +11,12 @@ use crate::{
 #[derive(Clone)]
 pub struct S3Storage {
     client: Client,
-    bucket: String,
     queue: Arc<MessageQueue>,
 }
 
 impl S3Storage {
-    pub fn new(client: Client, bucket: String, queue: Arc<MessageQueue>) -> Self {
-        Self {
-            client,
-            bucket,
-            queue,
-        }
+    pub fn new(client: Client, queue: Arc<MessageQueue>) -> Self {
+        Self { client, queue }
     }
 
     fn get_url(&self, key: &str) -> String {
@@ -38,11 +33,12 @@ impl S3Storage {
 impl super::StorageTrait for S3Storage {
     type StorageBytesStream =
         Pin<Box<dyn futures_util::stream::Stream<Item = bytes::Bytes> + Send + 'static>>;
-    async fn store(&self, data: Vec<u8>, key: &str) -> Result<String> {
+    async fn store(&self, bucket: &str, key: &str, data: Vec<u8>) -> Result<String> {
         // Push to queue instead of storing directly
         let message = QueuePayloadMessage {
             key: key.to_string(),
             data,
+            bucket: bucket.to_string(),
         };
 
         self.queue
@@ -57,11 +53,11 @@ impl super::StorageTrait for S3Storage {
         Ok(self.get_url(key))
     }
 
-    async fn store_direct(&self, data: Vec<u8>, key: &str) -> Result<String> {
+    async fn store_direct(&self, bucket: &str, key: &str, data: Vec<u8>) -> Result<String> {
         // Direct storage method used by the payload worker
         self.client
             .put_object()
-            .bucket(&self.bucket)
+            .bucket(bucket)
             .key(key)
             .body(data.into())
             .send()
@@ -70,15 +66,11 @@ impl super::StorageTrait for S3Storage {
         Ok(self.get_url(key))
     }
 
-    async fn get_stream(
-        &self,
-        key: &str,
-        bucket: &Option<String>,
-    ) -> Result<Self::StorageBytesStream> {
+    async fn get_stream(&self, bucket: &str, key: &str) -> Result<Self::StorageBytesStream> {
         let response = self
             .client
             .get_object()
-            .bucket(bucket.as_ref().unwrap_or(&self.bucket))
+            .bucket(bucket)
             .key(key)
             .send()
             .await?;
@@ -92,11 +84,11 @@ impl super::StorageTrait for S3Storage {
         )))
     }
 
-    async fn get_size(&self, key: &str, bucket: &Option<String>) -> Result<u64> {
+    async fn get_size(&self, bucket: &str, key: &str) -> Result<u64> {
         let response = self
             .client
             .head_object()
-            .bucket(bucket.as_ref().unwrap_or(&self.bucket))
+            .bucket(bucket)
             .key(key)
             .send()
             .await?;

app-server/src/traces/spans.rs

@@ -758,6 +758,10 @@ impl Span {
             .ok()
             .and_then(|s: String| s.parse::<usize>().ok())
             .unwrap_or(DEFAULT_PAYLOAD_SIZE_THRESHOLD);
+        let Ok(bucket) = std::env::var("S3_TRACE_PAYLOADS_BUCKET") else {
+            log::error!("S3_TRACE_PAYLOADS_BUCKET is not set");
+            return Err(anyhow::anyhow!("S3_TRACE_PAYLOADS_BUCKET is not set"));
+        };
         if let Some(input) = self.input.clone() {
             let span_input = serde_json::from_value::<Vec<ChatMessage>>(input);
             if let Ok(span_input) = span_input {
@@ -766,14 +770,16 @@ impl Span {
                     if let ChatMessageContent::ContentPartList(parts) = message.content {
                         let mut new_parts = Vec::new();
                         for part in parts {
-                            let stored_part =
-                                match part.store_media(project_id, storage.clone()).await {
-                                    Ok(stored_part) => stored_part,
-                                    Err(e) => {
-                                        log::error!("Error storing media: {e}");
-                                        part
-                                    }
-                                };
+                            let stored_part = match part
+                                .store_media(project_id, storage.clone(), &bucket)
+                                .await
+                            {
+                                Ok(stored_part) => stored_part,
+                                Err(e) => {
+                                    log::error!("Error storing media: {e}");
+                                    part
+                                }
+                            };
                             new_parts.push(stored_part);
                         }
                         message.content = ChatMessageContent::ContentPartList(new_parts);
@@ -801,7 +807,7 @@ impl Span {
                             data.len()
                         );
                     } else {
-                        let url = storage.store(data, &key).await?;
+                        let url = storage.store(&bucket, &key, data).await?;
                         self.input_url = Some(url);
                         self.input = Some(serde_json::Value::String(preview));
                     }
@@ -828,7 +834,7 @@ impl Span {
                         data.len()
                     );
                 } else {
-                    let url = storage.store(data, &key).await?;
+                    let url = storage.store(&bucket, &key, data).await?;
                     self.output_url = Some(url);
                     self.output = Some(serde_json::Value::String(
                         output_str.chars().take(100).collect(),