[azure][feat] Connect CosmosDB resources to the location instead of subscription (#2214)

Author: 1101-1

plugins/azure/fix_plugin_azure/collector.py

@@ -2,7 +2,7 @@
 from abc import abstractmethod
 from concurrent.futures import ThreadPoolExecutor, Future
 from datetime import datetime, timezone
-from typing import Any, Optional, Type, List, Dict
+from typing import Any, Optional, Type, List, Dict, Set
 
 from azure.core.utils import CaseInsensitiveDict
 
@@ -215,7 +215,11 @@ def after_collect(self) -> None:
 class AzureSubscriptionCollector(MicrosoftBaseCollector):
     def locations(self, builder: GraphBuilder) -> Dict[str, BaseRegion]:
         locations = AzureLocation.collect_resources(builder)
-        return CaseInsensitiveDict({loc.safe_name: loc for loc in locations})  # type: ignore
+        # сreate a location lookup map with lowercase name and display name of the locations
+        locations_map = CaseInsensitiveDict()
+        locations_map.update({loc.safe_name: loc for loc in locations})
+        locations_map.update({loc.display_name or loc.safe_name: loc for loc in locations})
+        return locations_map  # type: ignore
 
     def collect_with(self, builder: GraphBuilder, locations: Dict[str, BaseRegion]) -> None:
         # add deferred edge to organization
@@ -224,8 +228,11 @@ def collect_with(self, builder: GraphBuilder, locations: Dict[str, BaseRegion])
         regional_resources = [r for r in subscription_resources if resource_with_params(r, "location")]
         global_resources = list(set(subscription_resources) - set(regional_resources))
         self.collect_resource_list("subscription", builder, global_resources)
+        processed_locations: Set[str] = set()
         for location in locations.values():
-            self.collect_resource_list(location.safe_name, builder.with_location(location), regional_resources)
+            if location.safe_name not in processed_locations:
+                self.collect_resource_list(location.safe_name, builder.with_location(location), regional_resources)
+                processed_locations.add(location.safe_name)
 
     def remove_unused(self) -> None:
         remove_nodes = []
@@ -264,7 +271,7 @@ def remove_usage_zero_value() -> None:
         rm_nodes(AzureStorageSku, AzureLocation)
         rm_nodes(AzureMysqlServerType, AzureSubscription)
         rm_nodes(AzurePostgresqlServerType, AzureSubscription)
-        rm_nodes(AzureCosmosDBLocation, AzureSubscription, check_pred=False)
+        rm_nodes(AzureCosmosDBLocation, AzureLocation, check_pred=False)
         rm_nodes(AzureLocation, check_pred=False)
         remove_usage_zero_value()
 

plugins/azure/fix_plugin_azure/resource/base.py

@@ -857,6 +857,11 @@ def add_node(self, node: MicrosoftResourceType, source: Optional[Json] = None) -
             if location := self.location_lookup.get(source_location):
                 node._region = location
                 last_edge_key = self.add_edge(location, node=node)
+        elif (node_location := getattr(node, "location", None)) is not None:
+            # reference the location node if available in resource property
+            if location := self.location_lookup.get(node_location):
+                node._region = location
+                last_edge_key = self.add_edge(location, node=node)
         if source and "locations" in source:
             for loc in source["locations"]:
                 # reference the location node if available

plugins/azure/fix_plugin_azure/resource/cosmosdb.py

@@ -32,6 +32,16 @@
 log = logging.getLogger("fix.plugins.azure")
 
 
+class CosmosDBLocationSetter:
+    def __init__(self) -> None:
+        self.location: Optional[str] = None
+
+    def pre_process(self, graph_builder: GraphBuilder, source: Json) -> None:
+        if isinstance(self, MicrosoftResource):
+            if location := self.extract_part("locations"):
+                self.location = location
+
+
 @define(eq=False, slots=False)
 class AzureManagedCassandraReaperStatus:
     kind: ClassVar[str] = "azure_managed_cassandra_reaper_status"
@@ -1550,7 +1560,7 @@ class AzureRestorableLocationResource:
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableAccount(MicrosoftResource):
+class AzureCosmosDBRestorableAccount(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_account"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Account"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2032,7 +2042,7 @@ def collect_usages(cls, account_id: str, raw: List[Json], builder: GraphBuilder)
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBLocation(MicrosoftResource):
+class AzureCosmosDBLocation(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_location"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Location"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2163,7 +2173,7 @@ class AzureRestorableDatabase:
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableGremlinDatabase(MicrosoftResource):
+class AzureCosmosDBRestorableGremlinDatabase(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_gremlin_database"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Gremlin Database"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2211,7 +2221,7 @@ def collect_restorable_graphs() -> None:
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableGremlinGraph(MicrosoftResource):
+class AzureCosmosDBRestorableGremlinGraph(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_gremlin_graph"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Gremlin Graph"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2229,7 +2239,7 @@ class AzureCosmosDBRestorableGremlinGraph(MicrosoftResource):
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableMongoDBCollection(MicrosoftResource):
+class AzureCosmosDBRestorableMongoDBCollection(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_mongo_db_collection"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Mongo DB Collection"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2247,7 +2257,7 @@ class AzureCosmosDBRestorableMongoDBCollection(MicrosoftResource):
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableMongoDBDatabase(MicrosoftResource):
+class AzureCosmosDBRestorableMongoDBDatabase(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_mongo_db_database"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Mongo DB Database"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2295,7 +2305,7 @@ def collect_restorable_collections() -> None:
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableSqlContainer(MicrosoftResource):
+class AzureCosmosDBRestorableSqlContainer(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_sql_container"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable SQL Container"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2344,7 +2354,7 @@ class AzureExtendedPropertiesSqlDatabase(AzureSqlDatabaseResource):
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableSqlDatabase(MicrosoftResource):
+class AzureCosmosDBRestorableSqlDatabase(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_sql_database"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable SQL Database"
     _kind_service: ClassVar[Optional[str]] = service_name
@@ -2406,7 +2416,7 @@ def collect_restorable_containers() -> None:
 
 
 @define(eq=False, slots=False)
-class AzureCosmosDBRestorableTable(MicrosoftResource):
+class AzureCosmosDBRestorableTable(CosmosDBLocationSetter, MicrosoftResource):
     kind: ClassVar[str] = "azure_cosmos_db_restorable_table"
     _kind_display: ClassVar[str] = "Azure Cosmos DB Restorable Table"
     _kind_service: ClassVar[Optional[str]] = service_name

plugins/azure/test/collector_test.py

@@ -48,8 +48,8 @@ def test_collect(
         config, Cloud(id="azure"), azure_subscription, credentials, core_feedback
     )
     subscription_collector.collect()
-    assert len(subscription_collector.graph.nodes) == 645
-    assert len(subscription_collector.graph.edges) == 1034
+    assert len(subscription_collector.graph.nodes) == 654
+    assert len(subscription_collector.graph.edges) == 1043
 
     graph_collector = MicrosoftGraphOrganizationCollector(
         config, Cloud(id="azure"), MicrosoftGraphOrganization(id="test", name="test"), credentials, core_feedback