""" Aura API helper utilities to detect paused instances and resume them. This prevents Neo4j driver failures like: - socket.gaierror: [Errno -2] Name or service not known - Cannot resolve address .databases.neo4j.io:7687 Those typically happen when AuraDB Free is auto-paused and the bolt hostname stops resolving. """ from __future__ import annotations import os import time from dataclasses import dataclass from typing import Any, Optional import requests import rich @dataclass(frozen=True) class AuraCredentials: """OAuth client credentials for the Neo4j Aura API.""" client_id: str client_secret: str class AuraAPIError(RuntimeError): """Raised when Aura API interactions fail.""" class AuraAPI: """ Minimal Aura API client for: - fetching an OAuth bearer token - reading instance status - resuming an instance if paused Docs: - Base URL: https://api.neo4j.io - Token endpoint: POST https://api.neo4j.io/oauth/token """ def __init__( self, creds: AuraCredentials, *, base_url: str = "https://api.neo4j.io" ) -> None: self._creds = creds self._base_url = base_url.rstrip("/") # remove trailing slash if provided self._token: Optional[str] = None self._token_expires_at: float = 0.0 def _get_token(self) -> str: """ Fetch and cache a bearer token (client_credentials flow). The Aura docs specify: POST https://api.neo4j.io/oauth/token grant_type=client_credentials Basic auth: : """ now = time.time() if self._token and now < self._token_expires_at - 30: return self._token resp = requests.post( f"{self._base_url}/oauth/token", auth=(self._creds.client_id, self._creds.client_secret), headers={"Content-Type": "application/x-www-form-urlencoded"}, data={"grant_type": "client_credentials"}, timeout=30, ) if resp.status_code >= 400: raise AuraAPIError(f"Failed to obtain Aura API token: {resp.status_code} {resp.text}") payload = resp.json() access_token = payload["access_token"] expires_in = float(payload.get("expires_in", 3600)) self._token = access_token self._token_expires_at = now + expires_in return access_token def _headers(self) -> dict[str, str]: return {"Authorization": f"Bearer {self._get_token()}"} def get_instance(self, instance_id: str) -> dict[str, Any]: """ Get instance details (includes status). NOTE: Endpoint path may evolve across Aura API versions. If your tenant uses a different path, adjust according to the API spec. """ url = f"{self._base_url}/v1/instances/{instance_id}" resp = requests.get(url, headers=self._headers(), timeout=30) if resp.status_code >= 400: raise AuraAPIError(f"Failed to get instance {instance_id}: {resp.status_code} {resp.text}") return resp.json() def get_instance_status(self, instance_id: str) -> str: """ Get the status of an Aura instance. Expected values (observed): - "running" - "paused" - "resuming" - "provisioning" Returns ------- str Lowercased status string. Raises ------ AuraAPIError If the response does not contain a status field. """ info = self.get_instance(instance_id) try: status = info["data"]["status"] except KeyError as e: raise AuraAPIError( f"Unexpected Aura API response shape; missing status.\n" f"Response keys: {list(info.keys())}\n" f"Full response: {info}" ) from e return str(status).lower() def resume_instance(self, instance_id: str) -> None: """ Trigger a resume of a paused instance. NOTE: Endpoint path may evolve across Aura API versions. Adjust according to the API spec if needed. """ url = f"{self._base_url}/v1/instances/{instance_id}/resume" resp = requests.post(url, headers=self._headers(), timeout=30) resp = requests.post( url, headers={ **self._headers(), "Content-Type": "application/json", # needed here }, json={}, # optional but helps make intent explicit timeout=30, ) if resp.status_code >= 400: raise AuraAPIError(f"Failed to resume instance {instance_id}: {resp.status_code} {resp.text}") def ensure_running( self, instance_id: str, *, poll_seconds: float = 5.0, # time to wait between status checks while polling timeout_seconds: float = 180.0, verbose: bool = True, ) -> None: """ If instance is paused, resume it and wait until status is Running. This is intended to run BEFORE creating a Neo4j driver session. """ start = time.time() status = self.get_instance_status(instance_id) if verbose: rich.print(f"💧 [kbdebugger] Aura instance {instance_id}: status={status!r}") if "paused" in status: if verbose: rich.print(f"⚠ī¸ [kbdebugger] Aura instance {instance_id} is paused → resuming via Aura API...") self.resume_instance(instance_id) # Poll until running (or timeout) while True: if time.time() - start > timeout_seconds: raise AuraAPIError( f"🛑☹ī¸ Timed out waiting for Aura instance {instance_id} to become running " f"(waited {timeout_seconds}s)." ) status = self.get_instance_status(instance_id) if "running" in status: if verbose: rich.print(f"🏃💧 [kbdebugger] Aura instance {instance_id} is running.") return if verbose: rich.print(f"[kbdebugger] Aura instance {instance_id} not running yet (status={status!r}); polling...") time.sleep(poll_seconds) def ensure_aura_running_from_env(*, verbose: bool = True) -> None: """ Convenience entrypoint: uses env vars to ensure the Aura instance is running. Required env vars: - AURA_API_CLIENT_ID - AURA_API_CLIENT_SECRET - AURA_INSTANCE_ID """ client_id = os.getenv("AURA_API_CLIENT_ID", "").strip() client_secret = os.getenv("AURA_API_CLIENT_SECRET", "").strip() instance_id = os.getenv("AURA_INSTANCE_ID", "").strip() if not (client_id and client_secret and instance_id): # If you want this to be mandatory, raise instead of returning. if verbose: rich.print("⚠ī¸ [kbdebugger] Aura auto-resume is not configured; skipping (missing env vars).") return api = AuraAPI( AuraCredentials(client_id=client_id, client_secret=client_secret) ) api.ensure_running(instance_id, verbose=verbose)