mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-14 07:42:28 +00:00
feat: add persistent OAuth token cache and session APIs (#4182)
Stacked on #4173 (diff includes it until that merges; will rebase after). Addresses the token-cache part of [Colin's review](https://github.com/lancedb/lancedb/pull/4173#issuecomment-5674048100). Adds an explicit, opt-in persistent OAuth token cache shared by Rust, Python, and Node clients, plus `login` / `status` / `logout` session APIs, so short-lived processes (CLIs, scripts, notebooks) reuse one session instead of restarting a browser or device flow on every start. - **Opt-in and minimal**: existing callers stay memory-only and lazy. Only refresh tokens are persisted (never access tokens, never client secrets), so there are no local token-expiry decisions to get wrong when clocks move. Each process start performs one silent refresh grant. - **Hardened file backend**: private directory (`0700`), per-record files (`0600`), owner validation, symlink rejection, and atomic `rename` replacement. Corrupt, truncated, unknown-version, or permission-invalid records fail with actionable errors naming the file. Native keyring backends were evaluated (keyring crate routes Linux through D-Bus/zbus: heavy deps, headless/CI flakiness) and are deferred; the file store is the explicit opt-in, not a downgrade from a keyring. - **Cache key**: SHA-256 of the canonical identity (issuer, client ID, sorted/de-duplicated scopes, flow, public/confidential), so no secret appears in a filename and distinct identities never collide. Versioned record schema (`version: 1`). One record per identity: last login wins, documented. - **Cross-process rotation locking**: per-key `fs4` file lock (`flock` / `LockFileEx`) around the refresh critical section — acquire, reread the durable record, refresh exactly once, atomically store the rotated refresh token, release. The OS releases locks on process death, so crashes cannot strand stale locks. Only confirmed `invalid_grant`/`invalid_token` deletes a record and reauthenticates; transport, 5xx, 429, and parse failures retain it. - **Session APIs**: `OAuthSession::login/status/logout` in Rust, `lancedb.remote.OAuthSession` (async) in Python, `OAuthSession` class in Node. `status` returns non-secret metadata only. `logout` removes only the local credential — provider revocation (RFC 7009) is a deliberate follow-up, and local logout never terminates browser SSO. Azure managed identity is rejected for persistence (machine identity stays in memory); client credentials have nothing refreshable to persist and stay memory-only. - No CLI binary exists in this repo, so this ships library APIs plus doc examples in all three languages. Tests: Rust unit + mock-IdP integration (cache-key canonicalization/separation, record versioning/corruption/truncation/symlink/owner/perms, lock serialization + release, two concurrent providers proving no `invalid_grant` and correct rotation, transient-failure retention, `invalid_grant` delete + reauthenticate, login/status/logout lifecycle, client-credentials no-op, IMDS rejection, secret redaction); Python lifecycle + a true two-subprocess cross-process reuse test (second process refreshes once, never hits the device endpoint); Node lifecycle + device-flow login test. Local builds were skipped in development; CI validates all bindings. --------- Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
Jack Ye
Xuanwo
@@ -267,6 +267,33 @@ class JobInfo:
|
||||
@property
|
||||
def created_at_millis(self) -> int: ...
|
||||
|
||||
class SessionStatus:
|
||||
@property
|
||||
def refreshable(self) -> bool: ...
|
||||
@property
|
||||
def issuer_url(self) -> str: ...
|
||||
@property
|
||||
def client_id(self) -> str: ...
|
||||
@property
|
||||
def scopes(self) -> List[str]: ...
|
||||
@property
|
||||
def flow(self) -> str: ...
|
||||
@property
|
||||
def obtained_at(self) -> Optional[int]: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class SessionLogout:
|
||||
@property
|
||||
def removed(self) -> bool: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class OAuthSession:
|
||||
def __init__(self, config: Any) -> None: ...
|
||||
async def login(self) -> SessionStatus: ...
|
||||
async def status(self) -> SessionStatus: ...
|
||||
async def logout(self) -> SessionLogout: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class JobFailureInfo:
|
||||
@property
|
||||
def phase(self) -> Optional[str]: ...
|
||||
|
||||
@@ -9,7 +9,7 @@ from typing import List, Optional
|
||||
from lancedb import __version__
|
||||
|
||||
from .header import HeaderProvider
|
||||
from .oauth import OAuthConfig, OAuthFlowType
|
||||
from .oauth import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
|
||||
|
||||
# The API reference renders this module with a single mkdocstrings directive,
|
||||
# which only picks up names listed here. New public names must be added to this
|
||||
@@ -22,6 +22,8 @@ __all__ = [
|
||||
"HeaderProvider",
|
||||
"OAuthConfig",
|
||||
"OAuthFlowType",
|
||||
"OAuthSession",
|
||||
"TokenCacheOptions",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -12,10 +12,49 @@ class OAuthFlowType(str, Enum):
|
||||
CLIENT_CREDENTIALS = "client_credentials"
|
||||
"""Client Credentials grant (service-to-service / M2M)."""
|
||||
|
||||
AUTHORIZATION_CODE = "authorization_code"
|
||||
"""Interactive Authorization Code grant, using PKCE by default."""
|
||||
|
||||
DEVICE_CODE = "device_code"
|
||||
"""Device Authorization grant for CLI and headless environments."""
|
||||
|
||||
AZURE_MANAGED_IDENTITY = "azure_managed_identity"
|
||||
"""Azure Managed Identity via IMDS."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenCacheOptions:
|
||||
"""Options for the persistent OAuth token cache.
|
||||
|
||||
The cache is opt-in: it is only used when set as ``token_cache`` on
|
||||
:class:`OAuthConfig`. Only refresh tokens are persisted, in a private
|
||||
directory with owner-only permissions, so short-lived processes can reuse
|
||||
an authenticated session instead of re-prompting on every start.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cache_dir : Optional[str]
|
||||
Directory that holds cached credentials. Defaults to
|
||||
``$XDG_CACHE_HOME/lancedb/oauth``, ``$HOME/.cache/lancedb/oauth`` on
|
||||
Unix, or ``%LOCALAPPDATA%\\lancedb\\oauth`` on Windows. The directory
|
||||
is created with owner-only permissions (``0700``) when missing.
|
||||
lock_timeout_secs : Optional[int]
|
||||
How long to wait for the cross-process refresh lock before failing
|
||||
(default: 30 seconds).
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> opts = TokenCacheOptions(cache_dir="/tmp/my-app/oauth-cache")
|
||||
|
||||
Multiple identities (issuer, client, scopes, flow, client
|
||||
authentication) get separate cache entries. Within one identity the most
|
||||
recent login wins.
|
||||
"""
|
||||
|
||||
cache_dir: Optional[str] = None
|
||||
lock_timeout_secs: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OAuthConfig:
|
||||
"""OAuth configuration for LanceDB authentication.
|
||||
@@ -38,12 +77,23 @@ class OAuthConfig:
|
||||
Authentication flow to use. Default: CLIENT_CREDENTIALS.
|
||||
client_secret : Optional[str]
|
||||
Client secret (required for CLIENT_CREDENTIALS).
|
||||
redirect_uri : Optional[str]
|
||||
Loopback redirect URI for AUTHORIZATION_CODE. The default is
|
||||
``http://127.0.0.1:{callback_port}/callback``.
|
||||
callback_port : Optional[int]
|
||||
Port for the AUTHORIZATION_CODE loopback callback server (default: 8400).
|
||||
use_pkce : bool
|
||||
Protect AUTHORIZATION_CODE with S256 PKCE (default: True).
|
||||
managed_identity_client_id : Optional[str]
|
||||
Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
|
||||
refresh_buffer_secs : Optional[int]
|
||||
Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
Keep this well below the token TTL; if it is greater than or equal to
|
||||
the TTL, each request refreshes the token.
|
||||
token_cache : Optional[TokenCacheOptions]
|
||||
Opt in to the persistent token cache so short-lived processes reuse
|
||||
one session. Only supported by AUTHORIZATION_CODE and DEVICE_CODE;
|
||||
azure managed identity is rejected. Default: None (memory only).
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -64,6 +114,29 @@ class OAuthConfig:
|
||||
... scopes=["api://lancedb-api/.default"],
|
||||
... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY,
|
||||
... )
|
||||
|
||||
Authorization Code with PKCE:
|
||||
|
||||
The authorization URL is written to standard error before LanceDB tries to
|
||||
open a browser, so it can be copied in headless environments.
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... scopes=["openid", "api://lancedb-api/access"],
|
||||
... flow=OAuthFlowType.AUTHORIZATION_CODE,
|
||||
... )
|
||||
|
||||
Device Authorization with a persistent cache, so later processes reuse
|
||||
the session without a new device prompt:
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... scopes=["openid", "offline_access", "api://lancedb-api/access"],
|
||||
... flow=OAuthFlowType.DEVICE_CODE,
|
||||
... token_cache=TokenCacheOptions(),
|
||||
... )
|
||||
"""
|
||||
|
||||
issuer_url: str
|
||||
@@ -71,5 +144,70 @@ class OAuthConfig:
|
||||
scopes: List[str]
|
||||
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
|
||||
client_secret: Optional[str] = field(default=None, repr=False)
|
||||
redirect_uri: Optional[str] = None
|
||||
callback_port: Optional[int] = None
|
||||
use_pkce: bool = True
|
||||
managed_identity_client_id: Optional[str] = None
|
||||
refresh_buffer_secs: Optional[int] = None
|
||||
token_cache: Optional[TokenCacheOptions] = None
|
||||
|
||||
|
||||
class OAuthSession:
|
||||
"""Explicit OAuth session lifecycle for the persistent token cache.
|
||||
|
||||
Built from the same :class:`OAuthConfig` used for
|
||||
:func:`lancedb.connect_async` (including its ``token_cache`` options).
|
||||
A connection created with the same configuration shares the cache, so
|
||||
logging in here prepares tokens for later processes without any database
|
||||
request.
|
||||
|
||||
``login`` always runs the configured interactive flow and replaces the
|
||||
cached session (the most recent login wins). ``logout`` removes only the
|
||||
local credential; it does not revoke anything with the provider and does
|
||||
not sign out of a browser SSO session.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://issuer.example.com",
|
||||
... client_id="my-app",
|
||||
... scopes=["openid", "offline_access"],
|
||||
... flow=OAuthFlowType.DEVICE_CODE,
|
||||
... token_cache=TokenCacheOptions(),
|
||||
... )
|
||||
>>> session = OAuthSession(config) # doctest: +SKIP
|
||||
>>> status = await session.login() # doctest: +SKIP
|
||||
>>> status.refreshable # doctest: +SKIP
|
||||
True
|
||||
"""
|
||||
|
||||
def __init__(self, config: OAuthConfig):
|
||||
from lancedb._lancedb import OAuthSession as PyOAuthSession
|
||||
|
||||
self._inner: PyOAuthSession = PyOAuthSession(config)
|
||||
|
||||
async def login(self):
|
||||
"""Eagerly run the configured flow and store the session.
|
||||
|
||||
Returns a :class:`lancedb._lancedb.SessionStatus` describing the
|
||||
cached session. A successful login always replaces any prior cached
|
||||
session for this identity; if the provider does not issue a refresh
|
||||
token (for example without ``offline_access``), the previous record
|
||||
is removed and ``refreshable`` is ``False``.
|
||||
"""
|
||||
return await self._inner.login()
|
||||
|
||||
async def status(self):
|
||||
"""Report whether a cached session exists, with safe metadata.
|
||||
|
||||
Never contacts the identity provider and never exposes token values.
|
||||
"""
|
||||
return await self._inner.status()
|
||||
|
||||
async def logout(self):
|
||||
"""Remove the matching local cached credential.
|
||||
|
||||
Returns a :class:`lancedb._lancedb.SessionLogout` whose ``removed``
|
||||
flag reports whether a credential existed. Logout is idempotent.
|
||||
"""
|
||||
return await self._inner.logout()
|
||||
|
||||
Reference in New Issue
Block a user