Set up per-user OAuth

from cube import config
from datetime import datetime, timezone
import os
import time

# Don't hand the driver a token that is about to expire. Drivers cache their
# connection settings, and the connection pool opens new sessions long after
# the driver was built, so "valid right now" is not enough.
EXPIRY_SKEW_SECONDS = 120


def _parse_expiry(value):
    if not value:
        return None
    try:
        parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
    except ValueError:
        return None
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.timestamp()


def _credential(ctx: dict):
    """Resolve the credential for this request, plus the key identifying it.

    driver_factory and context_to_orchestrator_id both call this, so the cached
    connection can never drift from the token it was built with.
    """
    # For other data sources, swap "databricks" for "snowflake", etc.
    cube_cloud = (ctx.get("securityContext") or {}).get("cubeCloud") or {}
    creds = (cube_cloud.get("userCredentials") or {}).get("databricks") or {}

    access_token = creds.get("accessToken")
    expires_at = _parse_expiry(creds.get("accessTokenExpiresAt"))

    # Gate on the expiry rather than on `status`: a failed background refresh
    # can flag the record while the token already in hand is still valid, and
    # treating that as fatal drops the user onto the service account for no
    # reason.
    if access_token and expires_at and expires_at > time.time() + EXPIRY_SKEW_SECONDS:
        # Key on the expiry, never on the token itself — the orchestrator ID is
        # used as a cache prefix and appears in logs.
        return access_token, f"u{int(expires_at)}"

    return os.environ["CUBEJS_DB_DATABRICKS_TOKEN"], "service-account"


@config("driver_factory")
def driver_factory(ctx: dict) -> dict:
    token, _ = _credential(ctx)

    return {
        "type": "databricks-jdbc",
        "url": os.environ["CUBEJS_DB_DATABRICKS_URL"],
        "token": token,
        "acceptPolicy": True,
        "catalog": os.environ.get("CUBEJS_DB_DATABRICKS_CATALOG"),
    }


@config("context_to_orchestrator_id")
def context_to_orchestrator_id(ctx: dict) -> str:
    # One orchestrator per (user, credential): separate DB connections,
    # execution queues and pre-aggregation caches, and a cache key that changes
    # when the token rotates.
    cube_cloud = (ctx.get("securityContext") or {}).get("cubeCloud") or {}
    username = cube_cloud.get("username") or "default"
    _, cache_key = _credential(ctx)

    return f"CUBE_APP_{username}_{cache_key}"