weaverstack 0.1.1__py3-none-any.whl

weaver/fabric/auth.py

@@ -0,0 +1,137 @@
+"""Azure tokens for Fabric, OneLake and SQL.
+
+Core does **not** decide which credential to use. It accepts an injected
+credential and, absent one, falls back to ``DefaultAzureCredential`` — the
+library default — without pinning the chain. Choosing a specific identity is a
+caller's policy, not the core's.
+
+That policy matters in practice: ``DefaultAzureCredential`` walks a chain and
+does not always settle on the identity you are signed in as, so on a machine
+where ``az`` works a OneLake write can still fail
+``401 Access token validation failed``. ``azure-identity`` 1.23 honours
+``AZURE_TOKEN_CREDENTIALS`` to pin the chain, and :func:`prefer_cli_credential`
+sets it to ``AzureCliCredential`` — but a **caller** invokes that (the desktop
+CLI does; the Fabric test infrastructure does). Core never sets it as a side
+effect of asking for a token.
+"""
+
+from __future__ import annotations
+
+import os
+
+#: Scopes. Generic technical values, not environment-specific.
+FABRIC_SCOPE = "https://api.fabric.microsoft.com/.default"
+STORAGE_SCOPE = "https://storage.azure.com/.default"
+SQL_SCOPE = "https://database.windows.net/.default"
+
+#: Honoured by azure-identity >= 1.23 to pin DefaultAzureCredential's chain.
+CREDENTIAL_ENV = "AZURE_TOKEN_CREDENTIALS"
+DEFAULT_CREDENTIAL = "AzureCliCredential"
+
+
+def prefer_cli_credential() -> str:
+    """Pin the credential chain to the Azure CLI, unless already chosen.
+
+    Policy, so a **caller** invokes it — the desktop CLI before a Fabric
+    command, the test infrastructure before the Fabric suite. Core never calls
+    it, so importing or using the core imposes no credential choice.
+    """
+
+    existing = os.environ.get(CREDENTIAL_ENV)
+    if existing:
+        return existing
+    os.environ[CREDENTIAL_ENV] = DEFAULT_CREDENTIAL
+    return DEFAULT_CREDENTIAL
+
+
+def credential():
+    """A default credential. Callers that want a specific one inject it instead."""
+
+    from azure.identity import DefaultAzureCredential
+
+    return DefaultAzureCredential()
+
+
+def get_token(scope: str, cred=None) -> str:
+    """An access token for one scope, from an injected credential or the default.
+
+    Answers the string and drops the expiry, which suits a one-shot command and
+    nothing that outlives one. Anything long-lived wants :class:`TokenProvider`.
+    """
+
+    return (cred or credential()).get_token(scope).token
+
+
+#: Renew this long before a token lapses, so a call already in flight when the
+#: margin opens still carries a valid one.
+TOKEN_REFRESH_MARGIN_SECONDS = 300.0
+
+
+class TokenProvider:
+    """A token for one scope, renewed shortly before it expires.
+
+    This exists because holding the *string* is a bug that only shows up in long
+    runs, and then shows up as something else. A caller that snapshots
+    :func:`get_token` keeps sending the same bearer until the API starts
+    answering ``401`` — and because ``AzureCliCredential`` serves from the Azure
+    CLI's own cache, the string may already be most of the way through its life
+    when it arrives. The usable budget is therefore *not* the nominal lifetime
+    and cannot be assumed from it: a run that starts with a nearly-spent token
+    has minutes, not an hour.
+
+    That is not hypothetical. A Fabric suite whose session snapshotted its token
+    died twenty minutes in with ``401: no body``, taking every downstream test
+    with it, and looked like six unrelated failures.
+
+    Refetching on every call would also be correct, and is what the SQL path does
+    — but there a token is fetched per *connection*. A REST client fetches per
+    *request*, and a credential shells out to ``az``, so the expiry is kept and
+    the token renewed only when it is close.
+    """
+
+    def __init__(
+        self,
+        scope: str,
+        cred=None,
+        *,
+        margin: float = TOKEN_REFRESH_MARGIN_SECONDS,
+    ) -> None:
+        self.scope = scope
+        self._cred = cred
+        self._margin = margin
+        self._token: str | None = None
+        self._expires_on = 0.0
+
+    def _credential(self):
+        # Built once and kept: constructing one per call would shell out to the
+        # CLI every time, which is the cost this class exists to avoid.
+        if self._cred is None:
+            self._cred = credential()
+        return self._cred
+
+    def __call__(self) -> str:
+        import time
+
+        if self._token is None or time.time() >= self._expires_on - self._margin:
+            acquired = self._credential().get_token(self.scope)
+            self._token = acquired.token
+            # A credential that reports no expiry gets renewed every call. Slow
+            # rather than wrong, and no shipped credential does it.
+            self._expires_on = float(getattr(acquired, "expires_on", 0) or 0)
+        return self._token
+
+
+def token_source(token=None, *, scope: str, cred=None):
+    """Normalise what a caller supplied into a zero-argument token source.
+
+    ``None`` builds a renewing :class:`TokenProvider`. A **string** is honoured
+    exactly as given — the caller owns it and its lifetime, which is how a Fabric
+    session passes on the identity it was handed. A **callable** is used as-is,
+    so a caller with its own refresh keeps it.
+    """
+
+    if token is None:
+        return TokenProvider(scope, cred)
+    if callable(token):
+        return token
+    return lambda: token

weaver/fabric/capacity.py

@@ -0,0 +1,143 @@
+"""Turning a Fabric capacity on and off.
+
+Capacity is billed while it runs, so this is the first and last thing a session
+touches. It goes through the Azure CLI rather than a REST call because capacity
+lives in ARM rather than in the Fabric API, and ``az`` already holds the
+subscription context.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass
+from typing import Sequence
+
+from ..errors import WeaverError
+
+CAPACITY_ACTIONS = ("status", "resume", "suspend")
+
+_AZ_VERB = {"status": "show", "resume": "resume", "suspend": "suspend"}
+
+#: Environment fallback for the subscription, when az has more than one.
+SUBSCRIPTION_ENV = "FABRIC_SUBSCRIPTION_ID"
+
+
+class CapacityError(WeaverError):
+    """Raised when a capacity action cannot be run."""
+
+
+@dataclass(frozen=True)
+class CapacityAction:
+    """The outcome of one capacity action."""
+
+    action: str
+    capacity: str
+    state: str | None
+    sku: str | None = None
+    returncode: int = 0
+
+    @property
+    def running(self) -> bool:
+        return (self.state or "").lower() == "active"
+
+    def __str__(self) -> str:
+        detail = f"{self.state or 'unknown'}"
+        if self.sku:
+            detail += f", {self.sku}"
+        return f"{self.capacity}: {detail}"
+
+
+def capacity_command(
+    action: str,
+    *,
+    resource_group: str,
+    capacity_name: str,
+    subscription_id: str | None = None,
+    extra_args: Sequence[str] = (),
+) -> list[str]:
+    """The Azure CLI command for one capacity action, without running it."""
+
+    verb = _AZ_VERB.get(action)
+    if verb is None:
+        raise CapacityError(
+            f"unknown capacity action {action!r} — expected one of "
+            + ", ".join(CAPACITY_ACTIONS)
+        )
+    if not resource_group:
+        raise CapacityError("a capacity needs its resource group")
+    if not capacity_name:
+        raise CapacityError("a capacity needs its name")
+
+    command = [
+        "az", "fabric", "capacity", verb,
+        "--resource-group", resource_group,
+        "--capacity-name", capacity_name,
+    ]
+    if subscription_id:
+        command.extend(["--subscription", subscription_id])
+    command.extend(extra_args)
+    return command
+
+
+def run_capacity_action(
+    action: str,
+    *,
+    resource_group: str,
+    capacity_name: str,
+    subscription_id: str | None = None,
+    extra_args: Sequence[str] = (),
+) -> CapacityAction:
+    """Run a capacity action and report the resulting state."""
+
+    if shutil.which("az") is None:
+        from ..diagnostics import install_command
+
+        raise CapacityError(
+            f"the Azure CLI is not installed. Install it:  {install_command('azure-cli')}"
+        )
+
+    command = capacity_command(
+        action,
+        resource_group=resource_group,
+        capacity_name=capacity_name,
+        subscription_id=subscription_id or os.environ.get(SUBSCRIPTION_ENV),
+        extra_args=(*extra_args, "--output", "json"),
+    )
+    completed = subprocess.run(command, capture_output=True, text=True, check=False)
+    if completed.returncode != 0:
+        raise CapacityError(
+            f"az {action} failed for {capacity_name!r}: "
+            + (completed.stderr.strip() or completed.stdout.strip() or "no output")
+        )
+
+    payload = _payload(completed.stdout)
+    return CapacityAction(
+        action=action,
+        capacity=capacity_name,
+        state=_state(payload),
+        sku=(payload.get("sku") or {}).get("name") if payload else None,
+        returncode=completed.returncode,
+    )
+
+
+def _payload(stdout: str) -> dict:
+    text = (stdout or "").strip()
+    if not text:
+        return {}
+    try:
+        loaded = json.loads(text)
+    except json.JSONDecodeError:
+        return {}
+    return loaded if isinstance(loaded, dict) else {}
+
+
+def _state(payload: dict) -> str | None:
+    """The running state, which az reports in more than one place."""
+
+    if not payload:
+        return None
+    properties = payload.get("properties") or {}
+    return properties.get("state") or payload.get("state")

weaver/fabric/client.py

@@ -0,0 +1,147 @@
+"""Fabric REST transport.
+
+Thin on purpose: a token, a base URL, and enough error translation that a
+failure says what failed rather than surfacing a bare HTTP status.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+from ..errors import WeaverError
+from .auth import FABRIC_SCOPE, token_source
+
+#: Generic technical defaults, not environment-specific.
+FABRIC_API = "https://api.fabric.microsoft.com/v1"
+ONELAKE_DFS = "https://onelake.dfs.fabric.microsoft.com"
+DEFAULT_TIMEOUT = 60.0
+DEFAULT_OPERATION_TIMEOUT = 900.0
+DEFAULT_OPERATION_POLL_INTERVAL = 2.0
+
+
+class FabricError(WeaverError):
+    """Raised when a Fabric API call fails."""
+
+    def __init__(self, message: str, *, status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class FabricClient:
+    """Authenticated access to the Fabric REST API."""
+
+    def __init__(
+        self,
+        *,
+        api_base_url: str = FABRIC_API,
+        token: str | None = None,
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> None:
+        self.api_base_url = api_base_url.rstrip("/")
+        self.timeout = timeout
+        self._token_source = token_source(token, scope=FABRIC_SCOPE)
+
+    @property
+    def token(self) -> str:
+        """A currently-valid bearer, renewed when it is close to expiring.
+
+        Read per request rather than cached: a client outlives its token, and a
+        stale one surfaces as ``401`` in whatever call happens to be next.
+        """
+
+        return self._token_source()
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        payload: Any = None,
+        expected: tuple[int, ...] = (200, 201, 202),
+    ):
+        import requests
+
+        url = path if path.startswith("http") else f"{self.api_base_url}/{path.lstrip('/')}"
+        response = requests.request(
+            method,
+            url,
+            headers={
+                "Authorization": f"Bearer {self.token}",
+                "Content-Type": "application/json",
+            },
+            data=json.dumps(payload) if payload is not None else None,
+            timeout=self.timeout,
+        )
+        if response.status_code not in expected:
+            raise FabricError(
+                f"{method} {url} returned {response.status_code}: "
+                f"{response.text.strip()[:400] or 'no body'}",
+                status_code=response.status_code,
+            )
+        return response
+
+    def get_json(self, path: str) -> dict:
+        response = self.request("GET", path, expected=(200,))
+        return response.json() if response.content else {}
+
+    def paged(self, path: str, *, key: str = "value") -> list[dict]:
+        """Every item across a paged listing."""
+
+        items: list[dict] = []
+        next_path: str | None = path
+        while next_path:
+            payload = self.get_json(next_path)
+            items.extend(payload.get(key, []))
+            next_path = payload.get("continuationUri")
+        return items
+
+    def wait_for_operation(
+        self,
+        response,
+        *,
+        timeout: float = DEFAULT_OPERATION_TIMEOUT,
+        poll_interval: float = DEFAULT_OPERATION_POLL_INTERVAL,
+    ) -> dict:
+        """Wait for a Fabric long-running-operation response to settle."""
+
+        if response.status_code != 202:
+            return response.json() if response.content else {}
+
+        location = response.headers.get("Location")
+        operation_id = response.headers.get("x-ms-operation-id")
+        if not location and operation_id:
+            location = f"operations/{operation_id}"
+        if not location:
+            raise FabricError(
+                "Fabric accepted a long-running operation without a polling location"
+            )
+
+        deadline = time.monotonic() + timeout
+        current = response
+        while time.monotonic() < deadline:
+            retry_after = current.headers.get("Retry-After")
+            try:
+                delay = float(retry_after) if retry_after is not None else poll_interval
+            except (TypeError, ValueError):
+                delay = poll_interval
+            time.sleep(max(0.0, delay))
+            current = self.request("GET", location, expected=(200,))
+            body = current.json() if current.content else {}
+            status = str(body.get("status") or "").casefold()
+            if status == "succeeded":
+                return body
+            if status in {"failed", "cancelled", "canceled"}:
+                error = body.get("error") or {}
+                message = error.get("message") if isinstance(error, dict) else None
+                raise FabricError(
+                    f"Fabric operation {operation_id or location} {status}"
+                    + (f": {message}" if message else "")
+                )
+            location = current.headers.get("Location") or location
+
+        raise FabricError(
+            f"Fabric operation {operation_id or location} did not finish within "
+            f"{int(timeout)}s"
+        )