[SVLS-6337] Set resource group correctly for Azure functions on flex consumption plans (#1241)

nit: fix WEBSITE_OWNER_NAME mispelling

Check DD_AZURE_RESOURCE_GROUP env variable when getting resource group for flex consumption functions

nit: Update panic message

ran cargo fmt

Changed logic order to check DD_AZURE_RESOURCE_GROUP first

Set resource group to unknown instead of panicking

nit: fix formatting

Use WEBSITE_SKU instead of WEBSITE_OWNER_NAME to check for flex consumption function

Add comment to ensure future developers know to update other repos if flex consumption rg detection logic changes

Merge branch 'main' into kathie.huang/fix-flex-rg-span-attribute

Merge branch 'main' into kathie.huang/fix-flex-rg-span-attribute

Author: kathiehuang

ddcommon/src/azure_app_services.rs

@@ -5,7 +5,7 @@ use regex::Regex;
 use std::env;
 use std::sync::LazyLock;
 
-const WEBSITE_ONWER_NAME: &str = "WEBSITE_OWNER_NAME";
+const WEBSITE_OWNER_NAME: &str = "WEBSITE_OWNER_NAME";
 const WEBSITE_SITE_NAME: &str = "WEBSITE_SITE_NAME";
 const WEBSITE_RESOURCE_GROUP: &str = "WEBSITE_RESOURCE_GROUP";
 const SITE_EXTENSION_VERSION: &str = "DD_AAS_DOTNET_EXTENSION_VERSION";
@@ -16,6 +16,8 @@ const SERVICE_CONTEXT: &str = "DD_AZURE_APP_SERVICES";
 const FUNCTIONS_WORKER_RUNTIME: &str = "FUNCTIONS_WORKER_RUNTIME";
 const FUNCTIONS_WORKER_RUNTIME_VERSION: &str = "FUNCTIONS_WORKER_RUNTIME_VERSION";
 const FUNCTIONS_EXTENSION_VERSION: &str = "FUNCTIONS_EXTENSION_VERSION";
+const DD_AZURE_RESOURCE_GROUP: &str = "DD_AZURE_RESOURCE_GROUP";
+const WEBSITE_SKU: &str = "WEBSITE_SKU";
 
 const UNKNOWN_VALUE: &str = "unknown";
 
@@ -131,7 +133,7 @@ impl AzureMetadata {
 
     fn build_metadata<T: QueryEnv>(query: T) -> Option<Self> {
         let subscription_id =
-            AzureMetadata::extract_subscription_id(query.get_var(WEBSITE_ONWER_NAME));
+            AzureMetadata::extract_subscription_id(query.get_var(WEBSITE_OWNER_NAME));
         let site_name = query.get_var(WEBSITE_SITE_NAME);
 
         let (site_kind, site_type) = match AzureMetadata::get_azure_context(&query) {
@@ -140,8 +142,20 @@ impl AzureMetadata {
         };
 
         let resource_group = query
-            .get_var(WEBSITE_RESOURCE_GROUP)
-            .or_else(|| AzureMetadata::extract_resource_group(query.get_var(WEBSITE_ONWER_NAME)));
+            .get_var(DD_AZURE_RESOURCE_GROUP)
+            .or_else(|| query.get_var(WEBSITE_RESOURCE_GROUP))
+            .or_else(|| {
+                // Check if we're in flex consumption plan first
+                match query.get_var(WEBSITE_SKU).as_deref() {
+                    Some("FlexConsumption") => None,
+                    /* Flex Consumption plans need the `DD_AZURE_RESOURCE_GROUP` env var. If this
+                     * logic ever changes, update the logic in
+                     * `serverless-components/src/datadog-trace-agent` and the serverless compat
+                     * layers accordingly. */
+                    _ => AzureMetadata::extract_resource_group(query.get_var(WEBSITE_OWNER_NAME)),
+                }
+            });
+
         let resource_id = AzureMetadata::build_resource_id(
             subscription_id.as_ref(),
             site_name.as_ref(),
@@ -262,7 +276,7 @@ mod tests {
 
     use indexmap::IndexMap;
 
-    use crate::azure_app_services::{QueryEnv, WEBSITE_ONWER_NAME};
+    use crate::azure_app_services::{QueryEnv, WEBSITE_OWNER_NAME};
 
     use super::*;
 
@@ -353,7 +367,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_without_plus_sign() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "foo"), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "foo"), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -364,7 +378,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_with_plus_sign() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "foo+bar"), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "foo+bar"), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -374,7 +388,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_with_empty_string() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, ""), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, ""), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -383,7 +397,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_with_only_whitespaces() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "    "), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "    "), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -392,7 +406,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_with_only_plus_sign() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "+"), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "+"), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -401,7 +415,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_with_whitespaces_separated_by_plus() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "   + "), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "   + "), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -410,7 +424,7 @@ mod tests {
 
     #[test]
     fn test_extract_subscription_plus_sign_and_other_string() {
-        let mocked_env = MockEnv::new(&[(WEBSITE_ONWER_NAME, "+other"), (SERVICE_CONTEXT, "1")]);
+        let mocked_env = MockEnv::new(&[(WEBSITE_OWNER_NAME, "+other"), (SERVICE_CONTEXT, "1")]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
 
@@ -421,7 +435,7 @@ mod tests {
     fn test_extract_resource_group_pattern_match_linux() {
         let mocked_env = MockEnv::new(&[
             (
-                WEBSITE_ONWER_NAME,
+                WEBSITE_OWNER_NAME,
                 "00000000-0000-0000-0000-000000000000+test-rg-EastUSwebspace-Linux",
             ),
             ("FUNCTIONS_WORKER_RUNTIME", "node"),
@@ -439,7 +453,7 @@ mod tests {
     fn test_extract_resource_group_pattern_match_windows() {
         let mocked_env = MockEnv::new(&[
             (
-                WEBSITE_ONWER_NAME,
+                WEBSITE_OWNER_NAME,
                 "00000000-0000-0000-0000-000000000000+test-rg-EastUSwebspace",
             ),
             ("FUNCTIONS_WORKER_RUNTIME", "node"),
@@ -456,7 +470,7 @@ mod tests {
     #[test]
     fn test_extract_resource_group_no_pattern_match() {
         let mocked_env = MockEnv::new(&[
-            (WEBSITE_ONWER_NAME, "foo"),
+            (WEBSITE_OWNER_NAME, "foo"),
             (FUNCTIONS_WORKER_RUNTIME, "node"),
             (FUNCTIONS_EXTENSION_VERSION, "~4"),
         ]);
@@ -471,10 +485,11 @@ mod tests {
         let mocked_env = MockEnv::new(&[
             (WEBSITE_RESOURCE_GROUP, "test-rg-env-var"),
             (
-                WEBSITE_ONWER_NAME,
+                WEBSITE_OWNER_NAME,
                 "00000000-0000-0000-0000-000000000000+test-rg-EastUSwebspace-Linux",
             ),
             (SERVICE_CONTEXT, "1"),
+            (WEBSITE_SKU, "ElasticPremium"),
         ]);
 
         let metadata = AzureMetadata::new(mocked_env).unwrap();
@@ -484,10 +499,66 @@ mod tests {
         assert_eq!(metadata.get_resource_group(), expected_resource_group);
     }
 
+    #[test]
+    fn test_flex_consumption_resource_group_is_none_without_dd_azure_resource_group() {
+        let mocked_env = MockEnv::new(&[
+            (
+                WEBSITE_OWNER_NAME,
+                "00000000-0000-0000-0000-000000000000+flex-EastUSwebspace-Linux",
+            ),
+            (WEBSITE_SKU, "FlexConsumption"),
+            (SERVICE_CONTEXT, "1"),
+        ]);
+
+        let metadata = AzureMetadata::new(mocked_env).unwrap();
+
+        assert_eq!(metadata.get_resource_group(), UNKNOWN_VALUE);
+    }
+
+    #[test]
+    fn test_flex_consumption_uses_dd_azure_resource_group() {
+        let mocked_env = MockEnv::new(&[
+            (
+                WEBSITE_OWNER_NAME,
+                "00000000-0000-0000-0000-000000000000+flex-EastUSwebspace-Linux",
+            ),
+            (DD_AZURE_RESOURCE_GROUP, "test-flex-rg"),
+            (WEBSITE_SKU, "FlexConsumption"),
+            (SERVICE_CONTEXT, "1"),
+        ]);
+
+        let metadata = AzureMetadata::new(mocked_env).unwrap();
+
+        // Should use the DD_AZURE_RESOURCE_GROUP value instead of extracting from
+        // WEBSITE_OWNER_NAME
+        assert_eq!(metadata.get_resource_group(), "test-flex-rg");
+    }
+
+    #[test]
+    fn test_dd_azure_resource_group_has_highest_priority() {
+        let mocked_env = MockEnv::new(&[
+            (WEBSITE_RESOURCE_GROUP, "test-rg-env-var"),
+            (
+                WEBSITE_OWNER_NAME,
+                "00000000-0000-0000-0000-000000000000+test-rg-EastUSwebspace-Linux",
+            ),
+            (DD_AZURE_RESOURCE_GROUP, "dd-azure-rg-override"),
+            (SERVICE_CONTEXT, "1"),
+        ]);
+
+        let metadata = AzureMetadata::new(mocked_env).unwrap();
+
+        // DD_AZURE_RESOURCE_GROUP should have highest priority over WEBSITE_RESOURCE_GROUP and
+        // WEBSITE_OWNER_NAME
+        let expected_resource_group = "dd-azure-rg-override";
+
+        assert_eq!(metadata.get_resource_group(), expected_resource_group);
+    }
+
     #[test]
     fn test_build_resource_id() {
         let mocked_env = MockEnv::new(&[
-            (WEBSITE_ONWER_NAME, "foo"),
+            (WEBSITE_OWNER_NAME, "foo"),
             (WEBSITE_SITE_NAME, "my_website"),
             (WEBSITE_RESOURCE_GROUP, "resource_group"),
             (SERVICE_CONTEXT, "1"),
@@ -517,7 +588,7 @@ mod tests {
     #[test]
     fn test_build_resource_id_with_missing_site_name() {
         let mocked_env = MockEnv::new(&[
-            (WEBSITE_ONWER_NAME, "foo"),
+            (WEBSITE_OWNER_NAME, "foo"),
             (WEBSITE_RESOURCE_GROUP, "resource_group"),
             (SERVICE_CONTEXT, "1"),
         ]);
@@ -530,7 +601,7 @@ mod tests {
     #[test]
     fn test_build_resource_id_with_missing_resource_group() {
         let mocked_env = MockEnv::new(&[
-            (WEBSITE_ONWER_NAME, "foo"),
+            (WEBSITE_OWNER_NAME, "foo"),
             (WEBSITE_SITE_NAME, "my_website"),
             (SERVICE_CONTEXT, "1"),
         ]);