Auth and cache¶

def auth_status(validate: bool = True):
    """Which credentials connecto can see, whether they *work*, and who you are.

    ``validate=True`` (the default) asks each service. It costs about a second, and
    it is the whole point: a presence-only check reports ``found`` for a token that
    expired last week, which is precisely the wrong answer for the one person who
    ever runs this function.
    """
    import pandas as pd

    rows = []

    def add(service, server, info, ident=None):
        if not info.ok:
            status = "MISSING"
        elif ident is None or not ident.validated:
            # Present, but nobody asked the server. Not the same as OK.
            status = "FOUND"
        elif ident.valid:
            status = "OK"
        else:
            status = "INVALID"
        rows.append(
            {
                "service": service,
                "server": server or "-",
                "status": status,
                "source": info.source,
                "identity": (ident.email if ident else None) or "-",
                "access": (ident.access if ident else None)
                or (ident.detail if ident else "")
                or "-",
            }
        )

    cave = get_token("cave")
    add("cave", "global.daf-apis.com", cave,
        validate_token("cave") if (validate and cave.ok) else None)

    for server in _neuprint_servers():
        info = get_token("neuprint", server=server)
        add("neuprint", server, info,
            validate_token("neuprint", server=server) if (validate and info.ok) else None)

    for service in ("clio", "seatable"):
        info = get_token(service)
        add(service, None, info,
            validate_token(service) if (validate and info.ok) else None)

    df = pd.DataFrame(rows)

    # If connecto and caveclient have resolved *different* CAVE tokens, say so.
    # Two secret stores that quietly disagree is the fafbseg bug; it should never
    # be something you have to discover for yourself.
    native, _ = _cave_token()
    if cave.ok and native and native != cave.token:
        df.attrs["warning"] = (
            f"connecto is using the CAVE token from {cave.source}, but caveclient "
            f"resolves a different one from ~/.cloudvolume/secrets. They disagree; "
            f'run co.set_token("cave", ...) to make them match.'
        )
        print("!! " + df.attrs["warning"])

    return df

def get_token(
    service: str, *, token: str | None = None, server: str | None = None
) -> TokenInfo:
    """Resolve a token for ``service`` ("cave", "neuprint", "clio", "seatable")."""
    service = service.lower()

    if token:
        return TokenInfo(service, token, "argument")

    env = f"CONNECTO_{service.upper()}_TOKEN"
    if os.environ.get(env):
        return TokenInfo(service, os.environ[env], f"${env}")

    native_env = {
        "neuprint": "NEUPRINT_APPLICATION_CREDENTIALS",
        "clio": "CLIO_TOKEN",
        "seatable": "SEATABLE_TOKEN",
    }.get(service)
    if native_env and os.environ.get(native_env):
        raw = os.environ[native_env]
        # neuprint accepts either the bare token or the whole JSON doc.
        if raw.strip().startswith("{"):
            try:
                raw = json.loads(raw)["token"]
            except (json.JSONDecodeError, KeyError):
                pass
        return TokenInfo(service, raw, f"${native_env}")

    cfg = _read_config().get("auth", {})
    if server and cfg.get(server):
        return TokenInfo(service, cfg[server], f"{CONFIG_PATH.name} [{server}]")
    if cfg.get(service):
        return TokenInfo(service, cfg[service], CONFIG_PATH.name)

    if service == "clio":
        # clio piggybacks on the Google-auth token clio-py manages itself.
        return TokenInfo(service, None, "clio-py (managed)")

    if service == "cave":
        tok, src = _cave_token(server)
        return TokenInfo(service, tok, src)

    return TokenInfo(service, None, "not found")

def set_token(service: str, token: str, *, write_native: bool = True):
    """Store a token.

    With ``write_native=True`` (the default) the token is written to *both*
    connecto's config and the service's native location, so that caveclient,
    cloud-volume and connecto can never disagree about which token is current -
    which is the failure mode fafbseg papers over by reading from cloud-volume and
    writing via caveclient.
    """
    service = service.lower()

    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    cfg = _read_config()
    cfg.setdefault("auth", {})[service] = token

    lines = ["[auth]"]
    lines += [f'{k} = "{v}"' for k, v in cfg["auth"].items()]
    for section, values in cfg.items():
        if section == "auth":
            continue
        lines.append(f"\n[{section}]")
        lines += [f"{k} = {v!r}" for k, v in values.items()]
    CONFIG_PATH.write_text("\n".join(lines) + "\n")

    if write_native and service == "cave":
        from caveclient import CAVEclient

        CAVEclient().auth.save_token(token, overwrite=True)

def validate_token(service: str, *, token=None, server=None) -> Identity:
    """Ask the service whether a token actually works, and who it belongs to.

    One cheap request (~0.4s). This is the difference between "a string exists"
    and "you are authenticated", and only the second one is worth reporting.
    """
    service = service.lower()
    info = get_token(service, token=token, server=server)

    if not info.ok:
        return Identity(False, detail="no token found")

    headers = {"Authorization": f"Bearer {info.token}"}

    if service == "cave":
        url = CAVE_IDENTITY.format(server=(server or CAVE_GLOBAL_SERVER).rstrip("/"))
    elif service == "neuprint":
        url = NEUPRINT_IDENTITY.format(server=(server or "neuprint.janelia.org"))
    else:
        # No cheap probe for clio/seatable. Say so rather than implying one ran.
        return Identity(True, detail="present, not validated", validated=False)

    try:
        r = requests.get(url, headers=headers, timeout=TIMEOUT)
    except requests.RequestException as e:
        return Identity(False, detail=f"could not reach server ({type(e).__name__})")

    if r.status_code in (401, 403):
        return Identity(False, detail="server rejected this token")
    if not r.ok:
        return Identity(False, detail=f"HTTP {r.status_code}")

    try:
        data = r.json()
    except ValueError:
        return Identity(True, detail="ok")

    if service == "cave":
        groups = data.get("groups") or []
        return Identity(
            True,
            email=data.get("email"),
            access=", ".join(groups) if groups else None,
        )

    return Identity(True, email=data.get("Email"), access=data.get("AuthLevel"))

@contextmanager
def auth_errors(service: str, *, server=None, resource=None, dataset=None):
    """Translate upstream 401/403s into something the user can act on.

    caveclient raises its own ``AuthException`` when no token is configured, and a
    bare ``requests.HTTPError`` when the server rejects one. neuprint-python (and
    connecto's own ``server_datasets``) raise a bare 401 either way, so a missing
    token and a wrong token are indistinguishable. Both become a
    :class:`ConnectoAuthError` here, and the three kinds stay distinct.
    """
    try:
        yield
    except ConnectoAuthError:
        raise
    except requests.HTTPError as e:
        code = getattr(e.response, "status_code", None)
        if code not in (401, 403):
            raise
        kind = (
            ConnectoAuthError.FORBIDDEN
            if code == 403
            else ConnectoAuthError.INVALID
        )
        info = get_token(service, server=server)
        # A 401 with no token at all is a *missing* token, not an invalid one.
        if kind == ConnectoAuthError.INVALID and not info.ok:
            kind = ConnectoAuthError.MISSING
        raise ConnectoAuthError(
            _message(kind, service, info, server, resource, dataset),
            kind=kind, service=service, server=server,
        ) from e
    except Exception as e:  # noqa: BLE001
        if type(e).__name__ != "AuthException":  # caveclient's "no token" path
            raise
        info = get_token(service, server=server)
        raise ConnectoAuthError(
            _message(
                ConnectoAuthError.MISSING, service, info, server, resource, dataset
            ),
            kind=ConnectoAuthError.MISSING, service=service, server=server,
        ) from e

def server_status(dataset, *, timeout: int = PROBE_TIMEOUT):
    """Which servers behind ``dataset`` are answering, right now.

    Probes each service that the dataset's backend depends on and reports what it
    said. The point is the *shape* of the result rather than any single row: one
    service down means that service is having a bad day (on CAVE, almost always a
    materialization in progress); every service down means the deployment or your
    own network is gone, which is a different problem with a different fix.

    Parameters
    ----------
    dataset :   str | Dataset
                A registered dataset name (``"flywire"``) or a dataset object.
    timeout :   int
                Seconds to wait for each probe.

    Returns
    -------
    pandas.DataFrame
                One row per service: ``service, server, status, http, seconds, backs``.
                ``status`` is ``OK``, ``UNAVAILABLE``, ``ERROR``, ``UNREACHABLE`` or
                ``NO AUTH`` - never a bare boolean, because "answered" and "answered
                usefully" are different claims.

    Examples
    --------
    >>> co.server_status("flywire")                             # doctest: +SKIP
       service              server  status  http  seconds  backs
    0     info  global.daf-apis.com     OK   200     0.31  client setup, ...
    1  materialize  prod.flywire-daf.com  UNAVAILABLE  503  0.44  annotations, ...
    """
    import pandas as pd

    from .core.registry import get_dataset

    ds = get_dataset(dataset) if isinstance(dataset, str) else dataset
    service = ds.backend_kind

    client = None
    if service == "cave":
        # Unlike the error path, a cold status check may construct a client - but it
        # must not explode when the thing it is diagnosing is exactly that failing.
        try:
            client = ds.client
        except Exception:  # noqa: BLE001
            client = None

    token = get_token(service, server=ds._auth_server).token
    services = _services_for(service, dataset=ds, client=client)
    return pd.DataFrame(_probe_all(services, token, timeout))

def wait_until_available(
    dataset,
    *,
    service: str | None = None,
    interval: int = 300,
    timeout: int = 6 * 3600,
    verbose: bool = True,
):
    """Block until ``dataset``'s servers are serving again.

    This is the loop you would otherwise write yourself around a materialization,
    with a poll interval that will not get you rate-limited. It is a separate,
    explicit call rather than an automatic retry precisely because it can take
    hours: a library that quietly sleeps that long, pretending to work, is worse
    than one that raises and lets you decide.

    Parameters
    ----------
    dataset :   str | Dataset
    service :   str, optional
                Wait only for this service (e.g. ``"materialize"``). By default,
                waits until *every* service the dataset needs is answering.
    interval :  int
                Seconds between polls. Floored at 30 - a materialization does not
                finish in less time than that, and hammering a struggling server is
                not neighbourly.
    timeout :   int
                Give up after this many seconds. Default 6 hours.

    Returns
    -------
    pandas.DataFrame
                The final (healthy) status table.

    Raises
    ------
    TimeoutError
                If the servers are still down when ``timeout`` expires.
    """
    interval = max(30, int(interval))
    deadline = time.monotonic() + timeout
    n = 0

    while True:
        df = server_status(dataset)

        rows = df
        if service:
            rows = df[df["service"] == service]
            if rows.empty:
                # Otherwise a typo filters the table down to nothing, nothing is
                # down, and we cheerfully report "healthy" and return - having
                # waited for a service that does not exist.
                raise ValueError(
                    f"{_name(dataset)} has no service {service!r}. "
                    f"Known: {', '.join(df['service'])}."
                )

        # NO AUTH means the service is up and your token is the problem; waiting will
        # not fix that, so it does not count as "still down".
        down = rows[~rows["status"].isin([OK, NO_AUTH])]

        if down.empty:
            if verbose and n:
                print(f"\n{_name(dataset)} is back ({_hms(n * interval)} waited).")
            return df

        n += 1
        left = deadline - time.monotonic()
        if left <= 0:
            raise TimeoutError(
                f"{_name(dataset)}: still {', '.join(sorted(set(down['status'])))} "
                f"after {_hms(timeout)} "
                f"({', '.join(down['service'])}). Not waiting any longer."
            )

        if verbose:
            print(
                f"[{time.strftime('%H:%M:%S')}] "
                f"{', '.join(f'{r.service}={r.status}' for r in down.itertuples())} "
                f"- retrying in {interval}s "
                f"(giving up in {_hms(left)})",
                flush=True,
            )
        time.sleep(min(interval, max(1, left)))

def cache_dir(*parts) -> Path:
    d = CACHE_ROOT.joinpath(*[_safe(p) for p in parts])
    d.mkdir(parents=True, exist_ok=True)
    return d

def clear(dataset: str | None = None):
    """Delete cached data - for one dataset, or all of it."""
    target = CACHE_ROOT / dataset if dataset else CACHE_ROOT
    if target.exists():
        shutil.rmtree(target)

def size() -> int:
    """Total size of the cache in bytes."""
    if not CACHE_ROOT.exists():
        return 0
    return sum(f.stat().st_size for f in CACHE_ROOT.rglob("*") if f.is_file())

def enabled() -> bool:
    return not os.environ.get("CONNECTO_NO_CACHE")

def download(url: str, *, force: bool = False, verbose: bool = True) -> Path:
    """Download a URL into the cache and return the local path.

    The filename is derived from a hash of the *whole* URL, not its basename, so
    two sources that both serve ``annotations.tsv`` don't collide.
    """
    suffix = "".join(Path(urlparse(url).path).suffixes[-1:]) or ".dat"
    fp = cache_dir("_http") / f"{_hash(url)}{suffix}"

    if fp.exists() and not force and enabled():
        return fp

    if verbose:
        print(f"Caching {Path(urlparse(url).path).name} from {urlparse(url).netloc}...",
              end=" ", flush=True)

    with requests.get(url, allow_redirects=True, stream=True, timeout=60) as r:
        r.raise_for_status()
        tmp = fp.with_suffix(fp.suffix + ".part")
        with open(tmp, "wb") as f:  # always binary: text mode corrupts anything else
            for chunk in r.iter_content(chunk_size=1 << 20):
                f.write(chunk)
        tmp.replace(fp)

    if verbose:
        print("Done.")
    return fp