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>
2026-09-22 04:55:39 +00:00 · 2026-09-16 02:01:10 +08:00
parent 575286922b
commit 2f88b71c21
25 changed files with 5464 additions and 97 deletions
@@ -170,7 +170,14 @@ export {
  TokenResponse,
 } from "./header";

-export { OAuthConfig, OAuthFlowType } from "./oauth";
+export {
+  OAuthConfig,
+  OAuthFlowType,
+  OAuthSession,
+  SessionLogout,
+  SessionStatus,
+  TokenCacheOptions,
+} from "./oauth";

 export { MergeInsertBuilder, WriteExecutionOptions } from "./merge";

@@ -1,16 +1,52 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors

+import {
+  OAuthConfig as NativeOAuthConfig,
+  OAuthSession as NativeOAuthSession,
+} from "./native";
+
 /**
 * OAuth authentication flow types.
 */
 export enum OAuthFlowType {
  /** Client Credentials grant (service-to-service / M2M). */
  ClientCredentials = "client_credentials",
+  /** Interactive Authorization Code grant, using PKCE by default. */
+  AuthorizationCode = "authorization_code",
+  /** Device Authorization grant for CLI and headless environments. */
+  DeviceCode = "device_code",
  /** Azure Managed Identity via IMDS. */
  AzureManagedIdentity = "azure_managed_identity",
 }

+/**
+ * Options for the persistent OAuth token cache.
+ *
+ * The cache is opt-in: it is only used when set as `tokenCache` on
+ * {@link 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.
+ *
+ * Multiple identities (issuer, client, scopes, flow, client authentication)
+ * get separate cache entries. Within one identity the most recent login wins.
+ */
+export interface TokenCacheOptions {
+  /**
+   * 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.
+   */
+  cacheDir?: string;
+
+  /**
+   * How long to wait for the cross-process refresh lock before failing, in
+   * seconds (default: 30).
+   */
+  lockTimeoutSecs?: number;
+}
+
 /**
 * OAuth configuration for LanceDB authentication.
 *
@@ -40,6 +76,21 @@ export enum OAuthFlowType {
 *   flow: OAuthFlowType.AzureManagedIdentity,
 * };
 * ```
+ *
+ * @example Authorization Code with PKCE:
+ * The authorization URL is written to stderr before LanceDB tries to open a
+ * browser, so it can be copied in headless environments.
+ * ```typescript
+ * const config: OAuthConfig = {
+ *   issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
+ *   clientId: "app-id",
+ *   scopes: ["openid", "api://lancedb-api/access"],
+ *   flow: OAuthFlowType.AuthorizationCode,
+ * };
+ * ```
+ *
+ * Device Authorization writes the verification URL and user code to stderr
+ * before polling begins.
 */
 export interface OAuthConfig {
  /**
@@ -64,6 +115,15 @@ export interface OAuthConfig {
  /** Client secret (required for ClientCredentials). */
  clientSecret?: string;

+  /** Loopback redirect URI for AuthorizationCode. */
+  redirectUri?: string;
+
+  /** Port for the AuthorizationCode loopback callback server (default: 8400). */
+  callbackPort?: number;
+
+  /** Protect AuthorizationCode with S256 PKCE (default: true). */
+  usePkce?: boolean;
+
  /** Client ID for user-assigned managed identity (AzureManagedIdentity). */
  managedIdentityClientId?: string;

@@ -73,4 +133,119 @@ export interface OAuthConfig {
   * the TTL, each request refreshes the token.
   */
  refreshBufferSecs?: number;
+
+  /**
+   * Opt in to the persistent token cache so short-lived processes reuse one
+   * session. Only refresh tokens are persisted. Only supported by
+   * AuthorizationCode and DeviceCode; Azure managed identity is rejected.
+   * Default: unset (memory only).
+   */
+  tokenCache?: TokenCacheOptions;
+}
+
+/**
+ * Safe, non-secret view of a cached OAuth session, returned by
+ * {@link OAuthSession.status} and {@link OAuthSession.login}.
+ */
+export interface SessionStatus {
+  /**
+   * Whether a cached session exists that can obtain tokens without
+   * interactive authentication. Because access tokens are not persisted,
+   * this is `true` exactly when a refresh token is cached; the next
+   * connection refreshes with it rather than opening a browser or device
+   * prompt.
+   */
+  refreshable: boolean;
+
+  /** Canonical issuer URL of the cached session. */
+  issuerUrl: string;
+
+  /** Client ID of the cached session. */
+  clientId: string;
+
+  /** Canonical (sorted, de-duplicated) scope set of the cached session. */
+  scopes: string[];
+
+  /** Flow that produced the cached session. */
+  flow: string;
+
+  /** When the cached session was obtained, as Unix seconds. */
+  obtainedAt?: number;
+}
+
+/** Result of {@link OAuthSession.logout}. */
+export interface SessionLogout {
+  /**
+   * Whether a cached credential was removed. `false` means no matching
+   * session was cached; logout is idempotent.
+   */
+  removed: boolean;
+}
+
+/**
+ * Explicit OAuth session lifecycle for the persistent token cache: eager
+ * `login`, non-secret `status`, and local `logout`.
+ *
+ * A session is built from the same {@link OAuthConfig} used to connect
+ * (including its `tokenCache` 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.
+ *
+ * @example
+ * ```typescript
+ * const config: OAuthConfig = {
+ *   issuerUrl: "https://issuer.example.com",
+ *   clientId: "my-app",
+ *   scopes: ["openid", "offline_access"],
+ *   flow: OAuthFlowType.DeviceCode,
+ *   tokenCache: { cacheDir: "/tmp/my-app/oauth-cache" },
+ * };
+ * const session = new OAuthSession(config);
+ * const status = await session.login();
+ * ```
+ */
+export class OAuthSession {
+  private readonly inner: NativeOAuthSession;
+
+  /** Create a session manager for the given OAuth configuration. */
+  constructor(config: OAuthConfig) {
+    this.inner = new NativeOAuthSession(config as unknown as NativeOAuthConfig);
+  }
+
+  /**
+   * Eagerly run the configured authentication flow and store the 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 the status
+   * reports `refreshable == false`.
+   */
+  async login(): Promise<SessionStatus> {
+    return this.inner.login();
+  }
+
+  /**
+   * Report whether a matching cached session exists, with safe metadata.
+   *
+   * This never contacts the identity provider and never exposes token values.
+   */
+  async status(): Promise<SessionStatus> {
+    return this.inner.status();
+  }
+
+  /**
+   * Remove the matching local cached credential.
+   *
+   * This only deletes the local cache entry. It does not revoke the refresh
+   * token with the provider and does not sign out of a browser SSO session.
+   * Repeated calls succeed; `removed` reports whether a credential existed.
+   */
+  async logout(): Promise<SessionLogout> {
+    return this.inner.logout();
+  }
 }