Changes for fixes for Bucket Region and Okta Poller

matanolabs · Jun 5, 2024 · 819d8c1 · 819d8c1
1 parent 739b4f2
commit 819d8c1
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/lib/rust/log_puller/src/pullers/okta.rs b/lib/rust/log_puller/src/pullers/okta.rs
@@ -157,6 +157,18 @@ impl PullLogs for OktaPuller {
                 .await?;
 
             let headers = response.headers().clone();
+
+            let status = response.status();
+            if !response.status().is_success() {
+                let msg = response.text().await.unwrap_or_default();
+                log::error!("{}", msg);
+                if status.is_client_error() {
+                    anyhow::bail!(format!("Client error: {}", msg));
+                } else {
+                    anyhow::bail!("Internal error")
+                }
+            }
+
             let response_json: Vec<serde_json::Value> = response.json().await?;
 
             // handle Okta paged responses containing `Link` header for 'self' and 'next'

diff --git a/lib/rust/transformer/src/main.rs b/lib/rust/transformer/src/main.rs
@@ -99,6 +99,14 @@ lazy_static! {
     };
 }
 
+fn bucket_location_constraint_to_aws_region(blc: &str) -> &str {
+    match blc {
+        "EU" => "eu-west-1",
+        "" => "us-east-1",
+        _ => blc,
+    }
+}
+
 async fn get_s3_client_using_access_role_cached(bucket: &str) -> aws_sdk_s3::Client {
     if let Some(access_role_arn) = CUSTOM_BUCKET_TO_ACCESS_ROLE_ARN_MAP.get(bucket) {
         let mut custom_bucket_to_region_cache = CUSTOM_BUCKET_TO_REGION_CACHE.lock().unwrap();
@@ -121,12 +129,10 @@ async fn get_s3_client_using_access_role_cached(bucket: &str) -> aws_sdk_s3::Cli
                     .await
                     .unwrap()
                     .location_constraint;
-                let region = match location_constraint {
-                    Some(location_constraint) => {
-                        Region::new(location_constraint.as_str().to_owned())
-                    }
-                    None => Region::new("us-east-1"),
-                };
+                let region = location_constraint
+                    .map(|lc| bucket_location_constraint_to_aws_region(lc.as_str()).to_string())
+                    .unwrap_or("us-east-1".to_string());
+                let region = Region::new(region);
                 custom_bucket_to_region_cache.insert(bucket.to_string(), region);
                 custom_bucket_to_region_cache.get(bucket).unwrap()
             }