fix: reuse function to detect right aws partition and support gov too

Author: alexgallotta

bottlecap/src/bin/bottlecap/main.rs

@@ -9,6 +9,7 @@
 #![deny(missing_copy_implementations)]
 #![deny(missing_debug_implementations)]
 
+use bottlecap::config::get_aws_partition_by_region;
 use bottlecap::{
     base_url,
     config::{self, flush_strategy::FlushStrategy, AwsConfig, Config},
@@ -156,13 +157,8 @@ async fn register(client: &reqwest::Client) -> Result<RegisterResponse> {
 }
 
 fn build_function_arn(account_id: &str, region: &str, function_name: &str) -> String {
-    let arn_prefix = if region.starts_with("cn-") {
-        "aws-cn"
-    } else {
-        "aws"
-    };
-
-    format!("arn:{arn_prefix}:lambda:{region}:{account_id}:function:{function_name}")
+    let aws_partition = get_aws_partition_by_region(region);
+    format!("arn:{aws_partition}:lambda:{region}:{account_id}:function:{function_name}")
 }
 
 #[tokio::main]

bottlecap/src/config/mod.rs

@@ -311,6 +311,15 @@ pub struct AwsConfig {
     pub sandbox_init_time: Instant,
 }
 
+#[must_use]
+pub fn get_aws_partition_by_region(region: &str) -> String {
+    match region {
+        r if r.starts_with("us-gov-") => "aws-us-gov".to_string(),
+        r if r.starts_with("cn-") => "aws-cn".to_string(),
+        _ => "aws".to_string(),
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use super::*;

bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs

@@ -1,17 +1,14 @@
+use crate::config::get_aws_partition_by_region;
+use crate::lifecycle::invocation::{
+    processor::MS_TO_NS,
+    triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},
+};
 use datadog_trace_protobuf::pb::Span;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
 use tracing::debug;
 
-use crate::lifecycle::invocation::{
-    processor::MS_TO_NS,
-    triggers::{
-        get_aws_partition_by_region, lowercase_key, ServiceNameResolver, Trigger,
-        FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
-    },
-};
-
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct APIGatewayHttpEvent {
     #[serde(rename = "routeKey")]

bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs

@@ -1,17 +1,14 @@
+use crate::config::get_aws_partition_by_region;
+use crate::lifecycle::invocation::{
+    processor::MS_TO_NS,
+    triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},
+};
 use datadog_trace_protobuf::pb::Span;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
 use tracing::debug;
 
-use crate::lifecycle::invocation::{
-    processor::MS_TO_NS,
-    triggers::{
-        get_aws_partition_by_region, lowercase_key, ServiceNameResolver, Trigger,
-        FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
-    },
-};
-
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct APIGatewayRestEvent {
     #[serde(deserialize_with = "lowercase_key")]

bottlecap/src/lifecycle/invocation/triggers/mod.rs

@@ -57,15 +57,6 @@ pub trait Trigger: ServiceNameResolver {
     }
 }
 
-#[must_use]
-pub fn get_aws_partition_by_region(region: &str) -> String {
-    match region {
-        r if r.starts_with("us-gov-") => "aws-us-gov".to_string(),
-        r if r.starts_with("cn-") => "aws-cn".to_string(),
-        _ => "aws".to_string(),
-    }
-}
-
 /// Serialize a `HashMap` with lowercase keys
 ///
 pub fn lowercase_key<'de, D, V>(deserializer: D) -> Result<HashMap<String, V>, D::Error>

bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs

@@ -1,19 +1,18 @@
-use datadog_trace_protobuf::pb::Span;
-use serde::{Deserialize, Serialize};
-use serde_json::Value;
-use std::collections::HashMap;
-use tracing::debug;
-
+use crate::config::get_aws_partition_by_region;
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{
         event_bridge_event::EventBridgeEvent,
-        get_aws_partition_by_region,
         sns_event::{SnsEntity, SnsRecord},
         ServiceNameResolver, Trigger, DATADOG_CARRIER_KEY, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
     },
 };
 use crate::traces::context::{Sampling, SpanContext};
+use datadog_trace_protobuf::pb::Span;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::HashMap;
+use tracing::debug;
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct SqsEvent {