"""The Fetch client."""
from __future__ import annotations
import os
import random
import re
import time
import uuid
from typing import TYPE_CHECKING, Any, Literal, Self, get_args
from urllib.parse import quote
import httpx
from ._version import PACKAGE_NAME, __version__
from .auth import BearerToken, BrowserLogin, ClientCredentials, PasswordLogin, TokenAuth, cached_session
from .env import Environment, resolve
from .errors import APIError, NetworkError
from .models import Dataset, QueryResult, Tool, ToolResult, Validation, _columns
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
_FENCE = re.compile(r"^```(?:sql)?\s*|\s*```$", re.IGNORECASE)
MAX_LIMIT = 500
"""The largest dataset page the API serves; it lowers anything above this."""
Layer = Literal["gold", "silver", "bronze"]
LAYERS: tuple[Layer, ...] = get_args(Layer)
def _default_auth(env: Environment) -> httpx.Auth:
if token := os.environ.get("AIC_TOKEN"):
return BearerToken(token)
if (cid := os.environ.get("AIC_CLIENT_ID")) and (secret := os.environ.get("AIC_CLIENT_SECRET")):
return ClientCredentials(cid, secret)
if (email := os.environ.get("AIC_EMAIL")) and (password := os.environ.get("AIC_PASSWORD")):
return PasswordLogin(email, password, mfa_code=os.environ.get("AIC_MFA_CODE"))
return cached_session(env) or BrowserLogin()
[docs]
class Fetch:
"""Client for AI Commons Fetch: read-only SQL over your tenant data lake.
Without `token` or `auth`, credentials come from `$AIC_TOKEN`; then `$AIC_CLIENT_ID` and
`$AIC_CLIENT_SECRET`; then `$AIC_EMAIL` and `$AIC_PASSWORD`; then the cached session;
then browser sign-in.
with Fetch("dev") as fetch:
result = fetch.query('SELECT count(*) FROM gold."<dataset-id>"')
"""
def __init__(
self,
env: str | Environment | None = None,
*,
token: str | None = None,
auth: httpx.Auth | None = None,
language: str = "en",
timeout: float = 30.0,
retries: int = 3,
sql_tool: str = "text_to_sql",
transport: httpx.BaseTransport | None = None,
) -> None:
self.env = resolve(env)
auth = auth or (BearerToken(token) if token else _default_auth(self.env))
if isinstance(auth, TokenAuth):
auth.bind(self.env)
self.auth = auth
self.retries = max(1, retries)
self.sql_tool = sql_tool
self.tools = Tools(self)
self._http = httpx.Client(
auth=auth,
timeout=timeout,
transport=transport,
headers={
"Accept-Language": language,
"User-Agent": f"{PACKAGE_NAME}/{__version__}",
},
)
[docs]
def query(self, sql: str, *, fresh: bool = False) -> QueryResult:
"""Run one SELECT. `fresh` reads S3 as it is now, bypassing the index and cache."""
r = self._request("POST", self._url("/query/"), json={"sql": sql}, fresh=fresh)
body = r.json()
return QueryResult(
sql=sql,
columns=_columns(body["columns"]),
rows=body["rows"],
row_count=body["row_count"],
truncated=body["truncated"],
elapsed_ms=body["elapsed_ms"],
cache=r.headers.get("X-Cache"),
request_id=r.headers.get("X-Request-ID"),
)
[docs]
def validate(self, sql: str, *, fresh: bool = False) -> Validation:
"""Check one SELECT as `query` would, without reading data. Invalid SQL raises `QueryError`."""
return Validation.parse(self._request("POST", self._url("/query/validate/"), json={"sql": sql}, fresh=fresh).json())
[docs]
def sql(self, question: str, *, datasets: Sequence[str] = ()) -> str:
"""Translate a natural-language question to SQL with ds-tools."""
payload: dict[str, Any] = {"question": question}
if datasets:
payload["datasets"] = list(datasets)
return _FENCE.sub("", self.tools.invoke(self.sql_tool, payload).output.strip())
[docs]
def ask(self, question: str, *, datasets: Sequence[str] = (), fresh: bool = False) -> QueryResult:
"""Answer a natural-language question: translate it to SQL, then run it."""
return self.query(self.sql(question, datasets=datasets), fresh=fresh)
[docs]
def datasets(
self,
*,
names: Sequence[str] = (),
layers: Sequence[Layer] = (),
limit: int = 100,
fresh: bool = False,
) -> Iterator[Dataset]:
"""Every dataset and pipeline run you are granted, across pages; only `names` and `layers` if given.
Names that do not exist or are not granted are left out. `limit` is the page size, held to
1…`MAX_LIMIT`; the server describes every dataset on a page within its timeout, so large pages risk a 504.
"""
if bad := sorted(set(layers) - set(LAYERS)):
raise ValueError(f"unknown layer {', '.join(bad)}; use {', '.join(LAYERS)}")
limit, offset, seen = max(1, min(limit, MAX_LIMIT)), 0, set[str]()
while True:
params: dict[str, Any] = {"limit": limit, "offset": offset, "name": list(names), "layer": list(layers)}
body = self._request("GET", self._url("/dataset/"), params=params, fresh=fresh).json()
for ds in map(Dataset.parse, body["data"]):
if ds.name not in seen:
seen.add(ds.name)
yield ds
if not body["page"]["has_next"] or not body["data"]:
return
offset += len(body["data"])
[docs]
def dataset(self, name: str, *, fresh: bool = False) -> Dataset:
"""One dataset by name, e.g. `gold.<dataset-id>` or `silver.<pipeline-id>/<job-id>`."""
url = self._url(f"/dataset/{quote(name, safe='')}/")
return Dataset.parse(self._request("GET", url, fresh=fresh).json())
def health(self) -> dict[str, Any]:
return self._public("/health-check/")
def version(self) -> str:
return str(self._public("/version/")["version"])
def close(self) -> None:
self._http.close()
def __enter__(self) -> Self:
return self
def __exit__(self, *_: object) -> None:
self.close()
def _url(self, path: str) -> str:
return f"{self.env.fetch_url}{path}"
def _public(self, path: str) -> dict[str, Any]:
return dict(self._request("GET", self._url(path), auth=None).json())
def _request(self, method: str, url: str, *, fresh: bool = False, **kwargs: Any) -> httpx.Response:
headers = {"Cache-Control": "no-cache"} if fresh else None
for attempt in range(self.retries):
last = attempt == self.retries - 1
try:
r = self._http.request(method, url, headers=headers, **kwargs)
except httpx.TransportError as e:
if last:
raise NetworkError(f"cannot reach {url}: {e}") from e
time.sleep(_backoff(None, attempt))
continue
if r.is_success:
return r
if last or not _retryable(r):
raise APIError.from_response(r)
time.sleep(_backoff(r, attempt))
raise AssertionError("unreachable")
def _retryable(r: httpx.Response) -> bool:
if r.status_code in (429, 500, 503):
return True
return r.status_code >= 500 and "json" not in r.headers.get("Content-Type", "")
def _backoff(r: httpx.Response | None, attempt: int) -> float:
retry_after = r.headers.get("Retry-After", "") if r is not None else ""
base = int(retry_after) if retry_after.isdigit() else 0.5 * 2**attempt
return base + random.uniform(0, 0.5)