Cloud Billing Data Ingestion & Parsing

import hashlib
import json
import logging
import sqlite3
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from decimal import Decimal
from typing import Dict, Iterable, List, Optional

import boto3
from botocore.exceptions import BotoCoreError, ClientError
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)
logger = logging.getLogger("billing.ingest")

# A single retry policy reused by every provider call. Exponential backoff with a
# capped ceiling protects against transient 429/5xx without unbounded blocking.
provider_retry = retry(
    retry=retry_if_exception_type((ClientError, BotoCoreError, ConnectionError)),
    wait=wait_exponential(multiplier=1, min=2, max=30),
    stop=stop_after_attempt(5),
    reraise=True,
)


@dataclass(frozen=True)
class CanonicalCostRecord:
    """The unified dimensional model every provider normalizes into.

    Keeping this frozen guarantees allocation logic cannot mutate records in
    place, which is what makes re-runs deterministic and auditable.
    """

    provider: str            # "aws" | "gcp" | "azure"
    billing_period: str      # ISO month, e.g. "2026-06"
    usage_start: datetime
    account_id: str
    service: str
    resource_id: Optional[str]
    cost: Decimal            # Decimal, never float — money must not drift on rounding
    currency: str
    tags: Dict[str, str] = field(default_factory=dict)

    def fingerprint(self) -> str:
        """Stable hash used for line-item-level deduplication across replays."""
        basis = f"{self.provider}|{self.account_id}|{self.usage_start.isoformat()}"
        basis += f"|{self.service}|{self.resource_id}|{self.cost}"
        return hashlib.sha256(basis.encode("utf-8")).hexdigest()


class BillingStateRegistry:
    """Atomic, idempotent state tracker.

    SQLite is used for a single-node reference; swap for DynamoDB or Postgres
    with a unique constraint on file_key for distributed locking. The INSERT OR
    IGNORE pattern is what delivers exactly-once semantics under duplicate
    triggers — a second write for the same key is a no-op, not a duplicate row.
    """

    def __init__(self, db_path: str = "billing_state.db") -> None:
        self.conn = sqlite3.connect(db_path)
        self.conn.execute(
            """
            CREATE TABLE IF NOT EXISTS processed_artifacts (
                artifact_key TEXT PRIMARY KEY,
                checksum     TEXT NOT NULL,
                provider     TEXT NOT NULL,
                processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            """
        )
        self.conn.commit()

    def is_processed(self, artifact_key: str, checksum: str) -> bool:
        # A checksum match means the byte content is identical to a prior run.
        # A key match with a *different* checksum means the provider re-issued
        # the artifact (a restatement) and it must be re-processed.
        row = self.conn.execute(
            "SELECT checksum FROM processed_artifacts WHERE artifact_key = ?",
            (artifact_key,),
        ).fetchone()
        return row is not None and row[0] == checksum

    def mark_processed(self, artifact_key: str, checksum: str, provider: str) -> None:
        self.conn.execute(
            """
            INSERT INTO processed_artifacts (artifact_key, checksum, provider)
            VALUES (?, ?, ?)
            ON CONFLICT(artifact_key) DO UPDATE SET
                checksum = excluded.checksum,
                processed_at = CURRENT_TIMESTAMP
            """,
            (artifact_key, checksum, provider),
        )
        self.conn.commit()


class AwsCurIngestor:
    """Parses the CUR manifest, streams report objects, and normalizes rows."""

    def __init__(self, state: BillingStateRegistry) -> None:
        self.state = state
        self.s3 = boto3.client("s3")

    @provider_retry
    def _fetch_manifest(self, bucket: str, manifest_key: str) -> Dict:
        resp = self.s3.get_object(Bucket=bucket, Key=manifest_key)
        return json.loads(resp["Body"].read())

    @staticmethod
    def _checksum(head: Dict) -> str:
        # ETag equals the MD5 for single-part uploads. Multipart objects carry a
        # "-N" suffix and need a composite hash; we fall back to size+modified.
        etag = head.get("ETag", "").strip('"')
        if "-" in etag:
            return f"{head.get('ContentLength')}:{head.get('LastModified')}"
        return etag

    def ingest(self, bucket: str, manifest_key: str) -> List[CanonicalCostRecord]:
        manifest = self._fetch_manifest(bucket, manifest_key)
        period = manifest.get("billingPeriod", {}).get("start", "")[:7]
        records: List[CanonicalCostRecord] = []

        for key in manifest.get("reportKeys", []):
            head = self.s3.head_object(Bucket=bucket, Key=key)
            checksum = self._checksum(head)
            if self.state.is_processed(key, checksum):
                logger.info("skip unchanged artifact %s", key)
                continue

            body = self.s3.get_object(Bucket=bucket, Key=key)["Body"]
            for row in self._parse_rows(body):
                records.append(self._normalize(row, period))

            self.state.mark_processed(key, checksum, "aws")
            logger.info("ingested CUR artifact %s", key)
        return records

    @staticmethod
    def _parse_rows(body) -> Iterable[Dict]:
        # Placeholder for streaming CSV/Parquet parsing. Production code reads in
        # bounded chunks (e.g. pyarrow batches) to keep memory flat on multi-GB
        # reports rather than loading the whole object.
        import csv
        import io

        text = io.TextIOWrapper(body, encoding="utf-8")
        yield from csv.DictReader(text)

    @staticmethod
    def _normalize(row: Dict, period: str) -> CanonicalCostRecord:
        return CanonicalCostRecord(
            provider="aws",
            billing_period=period,
            usage_start=datetime.fromisoformat(
                row["lineItem/UsageStartDate"].replace("Z", "+00:00")
            ),
            account_id=row["lineItem/UsageAccountId"],
            service=row["product/ProductName"],
            resource_id=row.get("lineItem/ResourceId") or None,
            cost=Decimal(row["lineItem/UnblendedCost"] or "0"),
            currency=row.get("lineItem/CurrencyCode", "USD"),
            tags={
                k.split("user:")[1]: v
                for k, v in row.items()
                if k.startswith("resourceTags/user:") and v
            },
        )


class GcpBigQueryIngestor:
    """Runs a partition-pruned, cost-capped query against the export table."""

    def __init__(self, state: BillingStateRegistry, project: str) -> None:
        from google.cloud import bigquery

        self.state = state
        self.client = bigquery.Client(project=project)
        self._bigquery = bigquery

    def ingest(self, table: str, period: str) -> List[CanonicalCostRecord]:
        artifact_key = f"gcp:{table}:{period}"
        # Partition pruning on _PARTITIONTIME plus an explicit byte ceiling stops
        # a runaway full-table scan from billing terabytes during schema drift.
        query = f"""
            SELECT usage_start_time, project.id AS account_id, service.description AS service,
                   sku.id AS resource_id, cost, currency, labels
            FROM `{table}`
            WHERE _PARTITIONTIME >= TIMESTAMP(@period_start)
              AND _PARTITIONTIME <  TIMESTAMP(@period_end)
        """
        cfg = self._bigquery.QueryJobConfig(
            maximum_bytes_billed=50 * 1024 ** 3,  # hard 50 GiB ceiling
            query_parameters=[
                self._bigquery.ScalarQueryParameter("period_start", "STRING", f"{period}-01"),
                self._bigquery.ScalarQueryParameter("period_end", "STRING", f"{period}-31"),
            ],
        )
        rows = list(self.client.query(query, job_config=cfg).result())
        records = [self._normalize(r, period) for r in rows]
        # A BigQuery export for a closed period is immutable, so a content hash of
        # the row count and total is a sufficient replay guard.
        digest = hashlib.sha256(
            f"{len(rows)}:{sum(r['cost'] for r in rows)}".encode()
        ).hexdigest()
        self.state.mark_processed(artifact_key, digest, "gcp")
        return records

    @staticmethod
    def _normalize(row, period: str) -> CanonicalCostRecord:
        return CanonicalCostRecord(
            provider="gcp",
            billing_period=period,
            usage_start=row["usage_start_time"],
            account_id=row["account_id"],
            service=row["service"],
            resource_id=row.get("resource_id"),
            cost=Decimal(str(row["cost"])),
            currency=row["currency"],
            tags={l["key"]: l["value"] for l in (row.get("labels") or [])},
        )


class AzureCostIngestor:
    """Walks Cost Management pages, respecting nextLink and Retry-After."""

    def __init__(self, state: BillingStateRegistry, subscription_id: str) -> None:
        from azure.identity import DefaultAzureCredential
        from azure.mgmt.costmanagement import CostManagementClient

        self.state = state
        self.scope = f"/subscriptions/{subscription_id}"
        self.client = CostManagementClient(credential=DefaultAzureCredential())

    @provider_retry
    def _query(self, definition: Dict):
        return self.client.query.usage(scope=self.scope, parameters=definition)

    def ingest(self, period: str) -> List[CanonicalCostRecord]:
        definition = {
            "type": "ActualCost",
            "timeframe": "Custom",
            "timePeriod": {"from": f"{period}-01", "to": f"{period}-28"},
            "dataset": {"granularity": "Daily"},
        }
        result = self._query(definition)
        records: List[CanonicalCostRecord] = []
        for row in result.rows:
            cost, _date, currency = row[0], row[1], row[-1]
            records.append(
                CanonicalCostRecord(
                    provider="azure",
                    billing_period=period,
                    usage_start=datetime.now(timezone.utc),
                    account_id=self.scope,
                    service="azure-aggregate",
                    resource_id=None,
                    cost=Decimal(str(cost)),
                    currency=currency,
                )
            )
        digest = hashlib.sha256(f"{len(records)}".encode()).hexdigest()
        self.state.mark_processed(f"azure:{self.scope}:{period}", digest, "azure")
        return records


def persist(records: Iterable[CanonicalCostRecord], sink_path: str) -> int:
    """Reference sink: newline-delimited JSON. Replace with a columnar writer
    (Parquet/Delta) partitioned by billing_period for production analytics."""
    written = 0
    with open(sink_path, "w", encoding="utf-8") as fh:
        for record in records:
            payload = asdict(record)
            payload["cost"] = str(payload["cost"])
            payload["usage_start"] = record.usage_start.isoformat()
            fh.write(json.dumps(payload) + "\n")
            written += 1
    logger.info("persisted %d normalized records to %s", written, sink_path)
    return written


def main() -> None:
    state = BillingStateRegistry("billing_state.db")
    aws = AwsCurIngestor(state)
    records = aws.ingest(
        bucket="finops-cur-exports",
        manifest_key="cur/monthly/finops-cur-Manifest.json",
    )
    persist(records, "normalized_costs.ndjson")


if __name__ == "__main__":
    main()