Region configuration to namespace blob stores (#1655)

* Added region to the remaining blob config and used it when initializing S3 blob stores

---------

Co-authored-by: Julio Martinez <julio@tensorflow.ai>
Co-authored-by: Julio Martínez <julio@tensorlake.ai>

Author: 3 people

server/src/blob_store/mod.rs

@@ -28,6 +28,7 @@ pub mod registry;
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct BlobStorageConfig {
     pub path: String,
+    pub region: String,
 }
 
 impl Default for BlobStorageConfig {
@@ -42,6 +43,7 @@ impl Default for BlobStorageConfig {
         );
         BlobStorageConfig {
             path: blob_store_path,
+            region: "".to_string(),
         }
     }
 }
@@ -65,7 +67,7 @@ impl BlobStorage {
     pub fn new(config: BlobStorageConfig) -> Result<Self> {
         let url = &config.path.clone();
         debug!("using blob store path: {}", url);
-        let (object_store, path) = Self::build_object_store(url)?;
+        let (object_store, path) = Self::build_object_store(url, &config.region)?;
         Ok(Self {
             object_store: Arc::new(object_store),
             url_scheme: url.parse::<Url>()?.scheme().to_string(),
@@ -75,7 +77,7 @@ impl BlobStorage {
         })
     }
 
-    pub fn build_object_store(url_str: &str) -> Result<(Box<dyn ObjectStore>, Path)> {
+    pub fn build_object_store(url_str: &str, region: &str) -> Result<(Box<dyn ObjectStore>, Path)> {
         let url = &url_str.parse::<Url>()?;
         let (scheme, _) = ObjectStoreScheme::parse(url)?;
         match scheme {
@@ -86,6 +88,7 @@ impl BlobStorage {
                     .with_url(url_str)
                     .with_allow_http(true)
                     .with_conditional_put(S3ConditionalPut::ETagMatch)
+                    .with_region(region)
                     .build()
                     .expect("failed to create object store");
                 let (_, path) = parse_url(url)?;

server/src/blob_store/registry.rs

@@ -13,19 +13,21 @@ pub struct BlobStorageRegistry {
 }
 
 impl BlobStorageRegistry {
-    pub fn new(default_blob_storage_path: &str) -> Result<Self> {
+    pub fn new(default_blob_storage_path: &str, default_blob_storage_region: &str) -> Result<Self> {
         let default_blob_storage = BlobStorage::new(BlobStorageConfig {
             path: default_blob_storage_path.to_string(),
+            region: default_blob_storage_region.to_string(),
         })?;
         Ok(Self {
             blob_storage_buckets: Mutex::new(HashMap::new()),
             default_blob_storage: Arc::new(default_blob_storage),
         })
     }
 
-    pub fn create_new_blob_store(&self, namespace: &str, path: &str) -> Result<()> {
+    pub fn create_new_blob_store(&self, namespace: &str, path: &str, region: &str) -> Result<()> {
         let blob_storage = BlobStorage::new(BlobStorageConfig {
             path: path.to_string(),
+            region: region.to_string(),
         })?;
         self.blob_storage_buckets
             .lock()

server/src/data_model/mod.rs

@@ -2162,6 +2162,7 @@ pub struct Namespace {
     pub name: String,
     pub created_at: u64,
     pub blob_storage_bucket: Option<String>,
+    pub blob_storage_region: Option<String>,
 }
 
 #[cfg(test)]

server/src/http_objects.rs

@@ -633,6 +633,7 @@ impl From<data_model::ComputeGraph> for ComputeGraph {
 pub struct CreateNamespace {
     pub name: String,
     pub blob_storage_bucket: Option<String>,
+    pub blob_storage_region: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, ToSchema)]

server/src/routes_internal.rs

@@ -282,6 +282,7 @@ async fn namespace_middleware(
                     payload: RequestPayload::CreateNameSpace(NamespaceRequest {
                         name: namespace.to_string(),
                         blob_storage_bucket: None,
+                        blob_storage_region: None,
                     }),
                     processed_state_changes: vec![],
                 })
@@ -311,10 +312,16 @@ async fn create_namespace(
     Json(namespace): Json<CreateNamespace>,
 ) -> Result<(), IndexifyAPIError> {
     if let Some(blob_storage_bucket) = &namespace.blob_storage_bucket {
-        if let Err(e) = state
-            .blob_storage
-            .create_new_blob_store(&namespace.name, blob_storage_bucket)
-        {
+        let Some(blob_storage_region) = &namespace.blob_storage_region else {
+            return Err(IndexifyAPIError::bad_request(
+                "blob storage region is required",
+            ));
+        };
+        if let Err(e) = state.blob_storage.create_new_blob_store(
+            &namespace.name,
+            blob_storage_bucket,
+            &blob_storage_region,
+        ) {
             error!("failed to create blob storage bucket: {:?}", e);
             return Err(IndexifyAPIError::internal_error(e));
         }
@@ -323,6 +330,7 @@ async fn create_namespace(
         payload: RequestPayload::CreateNameSpace(NamespaceRequest {
             name: namespace.name.clone(),
             blob_storage_bucket: namespace.blob_storage_bucket,
+            blob_storage_region: namespace.blob_storage_region,
         }),
         processed_state_changes: vec![],
     };

server/src/routes_v1.rs

@@ -196,6 +196,7 @@ async fn namespace_middleware(
                     payload: RequestPayload::CreateNameSpace(NamespaceRequest {
                         name: namespace.to_string(),
                         blob_storage_bucket: None,
+                        blob_storage_region: None,
                     }),
                     processed_state_changes: vec![],
                 })

server/src/service.rs

@@ -64,14 +64,23 @@ impl Service {
 
         let indexify_state = IndexifyState::new(config.state_store_path.parse()?).await?;
 
-        let blob_storage_registry =
-            Arc::new(BlobStorageRegistry::new(config.blob_storage.path.as_str())?);
+        let blob_storage_registry = Arc::new(BlobStorageRegistry::new(
+            config.blob_storage.path.as_str(),
+            config.blob_storage.region.as_str(),
+        )?);
 
         let namespaces = indexify_state.reader().get_all_namespaces()?;
         for namespace in namespaces {
             if let Some(blob_storage_bucket) = namespace.blob_storage_bucket {
-                blob_storage_registry
-                    .create_new_blob_store(&namespace.name, &blob_storage_bucket)?;
+                let blob_storage_region = namespace.blob_storage_region.unwrap_or_else(|| {
+                    std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string())
+                });
+
+                blob_storage_registry.create_new_blob_store(
+                    &namespace.name,
+                    &blob_storage_bucket,
+                    &blob_storage_region,
+                )?;
             }
         }
 

server/src/state_store/in_memory_state.rs

@@ -592,6 +592,7 @@ impl InMemoryState {
                         name: req.name.clone(),
                         created_at,
                         blob_storage_bucket: req.blob_storage_bucket.clone(),
+                        blob_storage_region: req.blob_storage_region.clone(),
                     }),
                 );
             }

server/src/state_store/kv.rs

@@ -91,6 +91,7 @@ mod tests {
                     "file://{}",
                     temp_dir.path().join("blob_store").to_str().unwrap()
                 ),
+                region: "local".to_string(),
             })?),
             "test",
         )

server/src/state_store/mod.rs

@@ -442,6 +442,7 @@ mod tests {
                 payload: RequestPayload::CreateNameSpace(NamespaceRequest {
                     name: "namespace1".to_string(),
                     blob_storage_bucket: None,
+                    blob_storage_region: None,
                 }),
                 processed_state_changes: vec![],
             })
@@ -451,6 +452,7 @@ mod tests {
                 payload: RequestPayload::CreateNameSpace(NamespaceRequest {
                     name: "namespace2".to_string(),
                     blob_storage_bucket: Some("bucket2".to_string()),
+                    blob_storage_region: Some("local".to_string()),
                 }),
                 processed_state_changes: vec![],
             })