"""Private datastore adapter for deterministic joiner demand forecasts. Raw identities, organisation labels, product names, cloud roles, and office names never leave this module. The returned forecast contains only stable package-local references; ``reference_map`` is retained by Helix for an authorised drill-down and must not be sent to a model. """ from __future__ import annotations from collections import defaultdict from dataclasses import asdict, dataclass from datetime import UTC, datetime, timedelta from typing import Any from uuid import UUID from zoneinfo import ZoneInfo from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.domains.asset.hardware_prep import COMPUTER_CATEGORIES, RESERVABLE_STATUSES from app.domains.asset.location_admin_sync import ( resolve_asset_location_ids_bulk, resolve_person_location_ids_bulk, ) from app.domains.asset.models import AssetItem, SoftwareLicensePlan, SoftwareSeatAllocation from app.domains.cmdb.models import GRANT_STATUS_ACTIVE, CmdbCloudAccessGrant from app.domains.onboarding.models import PROC_OPEN_STATUSES, OnboardingProcess from app.domains.vpn_application.models import VPN_APP_STATUS_APPROVED, VpnApplication from app.identity.person_models import IdentityOrgMembership, IdentityPerson from app.intelligence.demand_forecast import ( Candidate, DemandForecast, Peer, ResourceKey, SupplyPosition, build_demand_forecast, ) from app.intelligence.privacy import EvidenceObject, assert_llm_safe, build_reference_map _APP_TZ = ZoneInfo("Asia/Shanghai") @dataclass(frozen=True) class DemandForecastFactPackage: """Forecast plus the local-only evidence map and source quality counters.""" forecast: DemandForecast reference_map: dict[str, Any] excluded_counts: dict[str, int] data_as_of: datetime | None def aggregate_facts(self) -> dict[str, Any]: payload = { "forecast": asdict(self.forecast), "excluded_counts": dict(sorted(self.excluded_counts.items())), } assert_llm_safe(payload) return payload def _utc(value: datetime) -> datetime: return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC) def _reverse_reference_map( objects: list[EvidenceObject], ) -> tuple[dict[tuple[str, str], str], dict[str, Any]]: _, mapping = build_reference_map(objects) reverse = { (str(value["kind"]), str(value["object_id"])): ref for ref, value in mapping.items() } return reverse, mapping async def collect_demand_forecast_facts( db: AsyncSession, *, now: datetime | None = None, ) -> DemandForecastFactPackage: """Collect privacy-safe 30/60/90-day resource-demand facts. Historical prevalence establishes a default requirement. A resource already assigned or reserved for a candidate is an explicit commitment, so it remains demand even when prevalence is below 60%; it does not waive the five-person minimum required to infer a cohort. """ observed_at = _utc(now or datetime.now(UTC)) today = observed_at.astimezone(_APP_TZ).date() window_end = today + timedelta(days=90) excluded: defaultdict[str, int] = defaultdict(int) open_processes = list((await db.execute(select(OnboardingProcess).where( OnboardingProcess.status.in_(PROC_OPEN_STATUSES), ))).scalars()) processes = [ row for row in open_processes if today < row.hire_date <= window_end and not row.backfilled ] people = list((await db.execute(select(IdentityPerson).where( IdentityPerson.deleted_at.is_(None), IdentityPerson.employment_status.in_(("active", "pre_hire")), ))).scalars()) people_by_id = {row.id: row for row in people} candidate_processes = [row for row in processes if row.person_id in people_by_id] excluded["candidate_without_eligible_person"] += len(processes) - len(candidate_processes) memberships = list((await db.execute(select(IdentityOrgMembership).where( IdentityOrgMembership.deleted_at.is_(None), IdentityOrgMembership.is_primary.is_(True), IdentityOrgMembership.membership_status == "active", ).order_by(IdentityOrgMembership.person_id, IdentityOrgMembership.id))).scalars()) department_by_person: dict[UUID, str] = {} for row in memberships: department_by_person.setdefault(row.person_id, row.org_node_source_id) person_ids = set(people_by_id) location_by_person = await resolve_person_location_ids_bulk(db, person_ids) candidate_person_ids = {row.person_id for row in candidate_processes} open_onboarding_person_ids = {row.person_id for row in open_processes} eligible_person_ids: set[UUID] = set() for person_id, person in people_by_id.items(): complete = bool( location_by_person.get(person_id) and department_by_person.get(person_id) and person.job_title ) if complete: eligible_person_ids.add(person_id) elif person_id in candidate_person_ids: excluded["candidate_incomplete_cohort"] += 1 elif person.employment_status == "active": excluded["peer_incomplete_cohort"] += 1 candidate_processes = [ row for row in candidate_processes if row.person_id in eligible_person_ids ] candidate_person_ids = {row.person_id for row in candidate_processes} peer_person_ids = { row.id for row in people if row.employment_status == "active" and row.id in eligible_person_ids and row.id not in open_onboarding_person_ids } assets = list((await db.execute(select(AssetItem).where( AssetItem.deleted_at.is_(None), AssetItem.disposed_at.is_(None), AssetItem.status != "disposed", AssetItem.category.in_(COMPUTER_CATEGORIES), ))).scalars()) asset_locations = await resolve_asset_location_ids_bulk(db, assets) plans = list((await db.execute(select(SoftwareLicensePlan).where( SoftwareLicensePlan.status == "active", ))).scalars()) allocations = list((await db.execute(select(SoftwareSeatAllocation).where( SoftwareSeatAllocation.status == "active", ))).scalars()) vpn_rows = list((await db.execute(select(VpnApplication).where( VpnApplication.status == VPN_APP_STATUS_APPROVED, ))).scalars()) cloud_rows = list((await db.execute(select(CmdbCloudAccessGrant).where( CmdbCloudAccessGrant.status == GRANT_STATUS_ACTIVE, CmdbCloudAccessGrant.person_id.is_not(None), ))).scalars()) cloud_rows = [ row for row in cloud_rows if row.expires_at is None or _utc(row.expires_at) > observed_at ] objects: list[EvidenceObject] = [] for process in candidate_processes: person = people_by_id[process.person_id] objects.extend(( EvidenceObject("candidate", str(process.id)), EvidenceObject("location", str(location_by_person[person.id])), EvidenceObject("department", department_by_person[person.id]), EvidenceObject("job", person.job_title or ""), )) for person_id in peer_person_ids: person = people_by_id[person_id] objects.extend(( EvidenceObject("peer", str(person_id)), EvidenceObject("location", str(location_by_person[person_id])), EvidenceObject("department", department_by_person[person_id]), EvidenceObject("job", person.job_title or ""), )) for asset in assets: if asset_locations.get(asset.id): objects.append(EvidenceObject("hardware", asset.category)) objects.append(EvidenceObject("location", str(asset_locations[asset.id]))) for plan in plans: objects.append(EvidenceObject("software", str(plan.id))) for vpn in vpn_rows: objects.append(EvidenceObject("access", f"vpn:{vpn.vpn_node}")) for grant in cloud_rows: objects.append(EvidenceObject("access", f"cloud:{grant.scope_id}:{grant.role_key}")) refs, reference_map = _reverse_reference_map(objects) def ref(kind: str, object_id: object) -> str: return refs[(kind, str(object_id))] hardware_resources = { category: ResourceKey("hardware", ref("hardware", category), "location") for category in sorted({row.category for row in assets if asset_locations.get(row.id)}) } software_resources = { row.id: ResourceKey("software", ref("software", row.id), "global") for row in plans } vpn_resources = { row.vpn_node: ResourceKey("access", ref("access", f"vpn:{row.vpn_node}"), "global") for row in vpn_rows } cloud_resources = { (row.scope_id, row.role_key): ResourceKey( "access", ref("access", f"cloud:{row.scope_id}:{row.role_key}"), "global", ) for row in cloud_rows } resources_by_person: defaultdict[UUID, set[ResourceKey]] = defaultdict(set) for asset in assets: if asset.assigned_to in person_ids and asset_locations.get(asset.id): resources_by_person[asset.assigned_to].add(hardware_resources[asset.category]) active_plan_ids = set(software_resources) for allocation in allocations: if allocation.user_id in person_ids and allocation.plan_id in active_plan_ids: resources_by_person[allocation.user_id].add(software_resources[allocation.plan_id]) for vpn in vpn_rows: try: person_id = UUID(str(vpn.person_id)) except (AttributeError, TypeError, ValueError): excluded["invalid_vpn_person_reference"] += 1 continue if person_id in person_ids: resources_by_person[person_id].add(vpn_resources[vpn.vpn_node]) for grant in cloud_rows: try: person_id = UUID(str(grant.person_id or "")) except (AttributeError, TypeError, ValueError): excluded["invalid_cloud_person_reference"] += 1 continue if person_id in person_ids: resources_by_person[person_id].add( cloud_resources[(grant.scope_id, grant.role_key)], ) candidates = tuple(Candidate( ref=ref("candidate", process.id), hire_date=process.hire_date, location_key=ref("location", location_by_person[process.person_id]), department_key=ref("department", department_by_person[process.person_id]), job_key=ref("job", people_by_id[process.person_id].job_title or ""), employment_status=( "prehire" if people_by_id[process.person_id].employment_status == "pre_hire" else people_by_id[process.person_id].employment_status ), ) for process in sorted(candidate_processes, key=lambda row: str(row.id))) candidate_ref_by_person = { process.person_id: ref("candidate", process.id) for process in candidate_processes } peers = tuple(Peer( ref=ref("peer", person_id), location_key=ref("location", location_by_person[person_id]), department_key=ref("department", department_by_person[person_id]), job_key=ref("job", people_by_id[person_id].job_title or ""), resources=frozenset(resources_by_person[person_id]), ) for person_id in sorted(peer_person_ids, key=str)) supply_rows: list[SupplyPosition] = [] stock: defaultdict[tuple[ResourceKey, str], int] = defaultdict(int) reserved: defaultdict[tuple[ResourceKey, str], set[str]] = defaultdict(set) for asset in assets: location_id = asset_locations.get(asset.id) if location_id is None: continue resource = hardware_resources[asset.category] scope = ref("location", location_id) key = (resource, scope) if asset.status in RESERVABLE_STATUSES: stock[key] += 1 reserved_for = asset.reserved_for_person_id candidate_ref = candidate_ref_by_person.get(reserved_for) if reserved_for else None if candidate_ref and asset.status in RESERVABLE_STATUSES: reserved[key].add(candidate_ref) assigned_to = asset.assigned_to assigned_ref = candidate_ref_by_person.get(assigned_to) if assigned_to else None if assigned_ref: # An assigned device is a confirmed unit outside stock. Add it to # both sides so general stock is unchanged and this hire is covered. stock[key] += 1 reserved[key].add(assigned_ref) for key in sorted(set(stock) | set(reserved), key=lambda row: (row[0], row[1])): commitments = tuple(sorted(reserved[key])) supply_rows.append(SupplyPosition( resource=key[0], scope_key=key[1], available=stock[key], reserved=len(commitments), committed_candidate_refs=commitments, )) allocations_by_plan: defaultdict[UUID, list[SoftwareSeatAllocation]] = defaultdict(list) for allocation in allocations: if allocation.plan_id in active_plan_ids: allocations_by_plan[allocation.plan_id].append(allocation) for plan in sorted(plans, key=lambda row: str(row.id)): committed = tuple(sorted({ candidate_ref_by_person[row.user_id] for row in allocations_by_plan[plan.id] if row.user_id in candidate_ref_by_person })) active_count = len(allocations_by_plan[plan.id]) unlimited = plan.max_seats is None available = 0 if unlimited else max(0, plan.max_seats or 0) - active_count + len(committed) supply_rows.append(SupplyPosition( resource=software_resources[plan.id], scope_key="global", available=max(0, available), reserved=len(committed), unlimited=unlimited, committed_candidate_refs=committed, )) for resource in sorted(set(vpn_resources.values()) | set(cloud_resources.values())): committed = tuple(sorted( candidate_ref_by_person[person_id] for person_id in candidate_person_ids if resource in resources_by_person[person_id] )) supply_rows.append(SupplyPosition( resource=resource, scope_key="global", available=len(committed), reserved=len(committed), unlimited=True, committed_candidate_refs=committed, )) # Per-collection generators, not one heterogeneous tuple: the models share # only ``Base``, which does not declare ``updated_at``. updated_times: list[datetime | None] = [ *(row.updated_at for row in processes), *(row.updated_at for row in people), *(row.updated_at for row in memberships), *(row.updated_at for row in assets), *(row.updated_at for row in plans), *(row.updated_at for row in vpn_rows), *(row.updated_at for row in cloud_rows), ] data_times = [_utc(value) for value in updated_times if value is not None] data_times.extend(_utc(row.allocated_at) for row in allocations) forecast = build_demand_forecast( as_of_date=today, candidates=candidates, peers=peers, supply=tuple(supply_rows), ) package = DemandForecastFactPackage( forecast=forecast, reference_map=reference_map, excluded_counts=dict(excluded), data_as_of=max(data_times, default=None), ) package.aggregate_facts() return package