from __future__ import annotations import json import re import subprocess import time from dataclasses import asdict, dataclass from datetime import datetime, timezone from typing import Protocol PD_SSD_QUOTA_METRIC = "SSD_TOTAL_GB" PD_BALANCED_QUOTA_METRIC = "DISKS_TOTAL_GB" STORAGE_CLASS_PREMIUM = "premium-rwo" STORAGE_CLASS_STANDARD = "standard-rwo" _REGION_TOKEN_RE = re.compile(r"^[a-z]+-[a-z]+\d+$") @dataclass(slots=True, frozen=True) class RegionCatalogEntry: region: str timezone_group: str country: str @dataclass(slots=True, frozen=True) class QuotaMetricHeadroom: metric: str limit_gb: float | None usage_gb: float | None @property def headroom_gb(self) -> float | None: if self.limit_gb is None or self.usage_gb is None: return None return self.limit_gb - self.usage_gb @dataclass(slots=True, frozen=True) class RegionQuotaSnapshot: region: str premium: QuotaMetricHeadroom standard: QuotaMetricHeadroom @dataclass(slots=True, frozen=True) class RegionLatencyResult: region: str latency_ms: float | None ok: bool error: str = "" @dataclass(slots=True, frozen=True) class ClusterStorageRequest: pgdata_gb: float wal_gb: float @property def total_gb(self) -> float: return max(0.0, self.pgdata_gb) + max(0.0, self.wal_gb) @dataclass(slots=True, frozen=True) class RegionFeasibility: region: str timezone_group: str premium_headroom_gb: float | None standard_headroom_gb: float | None can_premium: bool can_standard: bool default_pgdata_class: str default_wal_class: str feasibility_rank: int notes: tuple[str, ...] @dataclass(slots=True, frozen=True) class RegionCandidate: feasibility: RegionFeasibility latency_ms: float | None @dataclass(slots=True, frozen=True) class ClusterStorageBrowserResult: generated_at: str project_id: str timezone_group: str request: ClusterStorageRequest candidates: tuple[RegionCandidate, ...] quota_link: str quota_console_link: str def to_dict(self) -> dict: return { "generatedAt": self.generated_at, "projectId": self.project_id, "timezoneGroup": self.timezone_group, "request": asdict(self.request), "candidates": [ { "region": c.feasibility.region, "timezoneGroup": c.feasibility.timezone_group, "feasibilityRank": c.feasibility.feasibility_rank, "premiumHeadroomGb": c.feasibility.premium_headroom_gb, "standardHeadroomGb": c.feasibility.standard_headroom_gb, "canPremium": c.feasibility.can_premium, "canStandard": c.feasibility.can_standard, "defaultPgdataClass": c.feasibility.default_pgdata_class, "defaultWalClass": c.feasibility.default_wal_class, "latencyMs": c.latency_ms, "notes": list(c.feasibility.notes), } for c in self.candidates ], "quotaLink": self.quota_link, "quotaConsoleLink": self.quota_console_link, } def to_json(self) -> str: return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), indent=2) class RegionQuotaProvider(Protocol): def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]: ... class RegionLatencyProvider(Protocol): def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]: ... def default_region_catalog() -> tuple[RegionCatalogEntry, ...]: return ( RegionCatalogEntry("us-central1", "americas", "us"), RegionCatalogEntry("us-east1", "americas", "us"), RegionCatalogEntry("us-east4", "americas", "us"), RegionCatalogEntry("us-west1", "americas", "us"), RegionCatalogEntry("us-west2", "americas", "us"), RegionCatalogEntry("us-west3", "americas", "us"), RegionCatalogEntry("us-west4", "americas", "us"), RegionCatalogEntry("northamerica-northeast1", "americas", "ca"), RegionCatalogEntry("southamerica-east1", "americas", "br"), RegionCatalogEntry("europe-west1", "emea", "be"), RegionCatalogEntry("europe-west2", "emea", "uk"), RegionCatalogEntry("europe-west3", "emea", "de"), RegionCatalogEntry("europe-west4", "emea", "nl"), RegionCatalogEntry("europe-west6", "emea", "ch"), RegionCatalogEntry("europe-west8", "emea", "it"), RegionCatalogEntry("europe-west9", "emea", "fr"), RegionCatalogEntry("europe-west10", "emea", "de"), RegionCatalogEntry("europe-central2", "emea", "pl"), RegionCatalogEntry("asia-east1", "apac", "tw"), RegionCatalogEntry("asia-east2", "apac", "hk"), RegionCatalogEntry("asia-northeast1", "apac", "jp"), RegionCatalogEntry("asia-northeast2", "apac", "jp"), RegionCatalogEntry("asia-northeast3", "apac", "kr"), RegionCatalogEntry("asia-south1", "apac", "in"), RegionCatalogEntry("asia-south2", "apac", "in"), RegionCatalogEntry("asia-southeast1", "apac", "sg"), RegionCatalogEntry("asia-southeast2", "apac", "id"), RegionCatalogEntry("australia-southeast1", "apac", "au"), RegionCatalogEntry("australia-southeast2", "apac", "au"), ) def detect_timezone_group() -> str: now_local = datetime.now().astimezone() offset_hours = float(now_local.utcoffset().total_seconds()) / 3600.0 if now_local.utcoffset() else 0.0 if -10.0 <= offset_hours <= -2.0: return "americas" if -1.0 <= offset_hours <= 4.0: return "emea" return "apac" def filter_regions_by_timezone( regions: list[str], timezone_group: str, catalog: tuple[RegionCatalogEntry, ...] | None = None, ) -> list[RegionCatalogEntry]: entries = catalog or default_region_catalog() allowed = {r for r in regions if _REGION_TOKEN_RE.match(r)} return [entry for entry in entries if entry.region in allowed and entry.timezone_group == timezone_group] def _bool_feasible(headroom_gb: float | None, requested_gb: float) -> bool: if headroom_gb is None: return False return headroom_gb >= max(0.0, requested_gb) def _choose_default_class(premium_ok: bool, standard_ok: bool) -> str: if premium_ok: return STORAGE_CLASS_PREMIUM if standard_ok: return STORAGE_CLASS_STANDARD return STORAGE_CLASS_PREMIUM def _feasibility_rank(can_premium: bool, can_standard: bool) -> int: if can_premium: return 0 if can_standard: return 1 return 2 def evaluate_region_feasibility( region: str, timezone_group: str, quota: RegionQuotaSnapshot | None, requested: ClusterStorageRequest, ) -> RegionFeasibility: premium_headroom = quota.premium.headroom_gb if quota is not None else None standard_headroom = quota.standard.headroom_gb if quota is not None else None can_premium = _bool_feasible(premium_headroom, requested.total_gb) can_standard = _bool_feasible(standard_headroom, requested.total_gb) notes: list[str] = [] if quota is None: notes.append("quota-unavailable") if not can_premium and not can_standard: notes.append("insufficient-quota") default_class = _choose_default_class(can_premium, can_standard) return RegionFeasibility( region=region, timezone_group=timezone_group, premium_headroom_gb=premium_headroom, standard_headroom_gb=standard_headroom, can_premium=can_premium, can_standard=can_standard, default_pgdata_class=default_class, default_wal_class=default_class, feasibility_rank=_feasibility_rank(can_premium, can_standard), notes=tuple(notes), ) def _quota_links(project_id: str) -> tuple[str, str]: project = project_id.strip() if not project: return ( "https://cloud.google.com/compute/quotas", "https://console.cloud.google.com/iam-admin/quotas", ) return ( "https://cloud.google.com/compute/quotas", f"https://console.cloud.google.com/iam-admin/quotas?project={project}", ) def _sort_candidates(candidates: list[RegionCandidate]) -> list[RegionCandidate]: def _key(candidate: RegionCandidate): latency = candidate.latency_ms if candidate.latency_ms is not None else float("inf") return (candidate.feasibility.feasibility_rank, latency, candidate.feasibility.region) return sorted(candidates, key=_key) def build_cluster_storage_browser_result( *, project_id: str, requested: ClusterStorageRequest, available_regions: list[str], timezone_group: str | None = None, catalog: tuple[RegionCatalogEntry, ...] | None = None, quota_provider: RegionQuotaProvider, latency_provider: RegionLatencyProvider, ) -> ClusterStorageBrowserResult: tz_group = (timezone_group or detect_timezone_group()).strip().lower() or detect_timezone_group() candidates = filter_regions_by_timezone(available_regions, tz_group, catalog) candidate_regions = [entry.region for entry in candidates] quota_by_region = quota_provider.fetch(project_id=project_id, regions=candidate_regions) feasibility: list[RegionFeasibility] = [ evaluate_region_feasibility( region=entry.region, timezone_group=entry.timezone_group, quota=quota_by_region.get(entry.region), requested=requested, ) for entry in candidates ] probe_regions = [f.region for f in feasibility if f.can_premium or f.can_standard] latency_by_region = latency_provider.probe(probe_regions) combined: list[RegionCandidate] = [] for f in feasibility: latency = latency_by_region.get(f.region) combined.append( RegionCandidate( feasibility=f, latency_ms=latency.latency_ms if latency and latency.ok else None, ) ) link, console_link = _quota_links(project_id) return ClusterStorageBrowserResult( generated_at=datetime.now(timezone.utc).isoformat(), project_id=project_id, timezone_group=tz_group, request=requested, candidates=tuple(_sort_candidates(combined)), quota_link=link, quota_console_link=console_link, ) def _safe_float(value: str | int | float | None) -> float | None: if value is None: return None try: return float(value) except Exception: return None def _parse_limit_usage(raw: str) -> tuple[float | None, float | None]: # gcloud output format from --format=value(limit,usage) parts = [p.strip() for p in raw.split() if p.strip()] if len(parts) >= 2: return _safe_float(parts[0]), _safe_float(parts[1]) if len(parts) == 1: return _safe_float(parts[0]), None return None, None class GcloudRegionQuotaProvider: def __init__(self, timeout_sec: int = 20): self.timeout_sec = timeout_sec def _metric_headroom(self, project_id: str, region: str, metric: str) -> QuotaMetricHeadroom: cmd = [ "gcloud", "compute", "regions", "describe", region, "--project", project_id, "--format=value(quotas[metric=%s].limit,quotas[metric=%s].usage)" % (metric, metric), ] try: result = subprocess.run( cmd, capture_output=True, text=True, timeout=self.timeout_sec, ) if result.returncode != 0: return QuotaMetricHeadroom(metric=metric, limit_gb=None, usage_gb=None) limit, usage = _parse_limit_usage(result.stdout.strip()) return QuotaMetricHeadroom(metric=metric, limit_gb=limit, usage_gb=usage) except Exception: return QuotaMetricHeadroom(metric=metric, limit_gb=None, usage_gb=None) def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]: project = project_id.strip() if not project: return {} out: dict[str, RegionQuotaSnapshot] = {} for region in regions: premium = self._metric_headroom(project, region, PD_SSD_QUOTA_METRIC) standard = self._metric_headroom(project, region, PD_BALANCED_QUOTA_METRIC) out[region] = RegionQuotaSnapshot(region=region, premium=premium, standard=standard) return out class PingLatencyProvider: def __init__(self, timeout_sec: int = 2, count: int = 2): self.timeout_sec = timeout_sec self.count = count def _probe_one(self, region: str) -> RegionLatencyResult: # Best-effort: use region DNS style endpoint as latency anchor. host = f"{region}.gcping.com" cmd = ["ping", "-c", str(self.count), "-t", str(self.timeout_sec), host] started = time.monotonic() try: result = subprocess.run( cmd, capture_output=True, text=True, timeout=max(2, self.timeout_sec + 2), ) elapsed = (time.monotonic() - started) * 1000.0 if result.returncode == 0: return RegionLatencyResult(region=region, latency_ms=round(elapsed, 2), ok=True) return RegionLatencyResult(region=region, latency_ms=None, ok=False, error="ping-failed") except Exception as exc: return RegionLatencyResult(region=region, latency_ms=None, ok=False, error=str(exc)) def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]: return {region: self._probe_one(region) for region in regions} class StaticQuotaProvider: def __init__(self, values: dict[str, RegionQuotaSnapshot]): self.values = dict(values) def fetch(self, *, project_id: str, regions: list[str]) -> dict[str, RegionQuotaSnapshot]: _ = project_id return {r: self.values[r] for r in regions if r in self.values} class StaticLatencyProvider: def __init__(self, values: dict[str, RegionLatencyResult]): self.values = dict(values) def probe(self, regions: list[str]) -> dict[str, RegionLatencyResult]: return {r: self.values[r] for r in regions if r in self.values}