Azure Cost API Pagination and Deduplication Guide

import asyncio
import hashlib
import logging
from typing import Any, AsyncGenerator, Dict, List, Optional, Set

import httpx

logger = logging.getLogger(__name__)


class AzureCostPaginator:
    """Async paginator for the Azure Cost Management Query API.

    The Query endpoint uses POST for the initial request and GET for nextLink
    pages. This class handles both, plus rate-limit backoff, cursor cycle
    detection, and in-stream SHA-256 deduplication.
    """

    def __init__(self, token: str, base_url: str = "https://management.azure.com"):
        self.client = httpx.AsyncClient(
            base_url=base_url,
            headers={
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json",
            },
            timeout=httpx.Timeout(60.0, connect=10.0),
            limits=httpx.Limits(max_connections=10),
        )
        self._seen_cursors: Set[str] = set()

    async def aclose(self) -> None:
        await self.client.aclose()

    async def _request_with_backoff(
        self,
        method: str,
        url: str,
        body: Optional[Dict[str, Any]] = None,
        max_retries: int = 5,
    ) -> Dict[str, Any]:
        retries = 0
        while retries < max_retries:
            try:
                if method.upper() == "POST":
                    resp = await self.client.post(url, json=body)
                else:
                    resp = await self.client.get(url)

                if resp.status_code == 429:
                    retry_after = float(resp.headers.get("Retry-After", 2 ** retries))
                    logger.warning("Rate limited. Backing off for %.2fs", retry_after)
                    await asyncio.sleep(retry_after)
                    retries += 1
                    continue
                resp.raise_for_status()
                return resp.json()
            except httpx.HTTPStatusError as exc:
                if exc.response.status_code == 429:
                    retry_after = float(exc.response.headers.get("Retry-After", 2 ** retries))
                    await asyncio.sleep(retry_after)
                    retries += 1
                    continue
                raise
            except httpx.RequestError as exc:
                logger.error("Request failed: %s", exc)
                retries += 1
                await asyncio.sleep(min(2 ** retries, 30))
        raise RuntimeError(f"Max retries ({max_retries}) exceeded for {url}")

    async def paginate(
        self, query_url: str, query_body: Dict[str, Any]
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """Yield raw pages: POST the first, GET each nextLink verbatim."""
        payload = await self._request_with_backoff("POST", query_url, body=query_body)
        yield payload

        while True:
            next_link = payload.get("properties", {}).get("nextLink")
            if not next_link:
                break
            if next_link in self._seen_cursors:
                logger.warning("Cycle detected in pagination at %s. Breaking.", next_link)
                break
            self._seen_cursors.add(next_link)
            payload = await self._request_with_backoff("GET", next_link)
            yield payload

    @staticmethod
    def generate_record_hash(columns: List[str], row: List[Any]) -> str:
        """Derive a deduplication hash from immutable, column-named dimensions."""
        row_dict = dict(zip(columns, row))
        composite = (
            f"{row_dict.get('SubscriptionId', '')}|"
            f"{row_dict.get('MeterId', '')}|"
            f"{row_dict.get('UsageDate', '')}|"
            f"{row_dict.get('InstanceId', '')}|"
            f"{row_dict.get('Quantity', 0)}"
        )
        return hashlib.sha256(composite.encode("utf-8")).hexdigest()

    async def stream_deduplicated(
        self, query_url: str, query_body: Dict[str, Any]
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """Yield unique, deduplicated row dicts from every page."""
        seen_hashes: Set[str] = set()
        async for page in self.paginate(query_url, query_body):
            props = page.get("properties", {})
            columns = [col["name"] for col in props.get("columns", [])]
            for row in props.get("rows", []):
                record_hash = self.generate_record_hash(columns, row)
                if record_hash not in seen_hashes:
                    seen_hashes.add(record_hash)
                    yield dict(zip(columns, row))

async def write_to_warehouse(records: List[Dict[str, Any]]) -> None:
    """Stub: replace with a MERGE/INSERT OVERWRITE keyed on the dedup dimensions."""
    logger.info("Writing %d records to warehouse.", len(records))


async def consume_billing_stream(
    query_url: str,
    query_body: Dict[str, Any],
    token: str,
    batch_size: int = 5000,
) -> None:
    paginator = AzureCostPaginator(token=token)
    buffer: List[Dict[str, Any]] = []
    try:
        async for record in paginator.stream_deduplicated(query_url, query_body):
            buffer.append(record)
            if len(buffer) >= batch_size:
                await write_to_warehouse(buffer)
                buffer.clear()
        if buffer:
            await write_to_warehouse(buffer)
    finally:
        await paginator.aclose()


async def main() -> None:
    import os

    subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
    token = os.environ["AZURE_BEARER_TOKEN"]  # Obtain via DefaultAzureCredential

    query_url = (
        f"https://management.azure.com/subscriptions/{subscription_id}"
        f"/providers/Microsoft.CostManagement/query?api-version=2023-11-01"
    )
    query_body = {
        "type": "ActualCost",
        "timeframe": "BillingMonthToDate",
        "dataset": {
            "granularity": "Daily",
            "aggregation": {"totalCost": {"name": "PreTaxCost", "function": "Sum"}},
            "grouping": [{"type": "Dimension", "name": "ResourceGroup"}],
        },
    }
    await consume_billing_stream(query_url, query_body, token=token)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(main())

def test_dedup_drops_overlapping_rows() -> None:
    columns = ["SubscriptionId", "MeterId", "UsageDate", "InstanceId", "Quantity"]
    row = ["sub-1", "meter-a", "2026-06-01", "vm-1", 4]
    h1 = AzureCostPaginator.generate_record_hash(columns, row)
    h2 = AzureCostPaginator.generate_record_hash(columns, list(row))  # same dims
    assert h1 == h2, "identical dimensions must hash identically"

    changed = ["sub-1", "meter-a", "2026-06-01", "vm-1", 5]  # quantity differs
    assert AzureCostPaginator.generate_record_hash(columns, changed) != h1


if __name__ == "__main__":
    test_dedup_drops_overlapping_rows()
    print("dedup invariants hold")