[feat] Allow filtering descendant summary

fixbackend/inventory/inventory_router.py

@@ -38,6 +38,7 @@
  TimeseriesRequest,
  Scatters,
  KindUsage,
+ KindUsageRequest,
 )
 from fixbackend.streaming_response import streaming_response, StreamOnSuccessResponse
 from fixbackend.workspaces.dependencies import UserWorkspaceDependency
@@ -413,8 +414,16 @@ async def timeseries(graph_db: CurrentGraphDbDependency, ts: TimeseriesRequest)
  aggregation=ts.aggregation,
  )
 
- @router.get("/descendant/summary", tags=["search"])
- async def descendant_summary_account(graph_db: CurrentGraphDbDependency) -> Dict[str, KindUsage]:
- return await inventory().descendant_summary(graph_db)
+ @router.post("/descendant/summary", tags=["search"])
+ async def descendant_summary_account(
+ graph_db: CurrentGraphDbDependency, request: KindUsageRequest
+ ) -> Dict[str, KindUsage]:
+ return await inventory().descendant_summary(
+ graph_db,
+ request.cloud_ids or [],
+ request.account_ids or [],
+ request.region_ids or [],
+ request.kinds or [],
+ )
 
  return router

fixbackend/inventory/inventory_schemas.py

@@ -272,6 +272,13 @@ class TimeseriesRequest(BaseModel):
  aggregation: Optional[str] = None
 
 
+class KindUsageRequest(BaseModel):
+ cloud_ids: Optional[List[str]] = Field(default=None, description="The cloud ids to filter by.")
+ account_ids: Optional[List[str]] = Field(default=None, description="The account ids to filter by.")
+ region_ids: Optional[List[str]] = Field(default=None, description="The region ids to filter by.")
+ kinds: Optional[List[str]] = Field(default=None, description="The kinds to filter by.")
+
+
 class KindUsage(BaseModel):
  accounts: int = Field(default=0, description="The number of accounts using this kind.")
  regions: int = Field(default=0, description="The number of regions using this kind.")

fixbackend/inventory/inventory_service.py

@@ -19,6 +19,7 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
+from __future__ import annotations
 import asyncio
 import json
 import logging
@@ -43,7 +44,8 @@
 
 from arq import func
 from arq.connections import RedisSettings
-from attr import frozen
+from attr import frozen, evolve
+from attrs import define, field
 from fixcloudutils.redis.cache import RedisCache
 from fixcloudutils.redis.worker_queue import WorkDispatcher, WorkerInstance
 from fixcloudutils.service import Service
@@ -155,6 +157,24 @@ class InventorySummary:
  buckets_size_bytes_progress: Tuple[int, int]
 
 
+@define
+class AccountRegionSummary:
+ account_id: str
+ cloud_id: str
+ count: Dict[str, int] = field(factory=lambda: defaultdict(int))
+ regions: Dict[str, Dict[str, int]] = field(factory=lambda: defaultdict(lambda: defaultdict(int)))
+
+ def for_regions(self, regions: List[str]) -> AccountRegionSummary:
+ return evolve(self, regions={r: count for r, count in self.regions.items() if r in regions})
+
+ def for_kinds(self, kinds: List[str]) -> AccountRegionSummary:
+ return evolve(
+ self,
+ count={k: count for k, count in self.count.items() if k in kinds},
+ regions={r: {k: v for k, v in kc.items() if k in kinds} for r, kc in self.regions.items()},
+ )
+
+
 class InventoryService(Service):
  def __init__(
  self,
@@ -861,29 +881,82 @@ async def overall_score() -> Tuple[int, int]:
 
  return await self.cache.call(compute_inventory_info, key=str(dba.workspace_id))(duration)
 
- async def descendant_summary(self, dba: GraphDatabaseAccess) -> Dict[str, KindUsage]:
+ async def descendant_summary(
+ self,
+ dba: GraphDatabaseAccess,
+ cloud_ids: List[str],
+ account_ids: List[str],
+ region_ids: List[str],
+ kinds: List[str],
+ ) -> Dict[str, KindUsage]:
+
+ async def account_region_summary() -> Dict[str, AccountRegionSummary]:
+ account_usage: Dict[str, AccountRegionSummary] = {}
+
+ def get_account_usage(account_id: str, cloud_id: str) -> AccountRegionSummary:
+ au = account_usage.get(account_id)
+ if not au:
+ au = AccountRegionSummary(account_id=account_id, cloud_id=cloud_id)
+ account_usage[account_id] = au
+ return au
 
- async def compute_descendant_summary() -> Dict[str, KindUsage]:
- account_usage: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
- region_usage: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
  async with self.client.search(dba, "is(account) and /metadata.descendant_count>0") as response:
  async for acc in response:
- if account_id := value_in_path(acc, "reported.id"):
+ if (cloud_id := value_in_path(acc, "ancestors.cloud.reported.id")) and (
+ account_id := value_in_path(acc, "reported.id")
+ ):
  descendant_summary = value_in_path(acc, "metadata.descendant_summary") or {}
+ account_summary = get_account_usage(account_id, cloud_id)
  for descendant_kind, count in descendant_summary.items():
- account_usage[descendant_kind][account_id] += count
+ account_summary.count[descendant_kind] += count
+
  async with self.client.search(dba, "is(region) and /metadata.descendant_count>0") as response:
  async for acc in response:
- if region_id := value_in_path(acc, "reported.id"):
+ if (
+ (cloud_id := value_in_path(acc, "ancestors.cloud.reported.id"))
+ and (account_id := value_in_path(acc, "ancestors.account.reported.id"))
+ and (region_id := value_in_path(acc, "reported.id"))
+ ):
+ account_summary = get_account_usage(account_id, cloud_id)
  descendant_summary = value_in_path(acc, "metadata.descendant_summary") or {}
  for descendant_kind, count in descendant_summary.items():
- region_usage[descendant_kind][region_id] += count
- kind_usage: Dict[str, KindUsage] = defaultdict(KindUsage)
- for kind, accounts in account_usage.items():
- kind_usage[kind].accounts = len(accounts)
- kind_usage[kind].resources = sum(accounts.values())
+ account_summary.regions[region_id][descendant_kind] += count
+ return account_usage
+
+ summary = await self.cache.call(account_region_summary, key=str(dba.workspace_id))()
+ if cloud_ids:
+ summary = {k: v for k, v in summary.items() if v.cloud_id in cloud_ids}
+ if account_ids:
+ summary = {k: v for k, v in summary.items() if v.account_id in account_ids}
+ if region_ids:
+ summary = {k: v.for_regions(region_ids) for k, v in summary.items()}
+ if kinds:
+ summary = {k: v.for_kinds(kinds) for k, v in summary.items()}
+
+ account_usage: Dict[str, int] = defaultdict(int)
+ region_usage: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
+ account_region_usage: Dict[str, Set[str]] = defaultdict(set)
+ for account in summary.values():
+ for kind, count in account.count.items():
+ account_usage[kind] += count
+ for region, region_kinds in account.regions.items():
+ for kind, count in region_kinds.items():
+ account_region_usage[kind].add(account.account_id)
+ region_usage[kind][region] += count
+ kind_usage: Dict[str, KindUsage] = defaultdict(KindUsage)
+ if region_ids:
  for kind, regions in region_usage.items():
- kind_usage[kind].regions = len(regions)
- return kind_usage
-
- return await self.cache.call(compute_descendant_summary, key=str(dba.workspace_id))()
+ kd = kind_usage[kind]
+ kd.accounts = len(account_region_usage[kind])
+ kd.regions = len(regions)
+ kd.resources = sum(regions.values())
+ else:
+ for kind, count in account_usage.items():
+ kd = kind_usage[kind]
+ kd.accounts += 1
+ kd.regions = len(region_usage[kind])
+ kd.resources += count
+
+ for kind, regions in region_usage.items():
+ kind_usage[kind].regions = len(regions)
+ return kind_usage