Skip to content

Commit

Permalink
[feat][gcp] Improve the way of collection for DiskTypes and MachineTy…
Browse files Browse the repository at this point in the history
…pes (#2284)
  • Loading branch information
1101-1 authored Dec 4, 2024
1 parent 73efbf3 commit 54024e5
Show file tree
Hide file tree
Showing 6 changed files with 183 additions and 115 deletions.
1 change: 1 addition & 0 deletions plugins/gcp/fix_plugin_gcp/gcp_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from fixlib.types import Json

InternalZoneProp = "_zone"
ZoneProp = "zone"
RegionProp = "region"

# Store the discovery function as separate variable.
Expand Down
57 changes: 36 additions & 21 deletions plugins/gcp/fix_plugin_gcp/resources/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from googleapiclient.errors import HttpError

from fix_plugin_gcp.config import GcpConfig
from fix_plugin_gcp.gcp_client import GcpClient, GcpApiSpec, InternalZoneProp, RegionProp
from fix_plugin_gcp.gcp_client import GcpClient, GcpApiSpec, InternalZoneProp, ZoneProp, RegionProp
from fix_plugin_gcp.utils import Credentials
from fixlib.baseresources import (
BaseResource,
Expand Down Expand Up @@ -187,27 +187,26 @@ def nodes(
result.append(n)
return result

def add_node(self, node: GcpResourceType, source: Optional[Json] = None) -> Optional[GcpResourceType]:
def add_node(self, node: GcpResourceType, source: Optional[Json] = None) -> GcpResourceType:
log.debug(f"{self.name}: add node {node}")
node._cloud = self.cloud
node._account = self.project

if self._standard_edges(node, source):
with self.graph_nodes_access:
self.graph.add_node(node, source=source or {})
return node
return None
self.add_region_to_node(node, source)
with self.graph_nodes_access:
self.graph.add_node(node, source=source or {})
return node

def _standard_edges(self, node: GcpResourceType, source: Optional[Json] = None) -> bool:
def add_region_to_node(self, node: GcpResourceType, source: Optional[Json] = None) -> None:
if isinstance(node, GcpRegion):
self.add_edge(node, node=self.project, reverse=True)
return True
return
if node._zone:
self.add_edge(node, node=node._zone, reverse=True)
return True
return
if node._region:
self.add_edge(node, node=node._region, reverse=True)
return True
return

parts = node.id.split("/", maxsplit=4)
if len(parts) > 3 and parts[0] == "projects":
Expand All @@ -218,45 +217,61 @@ def _standard_edges(self, node: GcpResourceType, source: Optional[Json] = None)
node._zone = zone
node._region = self.region_by_zone_name.get(zone.id)
self.add_edge(zone, node=node)
return True
return

# Then check for region
if region := self.region_by_name.get(location_name):
node._region = region
self.add_edge(region, node=node)
return True
return

if source is not None:
if ZoneProp in source:
zone_name = source[ZoneProp].rsplit("/", 1)[-1]
if zone := self.zone_by_name.get(zone_name):
node._zone = zone
node._region = self.region_by_zone_name[zone_name]
self.add_edge(node, node=zone, reverse=True)
return
else:
log.debug(
"Zone property '%s' found in the source but no corresponding zone object is available to associate with the node.",
zone_name,
)

if InternalZoneProp in source:
if zone := self.zone_by_name.get(source[InternalZoneProp]):
node._zone = zone
node._region = self.region_by_zone_name[source[InternalZoneProp]]
self.add_edge(node, node=zone, reverse=True)
return True
return
else:
log.debug(f"Zone {source[InternalZoneProp]} not found for node: {node}. Ignore resource.")
return False
log.debug(
"Internal zone property '%s' exists in the source but no corresponding zone object is available to associate with the node.",
source[InternalZoneProp],
)

if RegionProp in source:
region_name = source[RegionProp].rsplit("/", 1)[-1]
if region := self.region_by_name.get(region_name):
node._region = region
self.add_edge(node, node=region, reverse=True)
return True
return
else:
log.debug(f"Region {region_name} not found for node: {node}. Ignore resource.")
return False
log.debug(
"Region property '%s' found in the source but no corresponding region object is available to associate with the node.",
region_name,
)

# Fallback to GraphBuilder region, i.e. regional collection
if self.region is not None:
node._region = self.region
self.add_edge(node, node=self.region, reverse=True)
return True
return

# Fallback to global region
node._region = self.fallback_global_region
self.add_edge(node, node=self.fallback_global_region, reverse=True)
return True

def add_edge(
self,
Expand Down
20 changes: 12 additions & 8 deletions plugins/gcp/fix_plugin_gcp/resources/billing.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from datetime import datetime
from typing import ClassVar, Dict, Optional, List, Type, cast, Any
from typing import ClassVar, Dict, Optional, List, Type, Any

from attr import define, field

Expand Down Expand Up @@ -146,18 +146,22 @@ class GcpService(GcpResource):
display_name: Optional[str] = field(default=None)

@classmethod
def collect(cls: Type[GcpResource], raw: List[Json], builder: GraphBuilder) -> List[GcpResource]:
def collect(cls, raw: List[Json], builder: GraphBuilder) -> List[GcpResource]:
# Additional behavior: iterate over list of collected GcpService and for each:
# - collect related GcpSku
result: List[GcpResource] = super().collect(raw, builder) # type: ignore
SERVICES_COLLECT_LIST = [
"Compute Engine",
]
service_names = [
service.name for service in cast(List[GcpService], result) if service.display_name in SERVICES_COLLECT_LIST
]
for service_name in service_names:
builder.submit_work(GcpSku.collect_resources, builder, parent=service_name)
service_names: List[str] = []
services = []
for service in raw:
if service.get("displayName") in SERVICES_COLLECT_LIST:
service_names.append(str(service["name"]))
services.append(service)
result = super().collect(services, builder)

for s_name in service_names:
builder.submit_work(GcpSku.collect_resources, builder, parent=s_name)

return result

Expand Down
Loading

0 comments on commit 54024e5

Please sign in to comment.