mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-22 04:55:39 +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
Generated
+84
@@ -3088,6 +3088,16 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dispatch2"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"objc2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "displaydoc"
|
||||
version = "0.2.5"
|
||||
@@ -3483,6 +3493,16 @@ version = "1.20260804.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "82eb03a32a1d50555353c85a7b9d3279a6f1e91af9890b789acdf544ed57c8d7"
|
||||
|
||||
[[package]]
|
||||
name = "fs4"
|
||||
version = "0.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4"
|
||||
dependencies = [
|
||||
"rustix",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs_extra"
|
||||
version = "1.3.0"
|
||||
@@ -5477,6 +5497,7 @@ dependencies = [
|
||||
"aws-sdk-s3",
|
||||
"aws-smithy-runtime",
|
||||
"aws-smithy-types",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"candle-core",
|
||||
"candle-nn",
|
||||
@@ -5491,6 +5512,7 @@ dependencies = [
|
||||
"datafusion-physical-expr",
|
||||
"datafusion-physical-plan",
|
||||
"datafusion-sql",
|
||||
"fs4",
|
||||
"futures",
|
||||
"half",
|
||||
"hf-hub",
|
||||
@@ -5543,6 +5565,7 @@ dependencies = [
|
||||
"urlencoding",
|
||||
"uuid",
|
||||
"walkdir",
|
||||
"webbrowser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6227,6 +6250,12 @@ dependencies = [
|
||||
"rawpointer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ndk-context"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b"
|
||||
|
||||
[[package]]
|
||||
name = "nibble_vec"
|
||||
version = "0.1.0"
|
||||
@@ -6409,6 +6438,26 @@ version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
|
||||
|
||||
[[package]]
|
||||
name = "objc2"
|
||||
version = "0.6.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f"
|
||||
dependencies = [
|
||||
"objc2-encode",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-app-kit"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"objc2",
|
||||
"objc2-foundation",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-core-foundation"
|
||||
version = "0.3.2"
|
||||
@@ -6416,6 +6465,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"dispatch2",
|
||||
"objc2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "objc2-encode"
|
||||
version = "4.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33"
|
||||
|
||||
[[package]]
|
||||
name = "objc2-foundation"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"objc2",
|
||||
"objc2-core-foundation",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -10788,6 +10856,22 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webbrowser"
|
||||
version = "1.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "62c35be770821a214dbc362fc26908c853e776c0004294d0b10b8a6bad582f94"
|
||||
dependencies = [
|
||||
"jni",
|
||||
"log",
|
||||
"ndk-context",
|
||||
"objc2",
|
||||
"objc2-app-kit",
|
||||
"objc2-foundation",
|
||||
"url",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "webpki-root-certs"
|
||||
version = "1.0.7"
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / OAuthSession
|
||||
|
||||
# Class: OAuthSession
|
||||
|
||||
Explicit OAuth session lifecycle for the persistent token cache: eager
|
||||
`login`, non-secret `status`, and local `logout`.
|
||||
|
||||
A session is built from the same [OAuthConfig](../interfaces/OAuthConfig.md) 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();
|
||||
```
|
||||
|
||||
## Constructors
|
||||
|
||||
### new OAuthSession()
|
||||
|
||||
```ts
|
||||
new OAuthSession(config): OAuthSession
|
||||
```
|
||||
|
||||
Create a session manager for the given OAuth configuration.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **config**: [`OAuthConfig`](../interfaces/OAuthConfig.md)
|
||||
|
||||
#### Returns
|
||||
|
||||
[`OAuthSession`](OAuthSession.md)
|
||||
|
||||
## Methods
|
||||
|
||||
### login()
|
||||
|
||||
```ts
|
||||
login(): Promise<SessionStatus>
|
||||
```
|
||||
|
||||
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`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`SessionStatus`](../interfaces/SessionStatus.md)>
|
||||
|
||||
***
|
||||
|
||||
### logout()
|
||||
|
||||
```ts
|
||||
logout(): Promise<SessionLogout>
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`SessionLogout`](../interfaces/SessionLogout.md)>
|
||||
|
||||
***
|
||||
|
||||
### status()
|
||||
|
||||
```ts
|
||||
status(): Promise<SessionStatus>
|
||||
```
|
||||
|
||||
Report whether a matching cached session exists, with safe metadata.
|
||||
|
||||
This never contacts the identity provider and never exposes token values.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`SessionStatus`](../interfaces/SessionStatus.md)>
|
||||
@@ -10,6 +10,16 @@ OAuth authentication flow types.
|
||||
|
||||
## Enumeration Members
|
||||
|
||||
### AuthorizationCode
|
||||
|
||||
```ts
|
||||
AuthorizationCode: "authorization_code";
|
||||
```
|
||||
|
||||
Interactive Authorization Code grant, using PKCE by default.
|
||||
|
||||
***
|
||||
|
||||
### AzureManagedIdentity
|
||||
|
||||
```ts
|
||||
@@ -27,3 +37,13 @@ ClientCredentials: "client_credentials";
|
||||
```
|
||||
|
||||
Client Credentials grant (service-to-service / M2M).
|
||||
|
||||
***
|
||||
|
||||
### DeviceCode
|
||||
|
||||
```ts
|
||||
DeviceCode: "device_code";
|
||||
```
|
||||
|
||||
Device Authorization grant for CLI and headless environments.
|
||||
|
||||
@@ -35,6 +35,7 @@
|
||||
- [MultiMatchQuery](classes/MultiMatchQuery.md)
|
||||
- [NativeJsHeaderProvider](classes/NativeJsHeaderProvider.md)
|
||||
- [OAuthHeaderProvider](classes/OAuthHeaderProvider.md)
|
||||
- [OAuthSession](classes/OAuthSession.md)
|
||||
- [PermutationBuilder](classes/PermutationBuilder.md)
|
||||
- [PhraseQuery](classes/PhraseQuery.md)
|
||||
- [Query](classes/Query.md)
|
||||
@@ -122,6 +123,8 @@
|
||||
- [RestNamespaceConfig](interfaces/RestNamespaceConfig.md)
|
||||
- [RetryConfig](interfaces/RetryConfig.md)
|
||||
- [ScannableOptions](interfaces/ScannableOptions.md)
|
||||
- [SessionLogout](interfaces/SessionLogout.md)
|
||||
- [SessionStatus](interfaces/SessionStatus.md)
|
||||
- [ShuffleOptions](interfaces/ShuffleOptions.md)
|
||||
- [SplitCalculatedOptions](interfaces/SplitCalculatedOptions.md)
|
||||
- [SplitHashOptions](interfaces/SplitHashOptions.md)
|
||||
@@ -131,6 +134,7 @@
|
||||
- [TableStatistics](interfaces/TableStatistics.md)
|
||||
- [TimeoutConfig](interfaces/TimeoutConfig.md)
|
||||
- [TlsConfig](interfaces/TlsConfig.md)
|
||||
- [TokenCacheOptions](interfaces/TokenCacheOptions.md)
|
||||
- [TokenResponse](interfaces/TokenResponse.md)
|
||||
- [TokenizeOptions](interfaces/TokenizeOptions.md)
|
||||
- [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md)
|
||||
|
||||
@@ -15,6 +15,16 @@ All token acquisition and refresh is handled in the Rust layer.
|
||||
|
||||
## Properties
|
||||
|
||||
### callbackPort?
|
||||
|
||||
```ts
|
||||
optional callbackPort: number;
|
||||
```
|
||||
|
||||
Port for the authorization_code loopback callback server.
|
||||
|
||||
***
|
||||
|
||||
### clientId
|
||||
|
||||
```ts
|
||||
@@ -41,7 +51,8 @@ Client secret (required for client_credentials).
|
||||
optional flow: string;
|
||||
```
|
||||
|
||||
Authentication flow: "client_credentials" or "azure_managed_identity"
|
||||
Authentication flow: "client_credentials", "authorization_code",
|
||||
"device_code", or "azure_managed_identity"
|
||||
|
||||
***
|
||||
|
||||
@@ -66,6 +77,16 @@ Client ID for user-assigned managed identity (azure_managed_identity).
|
||||
|
||||
***
|
||||
|
||||
### redirectUri?
|
||||
|
||||
```ts
|
||||
optional redirectUri: string;
|
||||
```
|
||||
|
||||
Loopback redirect URI for authorization_code.
|
||||
|
||||
***
|
||||
|
||||
### refreshBufferSecs?
|
||||
|
||||
```ts
|
||||
@@ -86,3 +107,24 @@ scopes: string[];
|
||||
|
||||
OAuth scopes to request. For Azure managed identity, exactly one scope
|
||||
or resource is required. For example: `["api://{app_id}/.default"]`
|
||||
|
||||
***
|
||||
|
||||
### tokenCache?
|
||||
|
||||
```ts
|
||||
optional tokenCache: TokenCacheOptions;
|
||||
```
|
||||
|
||||
Opt in to the persistent token cache so short-lived processes reuse
|
||||
one session. Only refresh tokens are persisted.
|
||||
|
||||
***
|
||||
|
||||
### usePkce?
|
||||
|
||||
```ts
|
||||
optional usePkce: boolean;
|
||||
```
|
||||
|
||||
Whether authorization_code uses S256 PKCE (default: true).
|
||||
|
||||
@@ -35,8 +35,32 @@ const config: OAuthConfig = {
|
||||
};
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
## Properties
|
||||
|
||||
### callbackPort?
|
||||
|
||||
```ts
|
||||
optional callbackPort: number;
|
||||
```
|
||||
|
||||
Port for the AuthorizationCode loopback callback server (default: 8400).
|
||||
|
||||
***
|
||||
|
||||
### clientId
|
||||
|
||||
```ts
|
||||
@@ -88,6 +112,16 @@ Client ID for user-assigned managed identity (AzureManagedIdentity).
|
||||
|
||||
***
|
||||
|
||||
### redirectUri?
|
||||
|
||||
```ts
|
||||
optional redirectUri: string;
|
||||
```
|
||||
|
||||
Loopback redirect URI for AuthorizationCode.
|
||||
|
||||
***
|
||||
|
||||
### refreshBufferSecs?
|
||||
|
||||
```ts
|
||||
@@ -109,3 +143,26 @@ scopes: string[];
|
||||
OAuth scopes to request.
|
||||
For Azure managed identity, exactly one scope or resource is required.
|
||||
For example: `["api://{app_id}/.default"]`
|
||||
|
||||
***
|
||||
|
||||
### tokenCache?
|
||||
|
||||
```ts
|
||||
optional tokenCache: TokenCacheOptions;
|
||||
```
|
||||
|
||||
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).
|
||||
|
||||
***
|
||||
|
||||
### usePkce?
|
||||
|
||||
```ts
|
||||
optional usePkce: boolean;
|
||||
```
|
||||
|
||||
Protect AuthorizationCode with S256 PKCE (default: true).
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / SessionLogout
|
||||
|
||||
# Interface: SessionLogout
|
||||
|
||||
Result of [OAuthSession.logout](../classes/OAuthSession.md#logout).
|
||||
|
||||
## Properties
|
||||
|
||||
### removed
|
||||
|
||||
```ts
|
||||
removed: boolean;
|
||||
```
|
||||
|
||||
Whether a cached credential was removed. `false` means no matching
|
||||
session was cached; logout is idempotent.
|
||||
@@ -0,0 +1,74 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / SessionStatus
|
||||
|
||||
# Interface: SessionStatus
|
||||
|
||||
Safe, non-secret view of a cached OAuth session, returned by
|
||||
[OAuthSession.status](../classes/OAuthSession.md#status) and [OAuthSession.login](../classes/OAuthSession.md#login).
|
||||
|
||||
## Properties
|
||||
|
||||
### clientId
|
||||
|
||||
```ts
|
||||
clientId: string;
|
||||
```
|
||||
|
||||
Client ID of the cached session.
|
||||
|
||||
***
|
||||
|
||||
### flow
|
||||
|
||||
```ts
|
||||
flow: string;
|
||||
```
|
||||
|
||||
Flow that produced the cached session.
|
||||
|
||||
***
|
||||
|
||||
### issuerUrl
|
||||
|
||||
```ts
|
||||
issuerUrl: string;
|
||||
```
|
||||
|
||||
Canonical issuer URL of the cached session.
|
||||
|
||||
***
|
||||
|
||||
### obtainedAt?
|
||||
|
||||
```ts
|
||||
optional obtainedAt: number;
|
||||
```
|
||||
|
||||
When the cached session was obtained, as Unix seconds.
|
||||
|
||||
***
|
||||
|
||||
### refreshable
|
||||
|
||||
```ts
|
||||
refreshable: boolean;
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
***
|
||||
|
||||
### scopes
|
||||
|
||||
```ts
|
||||
scopes: string[];
|
||||
```
|
||||
|
||||
Canonical (sorted, de-duplicated) scope set of the cached session.
|
||||
@@ -0,0 +1,41 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenCacheOptions
|
||||
|
||||
# Interface: TokenCacheOptions
|
||||
|
||||
Options for the persistent OAuth token cache.
|
||||
|
||||
The cache is opt-in: it is only used when set as `tokenCache` on
|
||||
[OAuthConfig](OAuthConfig.md). 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.
|
||||
|
||||
## Properties
|
||||
|
||||
### cacheDir?
|
||||
|
||||
```ts
|
||||
optional cacheDir: string;
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
***
|
||||
|
||||
### lockTimeoutSecs?
|
||||
|
||||
```ts
|
||||
optional lockTimeoutSecs: number;
|
||||
```
|
||||
|
||||
How long to wait for the cross-process refresh lock before failing, in
|
||||
seconds (default: 30).
|
||||
@@ -0,0 +1,211 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as fs from "fs";
|
||||
import * as http from "http";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
import { OAuthConfig, OAuthFlowType, OAuthSession } from "../lancedb/oauth";
|
||||
|
||||
function tempCacheDir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "lancedb-oauth-cache-"));
|
||||
}
|
||||
|
||||
function deviceConfig(issuerUrl: string, cacheDir: string): OAuthConfig {
|
||||
return {
|
||||
issuerUrl,
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
tokenCache: { cacheDir },
|
||||
};
|
||||
}
|
||||
|
||||
describe("OAuthSession", () => {
|
||||
beforeAll(() => {
|
||||
// Point the Rust browser helper at a no-op so device-flow logins never
|
||||
// open a real browser window during tests.
|
||||
process.env.LANCEDB_OAUTH_BROWSER = "/usr/bin/true";
|
||||
});
|
||||
|
||||
it("reports an absent session and logout is idempotent", async () => {
|
||||
const cacheDir = tempCacheDir();
|
||||
const session = new OAuthSession(
|
||||
deviceConfig("https://issuer.example.com", cacheDir),
|
||||
);
|
||||
|
||||
const status = await session.status();
|
||||
expect(status.refreshable).toBe(false);
|
||||
expect(status.issuerUrl).toBe("https://issuer.example.com");
|
||||
expect(status.clientId).toBe("client-id");
|
||||
expect(status.scopes).toEqual(["openid"]);
|
||||
expect(status.flow).toBe("device_code");
|
||||
expect(status.obtainedAt).toBeUndefined();
|
||||
|
||||
const logout = await session.logout();
|
||||
expect(logout.removed).toBe(false);
|
||||
});
|
||||
|
||||
it("requires token cache options", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://issuer.example.com",
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/token/);
|
||||
});
|
||||
|
||||
it("rejects azure managed identity persistence", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://login.microsoftonline.com/tenant/v2.0",
|
||||
clientId: "app-id",
|
||||
scopes: ["api://app/.default"],
|
||||
flow: OAuthFlowType.AzureManagedIdentity,
|
||||
tokenCache: { cacheDir: tempCacheDir() },
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/AzureManagedIdentity/);
|
||||
});
|
||||
|
||||
it("logs in via device flow, caches, and logs out", async () => {
|
||||
const server = new MockIdp();
|
||||
await server.start();
|
||||
try {
|
||||
const cacheDir = tempCacheDir();
|
||||
const issuerUrl = server.issuerUrl();
|
||||
|
||||
const session = new OAuthSession(deviceConfig(issuerUrl, cacheDir));
|
||||
const status = await session.login();
|
||||
expect(status.refreshable).toBe(true);
|
||||
expect(status.obtainedAt).toBeGreaterThan(0);
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
|
||||
// An independent session (a fresh "process") sees the cached login.
|
||||
const other = new OAuthSession(deviceConfig(issuerUrl, cacheDir));
|
||||
const cached = await other.status();
|
||||
expect(cached.refreshable).toBe(true);
|
||||
|
||||
const logout = await other.logout();
|
||||
expect(logout.removed).toBe(true);
|
||||
const again = await session.logout();
|
||||
expect(again.removed).toBe(false);
|
||||
expect((await session.status()).refreshable).toBe(false);
|
||||
|
||||
// Only the initial login used the interactive device flow.
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
expect(server.state.refreshGrants).toBe(0);
|
||||
} finally {
|
||||
server.close();
|
||||
}
|
||||
}, 15000);
|
||||
});
|
||||
|
||||
/** Mock IdP with discovery, device authorization, and rotating refresh. */
|
||||
class MockIdp {
|
||||
readonly state = {
|
||||
deviceAuthorizations: 0,
|
||||
refreshGrants: 0,
|
||||
accessTokensIssued: 0,
|
||||
currentRefresh: null as string | null,
|
||||
};
|
||||
private server?: http.Server;
|
||||
private port = 0;
|
||||
|
||||
issuerUrl(): string {
|
||||
return `http://127.0.0.1:${this.port}`;
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
const server = http.createServer((req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
req.on("data", (chunk) => chunks.push(chunk));
|
||||
req.on("end", () => {
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
const params = new URLSearchParams(body);
|
||||
this.handle(req.url ?? "", params, res);
|
||||
});
|
||||
});
|
||||
await new Promise<void>((resolve) => {
|
||||
server.listen(0, "127.0.0.1", () => resolve());
|
||||
});
|
||||
const address = server.address();
|
||||
if (address && typeof address === "object") {
|
||||
this.port = address.port;
|
||||
}
|
||||
this.server = server;
|
||||
}
|
||||
|
||||
private handle(
|
||||
url: string,
|
||||
params: URLSearchParams,
|
||||
res: http.ServerResponse,
|
||||
): void {
|
||||
const respond = (status: number, payload: unknown): void => {
|
||||
const body = JSON.stringify(payload);
|
||||
res.writeHead(status, {
|
||||
"Content-Type": "application/json",
|
||||
"Content-Length": Buffer.byteLength(body),
|
||||
});
|
||||
res.end(body);
|
||||
};
|
||||
|
||||
if (url === "/.well-known/openid-configuration") {
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
token_endpoint: `${this.issuerUrl()}/token`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_authorization_endpoint: `${this.issuerUrl()}/device`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/device") {
|
||||
this.state.deviceAuthorizations += 1;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_code: "device-code",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
user_code: "ABCD-EFGH",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
verification_uri: `${this.issuerUrl()}/verify`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 60,
|
||||
interval: 1,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/token") {
|
||||
const grantType = params.get("grant_type") ?? "";
|
||||
if (grantType === "refresh_token") {
|
||||
this.state.refreshGrants += 1;
|
||||
if (params.get("refresh_token") !== this.state.currentRefresh) {
|
||||
respond(400, { error: "invalid_grant" });
|
||||
return;
|
||||
}
|
||||
} else if (!grantType.includes("device_code")) {
|
||||
respond(400, { error: "unsupported_grant_type" });
|
||||
return;
|
||||
}
|
||||
this.state.accessTokensIssued += 1;
|
||||
const number = this.state.accessTokensIssued;
|
||||
const refresh = `refresh-${number}`;
|
||||
this.state.currentRefresh = refresh;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
access_token: `access-${number}`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
refresh_token: refresh,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 3600,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
respond(404, {});
|
||||
}
|
||||
|
||||
close(): void {
|
||||
this.server?.close();
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
+222
-2
@@ -6,6 +6,8 @@ use std::collections::HashMap;
|
||||
use lancedb::error::Error;
|
||||
use napi_derive::*;
|
||||
|
||||
use crate::error::NapiErrorExt;
|
||||
|
||||
/// Timeout configuration for remote HTTP client.
|
||||
#[napi(object)]
|
||||
#[derive(Debug)]
|
||||
@@ -141,6 +143,34 @@ impl From<TlsConfig> for lancedb::remote::TlsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Options for the persistent OAuth token cache.
|
||||
///
|
||||
/// The cache is opt-in: it is only used when set as `tokenCache` on
|
||||
/// `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.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct 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.
|
||||
pub cache_dir: Option<String>,
|
||||
/// How long to wait for the cross-process refresh lock before failing,
|
||||
/// in seconds (default: 30).
|
||||
pub lock_timeout_secs: Option<u32>,
|
||||
}
|
||||
|
||||
impl From<TokenCacheOptions> for lancedb::remote::TokenCacheOptions {
|
||||
fn from(options: TokenCacheOptions) -> Self {
|
||||
Self {
|
||||
cache_dir: options.cache_dir.map(std::path::PathBuf::from),
|
||||
lock_timeout_secs: options.lock_timeout_secs.map(|secs| secs as u64),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// OAuth configuration for LanceDB authentication.
|
||||
///
|
||||
/// This is the generated napi-rs binding shape. TypeScript users should prefer
|
||||
@@ -158,16 +188,26 @@ pub struct OAuthConfig {
|
||||
/// OAuth scopes to request. For Azure managed identity, exactly one scope
|
||||
/// or resource is required. For example: `["api://{app_id}/.default"]`
|
||||
pub scopes: Vec<String>,
|
||||
/// Authentication flow: "client_credentials" or "azure_managed_identity"
|
||||
/// Authentication flow: "client_credentials", "authorization_code",
|
||||
/// "device_code", or "azure_managed_identity"
|
||||
pub flow: Option<String>,
|
||||
/// Client secret (required for client_credentials).
|
||||
pub client_secret: Option<String>,
|
||||
/// Loopback redirect URI for authorization_code.
|
||||
pub redirect_uri: Option<String>,
|
||||
/// Port for the authorization_code loopback callback server.
|
||||
pub callback_port: Option<u16>,
|
||||
/// Whether authorization_code uses S256 PKCE (default: true).
|
||||
pub use_pkce: Option<bool>,
|
||||
/// Client ID for user-assigned managed identity (azure_managed_identity).
|
||||
pub managed_identity_client_id: Option<String>,
|
||||
/// 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.
|
||||
pub refresh_buffer_secs: Option<u32>,
|
||||
/// Opt in to the persistent token cache so short-lived processes reuse
|
||||
/// one session. Only refresh tokens are persisted.
|
||||
pub token_cache: Option<TokenCacheOptions>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for OAuthConfig {
|
||||
@@ -181,11 +221,15 @@ impl std::fmt::Debug for OAuthConfig {
|
||||
"client_secret",
|
||||
&self.client_secret.as_deref().map(|_| "<redacted>"),
|
||||
)
|
||||
.field("redirect_uri", &self.redirect_uri)
|
||||
.field("callback_port", &self.callback_port)
|
||||
.field("use_pkce", &self.use_pkce)
|
||||
.field(
|
||||
"managed_identity_client_id",
|
||||
&self.managed_identity_client_id,
|
||||
)
|
||||
.field("refresh_buffer_secs", &self.refresh_buffer_secs)
|
||||
.field("token_cache", &self.token_cache)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -194,10 +238,22 @@ impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: OAuthConfig) -> Result<Self, Self::Error> {
|
||||
use lancedb::remote::oauth::OAuthFlow;
|
||||
use lancedb::remote::oauth::{AuthorizationCodeOptions, OAuthFlow};
|
||||
|
||||
let flow = match config.flow.as_deref().unwrap_or("client_credentials") {
|
||||
"client_credentials" => OAuthFlow::ClientCredentials,
|
||||
"authorization_code" => {
|
||||
let mut options =
|
||||
AuthorizationCodeOptions::new().use_pkce(config.use_pkce.unwrap_or(true));
|
||||
if let Some(redirect_uri) = config.redirect_uri {
|
||||
options = options.redirect_uri(redirect_uri);
|
||||
}
|
||||
if let Some(callback_port) = config.callback_port {
|
||||
options = options.callback_port(callback_port);
|
||||
}
|
||||
OAuthFlow::AuthorizationCode(options)
|
||||
}
|
||||
"device_code" => OAuthFlow::DeviceCode,
|
||||
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
|
||||
client_id: config.managed_identity_client_id,
|
||||
},
|
||||
@@ -215,10 +271,115 @@ impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
|
||||
scopes: config.scopes,
|
||||
flow,
|
||||
refresh_buffer_secs: config.refresh_buffer_secs.map(|v| v as u64),
|
||||
token_cache: config.token_cache.map(Into::into),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Safe, non-secret view of a cached OAuth session, returned by
|
||||
/// `OAuthSession.status()` and `OAuthSession.login()`.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionStatus {
|
||||
/// Whether a cached session exists that can obtain tokens without
|
||||
/// interactive authentication.
|
||||
pub refreshable: bool,
|
||||
/// Canonical issuer URL of the cached session.
|
||||
pub issuer_url: String,
|
||||
/// Client ID of the cached session.
|
||||
pub client_id: String,
|
||||
/// Canonical (sorted, de-duplicated) scopes of the cached session.
|
||||
pub scopes: Vec<String>,
|
||||
/// Flow that produced the cached session.
|
||||
pub flow: String,
|
||||
/// When the cached session was obtained, as Unix seconds.
|
||||
pub obtained_at: Option<f64>,
|
||||
}
|
||||
|
||||
/// Result of `OAuthSession.logout()`.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionLogout {
|
||||
/// Whether a cached credential was removed. `false` means no matching
|
||||
/// session was cached; logout is idempotent.
|
||||
pub removed: bool,
|
||||
}
|
||||
|
||||
/// Explicit OAuth session lifecycle for the persistent token cache: eager
|
||||
/// `login`, non-secret `status`, and local `logout`.
|
||||
///
|
||||
/// A session is built from the same `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.
|
||||
#[napi]
|
||||
pub struct OAuthSession {
|
||||
inner: lancedb::remote::OAuthSession,
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl OAuthSession {
|
||||
/// Create a session manager for the given OAuth configuration.
|
||||
///
|
||||
/// The configuration must enable `tokenCache` options and use a flow that
|
||||
/// supports persistent sessions (authorization code or device code).
|
||||
#[napi(constructor)]
|
||||
pub fn new(config: OAuthConfig) -> napi::Result<Self> {
|
||||
let config: lancedb::remote::oauth::OAuthConfig = config.try_into().default_error()?;
|
||||
let inner = lancedb::remote::OAuthSession::new(config).default_error()?;
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
/// 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`.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn login(&self) -> napi::Result<SessionStatus> {
|
||||
let status = self.inner.login().await.default_error()?;
|
||||
Ok(SessionStatus::from(status))
|
||||
}
|
||||
|
||||
/// Report whether a matching cached session exists, with safe metadata.
|
||||
///
|
||||
/// This never contacts the identity provider and never exposes token
|
||||
/// values.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn status(&self) -> napi::Result<SessionStatus> {
|
||||
let status = self.inner.status().await.default_error()?;
|
||||
Ok(SessionStatus::from(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.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn logout(&self) -> napi::Result<SessionLogout> {
|
||||
let logout = self.inner.logout().await.default_error()?;
|
||||
Ok(SessionLogout {
|
||||
removed: logout.removed,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::remote::SessionStatus> for SessionStatus {
|
||||
fn from(status: lancedb::remote::SessionStatus) -> Self {
|
||||
Self {
|
||||
refreshable: status.refreshable,
|
||||
issuer_url: status.issuer_url,
|
||||
client_id: status.client_id,
|
||||
scopes: status.scopes,
|
||||
flow: status.flow,
|
||||
obtained_at: status.obtained_at.map(|secs| secs as f64),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ClientConfig> for lancedb::remote::ClientConfig {
|
||||
fn from(config: ClientConfig) -> Self {
|
||||
Self {
|
||||
@@ -252,8 +413,12 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("typo".to_string()),
|
||||
client_secret: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let err = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap_err();
|
||||
@@ -272,12 +437,67 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("client_credentials".to_string()),
|
||||
client_secret: Some("super-secret".to_string()),
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let debug = format!("{config:?}");
|
||||
assert!(!debug.contains("super-secret"));
|
||||
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_authorization_code_conversion_preserves_options() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["openid".to_string()],
|
||||
flow: Some("authorization_code".to_string()),
|
||||
client_secret: Some("secret".to_string()),
|
||||
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
|
||||
callback_port: Some(9000),
|
||||
use_pkce: Some(false),
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let converted = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap();
|
||||
let lancedb::remote::oauth::OAuthFlow::AuthorizationCode(options) = converted.flow else {
|
||||
panic!("expected authorization code flow");
|
||||
};
|
||||
assert_eq!(
|
||||
options.redirect_uri.as_deref(),
|
||||
Some("http://127.0.0.1:9000/callback")
|
||||
);
|
||||
assert_eq!(options.callback_port, Some(9000));
|
||||
assert!(!options.use_pkce);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_device_code_conversion() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["openid".to_string()],
|
||||
flow: Some("device_code".to_string()),
|
||||
client_secret: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let converted = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap();
|
||||
assert!(matches!(
|
||||
converted.flow,
|
||||
lancedb::remote::oauth::OAuthFlow::DeviceCode
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
+2
-1
@@ -4,7 +4,8 @@
|
||||
"lancedb/native.d.ts:VectorQuery",
|
||||
"lancedb/native.d.ts:TakeQuery",
|
||||
"lancedb/native.d.ts:RecordBatchIterator",
|
||||
"lancedb/native.d.ts:NativeMergeInsertBuilder"
|
||||
"lancedb/native.d.ts:NativeMergeInsertBuilder",
|
||||
"lancedb/native.d.ts:TokenCacheOptions"
|
||||
],
|
||||
"useHTMLEncodedBrackets": true,
|
||||
"useCodeBlocks": true,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -47,6 +47,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<crate::oauth::PyOAuthSession>()?;
|
||||
m.add_class::<crate::oauth::PySessionStatus>()?;
|
||||
m.add_class::<crate::oauth::PySessionLogout>()?;
|
||||
m.add_class::<crate::job::Job>()?;
|
||||
m.add_class::<crate::job::JobInfo>()?;
|
||||
m.add_class::<crate::job::JobDescription>()?;
|
||||
|
||||
+246
-6
@@ -1,10 +1,33 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use pyo3::FromPyObject;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use pyo3::{FromPyObject, PyResult, Python, pyclass, pymethods};
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
use lancedb::error::Error;
|
||||
use lancedb::remote::oauth::{OAuthConfig, OAuthFlow};
|
||||
use lancedb::remote::oauth::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow};
|
||||
use lancedb::remote::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
|
||||
|
||||
/// Python-side persistent token cache options, extracted via FromPyObject.
|
||||
/// Maps to `lancedb.remote.oauth.TokenCacheOptions` Python dataclass.
|
||||
#[derive(FromPyObject, Default)]
|
||||
pub struct PyTokenCacheOptions {
|
||||
pub cache_dir: Option<String>,
|
||||
pub lock_timeout_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl From<PyTokenCacheOptions> for TokenCacheOptions {
|
||||
fn from(py: PyTokenCacheOptions) -> Self {
|
||||
TokenCacheOptions {
|
||||
cache_dir: py.cache_dir.map(PathBuf::from),
|
||||
lock_timeout_secs: py.lock_timeout_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python-side OAuth configuration, extracted via FromPyObject.
|
||||
/// Maps to `lancedb.remote.oauth.OAuthConfig` Python dataclass.
|
||||
@@ -15,8 +38,12 @@ pub struct PyOAuthConfig {
|
||||
pub scopes: Vec<String>,
|
||||
pub flow: String,
|
||||
pub client_secret: Option<String>,
|
||||
pub redirect_uri: Option<String>,
|
||||
pub callback_port: Option<u16>,
|
||||
pub use_pkce: bool,
|
||||
pub managed_identity_client_id: Option<String>,
|
||||
pub refresh_buffer_secs: Option<u64>,
|
||||
pub token_cache: Option<PyTokenCacheOptions>,
|
||||
}
|
||||
|
||||
impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
@@ -25,6 +52,17 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
fn try_from(py: PyOAuthConfig) -> Result<Self, Self::Error> {
|
||||
let flow = match py.flow.as_str() {
|
||||
"client_credentials" => OAuthFlow::ClientCredentials,
|
||||
"authorization_code" => {
|
||||
let mut options = AuthorizationCodeOptions::new().use_pkce(py.use_pkce);
|
||||
if let Some(redirect_uri) = py.redirect_uri {
|
||||
options = options.redirect_uri(redirect_uri);
|
||||
}
|
||||
if let Some(callback_port) = py.callback_port {
|
||||
options = options.callback_port(callback_port);
|
||||
}
|
||||
OAuthFlow::AuthorizationCode(options)
|
||||
}
|
||||
"device_code" => OAuthFlow::DeviceCode,
|
||||
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
|
||||
client_id: py.managed_identity_client_id,
|
||||
},
|
||||
@@ -42,6 +80,148 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
scopes: py.scopes,
|
||||
flow,
|
||||
refresh_buffer_secs: py.refresh_buffer_secs,
|
||||
token_cache: py.token_cache.map(TokenCacheOptions::from),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::SessionStatus`] exposing safe metadata.
|
||||
#[pyclass(name = "SessionStatus", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PySessionStatus {
|
||||
inner: SessionStatus,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySessionStatus {
|
||||
/// Whether a cached session exists that can obtain tokens without
|
||||
/// interactive authentication.
|
||||
#[getter]
|
||||
pub fn refreshable(&self) -> bool {
|
||||
self.inner.refreshable
|
||||
}
|
||||
|
||||
/// Canonical issuer URL of the cached session.
|
||||
#[getter]
|
||||
pub fn issuer_url(&self) -> String {
|
||||
self.inner.issuer_url.clone()
|
||||
}
|
||||
|
||||
/// Client ID of the cached session.
|
||||
#[getter]
|
||||
pub fn client_id(&self) -> String {
|
||||
self.inner.client_id.clone()
|
||||
}
|
||||
|
||||
/// Canonical (sorted, de-duplicated) scopes of the cached session.
|
||||
#[getter]
|
||||
pub fn scopes(&self) -> Vec<String> {
|
||||
self.inner.scopes.clone()
|
||||
}
|
||||
|
||||
/// Flow that produced the cached session.
|
||||
#[getter]
|
||||
pub fn flow(&self) -> String {
|
||||
self.inner.flow.clone()
|
||||
}
|
||||
|
||||
/// When the cached session was obtained, as Unix seconds.
|
||||
#[getter]
|
||||
pub fn obtained_at(&self) -> Option<u64> {
|
||||
self.inner.obtained_at
|
||||
}
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"SessionStatus(refreshable={}, issuer_url='{}', client_id='{}', flow='{}')",
|
||||
self.inner.refreshable, self.inner.issuer_url, self.inner.client_id, self.inner.flow
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionStatus> for PySessionStatus {
|
||||
fn from(inner: SessionStatus) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::SessionLogout`].
|
||||
#[pyclass(name = "SessionLogout", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PySessionLogout {
|
||||
inner: SessionLogout,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySessionLogout {
|
||||
/// Whether a cached credential was removed.
|
||||
#[getter]
|
||||
pub fn removed(&self) -> bool {
|
||||
self.inner.removed
|
||||
}
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!("SessionLogout(removed={})", self.inner.removed)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionLogout> for PySessionLogout {
|
||||
fn from(inner: SessionLogout) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::OAuthSession`].
|
||||
#[pyclass(name = "OAuthSession", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PyOAuthSession {
|
||||
inner: Arc<OAuthSession>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyOAuthSession {
|
||||
/// Create a session manager for the given OAuth configuration.
|
||||
///
|
||||
/// The configuration must set ``token_cache`` options and use a flow that
|
||||
/// supports persistent sessions (authorization code or device code).
|
||||
#[new]
|
||||
pub fn new(config: PyOAuthConfig) -> PyResult<Self> {
|
||||
let config: OAuthConfig = config.try_into().infer_error()?;
|
||||
let inner = OAuthSession::new(config).infer_error()?;
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
})
|
||||
}
|
||||
|
||||
/// Eagerly run the configured authentication flow and store the session.
|
||||
pub fn login<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner.login().await.map(PySessionStatus::from).infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Report whether a matching cached session exists, with safe metadata.
|
||||
pub fn status<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner
|
||||
.status()
|
||||
.await
|
||||
.map(PySessionStatus::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove the matching local cached credential.
|
||||
pub fn logout<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner
|
||||
.logout()
|
||||
.await
|
||||
.map(PySessionLogout::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -50,16 +230,27 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_unknown_oauth_flow_returns_invalid_input() {
|
||||
let config = PyOAuthConfig {
|
||||
fn base_config() -> PyOAuthConfig {
|
||||
PyOAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: "typo".to_string(),
|
||||
flow: "device_code".to_string(),
|
||||
client_secret: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: true,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unknown_oauth_flow_returns_invalid_input() {
|
||||
let config = PyOAuthConfig {
|
||||
flow: "typo".to_string(),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let err = OAuthConfig::try_from(config).unwrap_err();
|
||||
@@ -69,4 +260,53 @@ mod tests {
|
||||
if message == "Unknown OAuth flow type: typo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_authorization_code_conversion_preserves_options() {
|
||||
let config = PyOAuthConfig {
|
||||
flow: "authorization_code".to_string(),
|
||||
client_secret: Some("secret".to_string()),
|
||||
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
|
||||
callback_port: Some(9000),
|
||||
use_pkce: false,
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
let OAuthFlow::AuthorizationCode(options) = converted.flow else {
|
||||
panic!("expected authorization code flow");
|
||||
};
|
||||
assert_eq!(
|
||||
options.redirect_uri.as_deref(),
|
||||
Some("http://127.0.0.1:9000/callback")
|
||||
);
|
||||
assert_eq!(options.callback_port, Some(9000));
|
||||
assert!(!options.use_pkce);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_device_code_conversion() {
|
||||
let config = base_config();
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
assert!(matches!(converted.flow, OAuthFlow::DeviceCode));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_cache_conversion() {
|
||||
let config = PyOAuthConfig {
|
||||
token_cache: Some(PyTokenCacheOptions {
|
||||
cache_dir: Some("/tmp/oauth-cache".to_string()),
|
||||
lock_timeout_secs: Some(5),
|
||||
}),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
let cache = converted.token_cache.expect("token cache options");
|
||||
assert_eq!(
|
||||
cache.cache_dir.as_deref(),
|
||||
Some(std::path::Path::new("/tmp/oauth-cache"))
|
||||
);
|
||||
assert_eq!(cache.lock_timeout_secs, Some(5));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import urllib.parse
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_oauth_module():
|
||||
oauth_path = (
|
||||
@@ -31,3 +40,289 @@ def test_oauth_config_repr_redacts_client_secret():
|
||||
rendered = repr(config)
|
||||
assert "super-secret" not in rendered
|
||||
assert "client_secret" not in rendered
|
||||
|
||||
|
||||
def test_authorization_code_uses_pkce_by_default():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=oauth.OAuthFlowType.AUTHORIZATION_CODE,
|
||||
)
|
||||
|
||||
assert config.use_pkce is True
|
||||
assert config.redirect_uri is None
|
||||
assert config.callback_port is None
|
||||
|
||||
|
||||
def test_device_code_flow_value():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code"
|
||||
|
||||
|
||||
def test_token_cache_options_default_to_memory_only():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
)
|
||||
assert config.token_cache is None
|
||||
|
||||
options = oauth.TokenCacheOptions()
|
||||
assert options.cache_dir is None
|
||||
assert options.lock_timeout_secs is None
|
||||
|
||||
|
||||
def _remote_oauth():
|
||||
pytest.importorskip("lancedb")
|
||||
from lancedb.remote import oauth as remote_oauth
|
||||
|
||||
return remote_oauth
|
||||
|
||||
|
||||
def _device_config(remote_oauth, issuer_url, cache_dir):
|
||||
return remote_oauth.OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=remote_oauth.OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=remote_oauth.TokenCacheOptions(cache_dir=str(cache_dir)),
|
||||
)
|
||||
|
||||
|
||||
def test_oauth_session_status_and_logout_without_cache_entry(tmp_path):
|
||||
remote_oauth = _remote_oauth()
|
||||
config = _device_config(remote_oauth, "https://issuer.example.com", tmp_path)
|
||||
|
||||
session = remote_oauth.OAuthSession(config)
|
||||
status = asyncio.run(session.status())
|
||||
assert status.refreshable is False
|
||||
assert status.issuer_url == "https://issuer.example.com"
|
||||
assert status.client_id == "client-id"
|
||||
assert status.scopes == ["openid"]
|
||||
assert status.flow == "device_code"
|
||||
assert status.obtained_at is None
|
||||
|
||||
logout = asyncio.run(session.logout())
|
||||
assert logout.removed is False
|
||||
|
||||
|
||||
class _MockIdpState:
|
||||
def __init__(self, port):
|
||||
self.port = port
|
||||
self.lock = threading.Lock()
|
||||
self.device_authorizations = 0
|
||||
self.refresh_grants = 0
|
||||
self.invalid_grant_rejections = 0
|
||||
self.access_tokens_issued = 0
|
||||
self.current_refresh = None
|
||||
|
||||
|
||||
class _MockIdpHandler(BaseHTTPRequestHandler):
|
||||
@property
|
||||
def state(self) -> _MockIdpState:
|
||||
return self.server.state
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def _respond(self, status, payload):
|
||||
body = json.dumps(payload).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def do_GET(self):
|
||||
if self.path.startswith("/.well-known/openid-configuration"):
|
||||
base = f"http://127.0.0.1:{self.state.port}"
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"token_endpoint": f"{base}/token",
|
||||
"device_authorization_endpoint": f"{base}/device",
|
||||
},
|
||||
)
|
||||
else:
|
||||
self._respond(404, {})
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
body = self.rfile.read(length).decode()
|
||||
params = urllib.parse.parse_qs(body)
|
||||
|
||||
if self.path == "/device":
|
||||
with self.state.lock:
|
||||
self.state.device_authorizations += 1
|
||||
base = f"http://127.0.0.1:{self.state.port}"
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"device_code": "device-code",
|
||||
"user_code": "ABCD-EFGH",
|
||||
"verification_uri": f"{base}/verify",
|
||||
"expires_in": 60,
|
||||
"interval": 1,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
if self.path == "/token":
|
||||
grant_type = params.get("grant_type", [""])[0]
|
||||
with self.state.lock:
|
||||
if grant_type == "refresh_token":
|
||||
self.state.refresh_grants += 1
|
||||
offered = params.get("refresh_token", [""])[0]
|
||||
if offered != self.state.current_refresh:
|
||||
self.state.invalid_grant_rejections += 1
|
||||
self._respond(400, {"error": "invalid_grant"})
|
||||
return
|
||||
elif "device_code" not in grant_type:
|
||||
self._respond(400, {"error": "unsupported_grant_type"})
|
||||
return
|
||||
self.state.access_tokens_issued += 1
|
||||
number = self.state.access_tokens_issued
|
||||
refresh = f"refresh-{number}"
|
||||
self.state.current_refresh = refresh
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"access_token": f"access-{number}",
|
||||
"refresh_token": refresh,
|
||||
"expires_in": 3600,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
self._respond(404, {})
|
||||
|
||||
|
||||
def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]:
|
||||
server = HTTPServer(("127.0.0.1", 0), _MockIdpHandler)
|
||||
state = _MockIdpState(server.server_address[1])
|
||||
server.state = state
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
return state, server
|
||||
|
||||
|
||||
def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path):
|
||||
env = dict(os.environ)
|
||||
env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true"
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(script), issuer_url, str(cache_dir)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
env=env,
|
||||
)
|
||||
assert result.returncode == 0, (
|
||||
f"subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
LOGIN_SCRIPT = """
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
|
||||
|
||||
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
|
||||
config = OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=TokenCacheOptions(cache_dir=cache_dir),
|
||||
)
|
||||
session = OAuthSession(config)
|
||||
status = asyncio.run(session.login())
|
||||
assert status.refreshable, "login must cache a refresh token"
|
||||
print("LOGIN-OK")
|
||||
"""
|
||||
|
||||
REUSE_SCRIPT = """
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
import lancedb
|
||||
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
|
||||
|
||||
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
|
||||
config = OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=TokenCacheOptions(cache_dir=cache_dir),
|
||||
)
|
||||
|
||||
session = OAuthSession(config)
|
||||
status = asyncio.run(session.status())
|
||||
assert status.refreshable, "second process must see the cached session"
|
||||
|
||||
|
||||
async def main():
|
||||
# Point the database endpoint at a dead port. OAuth headers are fetched
|
||||
# before the request is sent, so a successful refresh proves the second
|
||||
# process reused the cached session; only the database call fails.
|
||||
db = await lancedb.connect_async(
|
||||
"db://e2e",
|
||||
host_override="http://127.0.0.1:1",
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
oauth_config=config,
|
||||
)
|
||||
try:
|
||||
await db.table_names()
|
||||
except Exception:
|
||||
print("DATABASE-UNREACHABLE-AS-EXPECTED")
|
||||
else:
|
||||
raise AssertionError("expected the database request to fail")
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
print("REUSE-OK")
|
||||
"""
|
||||
|
||||
|
||||
def test_cross_process_session_reuse_without_new_prompt(tmp_path):
|
||||
pytest.importorskip("lancedb")
|
||||
state, server = _start_mock_idp()
|
||||
try:
|
||||
issuer_url = f"http://127.0.0.1:{state.port}"
|
||||
login_script = tmp_path / "login.py"
|
||||
login_script.write_text(LOGIN_SCRIPT)
|
||||
reuse_script = tmp_path / "reuse.py"
|
||||
reuse_script.write_text(REUSE_SCRIPT)
|
||||
cache_dir = tmp_path / "oauth-cache"
|
||||
|
||||
result = _run_subprocess(login_script, issuer_url, cache_dir)
|
||||
assert "LOGIN-OK" in result.stdout
|
||||
assert state.device_authorizations == 1
|
||||
|
||||
result = _run_subprocess(reuse_script, issuer_url, cache_dir)
|
||||
assert "REUSE-OK" in result.stdout
|
||||
assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout
|
||||
|
||||
# The second process refreshed exactly once and never started a new
|
||||
# interactive device flow.
|
||||
assert state.refresh_grants == 1
|
||||
assert state.device_authorizations == 1
|
||||
assert state.invalid_grant_rejections == 0
|
||||
|
||||
logout = asyncio.run(
|
||||
_remote_oauth()
|
||||
.OAuthSession(_device_config(_remote_oauth(), issuer_url, cache_dir))
|
||||
.logout()
|
||||
)
|
||||
assert logout.removed is True
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
|
||||
@@ -53,7 +53,7 @@ metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio = { workspace = true, features = ["io-util", "net", "time"] }
|
||||
log.workspace = true
|
||||
async-trait = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
@@ -82,6 +82,9 @@ reqwest = { version = "0.12.0", default-features = false, features = [
|
||||
tonic = { workspace = true, optional = true }
|
||||
http = { version = "1", optional = true } # Matching what is in reqwest
|
||||
urlencoding = { version = "2", optional = true }
|
||||
base64 = { version = "0.22", optional = true }
|
||||
fs4 = { version = "0.13", optional = true }
|
||||
webbrowser = { version = "1", optional = true }
|
||||
uuid = { workspace = true, features = ["v5"] }
|
||||
polars-arrow = { version = ">=0.37,<0.40.0", optional = true }
|
||||
polars = { version = ">=0.37,<0.40.0", optional = true }
|
||||
@@ -159,6 +162,9 @@ remote = [
|
||||
"dep:http",
|
||||
"dep:tonic",
|
||||
"dep:urlencoding",
|
||||
"dep:base64",
|
||||
"dep:webbrowser",
|
||||
"dep:fs4",
|
||||
"lance-namespace-impls/rest",
|
||||
"lance-namespace-impls/rest-adapter",
|
||||
]
|
||||
|
||||
@@ -1546,6 +1546,7 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: crate::remote::OAuthFlow::ClientCredentials,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let result = ConnectBuilder::new("db://my-container/my-prefix")
|
||||
@@ -1588,6 +1589,7 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: crate::remote::OAuthFlow::ClientCredentials,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
let client_config = crate::remote::ClientConfig {
|
||||
header_provider: Some(
|
||||
|
||||
@@ -13,6 +13,7 @@ pub mod oauth;
|
||||
mod retry;
|
||||
pub(crate) mod sql;
|
||||
pub(crate) mod table;
|
||||
pub(crate) mod token_cache;
|
||||
pub(crate) mod util;
|
||||
|
||||
const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream";
|
||||
@@ -31,4 +32,5 @@ fn extract_job_id(body: &str) -> Option<String> {
|
||||
|
||||
pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig};
|
||||
pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder};
|
||||
pub use oauth::{OAuthConfig, OAuthFlow, OAuthHeaderProvider};
|
||||
pub use oauth::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow, OAuthHeaderProvider};
|
||||
pub use token_cache::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
|
||||
|
||||
+1921
-83
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user