rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-21 12:35:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: use oauth2 crate for OAuth with configurable client auth (#4181)

Browse Source 
Stacked on #4173 (`jack/restore-oidc-flows`, base branch mirrored to
this repo so the diff shows only this change); context from review:
https://github.com/lancedb/lancedb/pull/4173#issuecomment-5674048100.
Rebase to `main` once #4173 and #4179 merge.

## What moved to the `oauth2` crate (5.0, no default features)

- Authorization URL generation and CSRF state (`authorize_url`,
`CsrfToken`)
- PKCE S256 challenge/verifier generation and code exchange
- Client-credentials, authorization-code, refresh-token, and device-code
grant request construction
- Device authorization request and the device token polling loop
(`authorization_pending`, `slow_down` +5s, expiry deadline, denial,
network backoff capped at 10s)
- Standard success/error response parsing (`RequestTokenError`)
- Token-endpoint client authentication and standards-compliant parameter
encoding (RFC 6749 2.3.1 Basic encoding)

LanceDB keeps ownership of OIDC discovery (compared `openidconnect`: no
measurable win for our 3-field metadata + strict validation, at real
dependency cost), HTTPS-or-loopback endpoint enforcement, the loopback
callback server, browser/stderr prompts, token caching and refresh
orchestration, and the dedicated hardened Azure IMDS source, which is
unchanged.

## Client authentication methods

New `ClientAuthMethod` enum (`none` | `client_secret_basic` |
`client_secret_post`), exposed in Rust, Python, and Node. Unset resolves
to `client_secret_basic` when a secret is present (RFC 6749 2.3.1
recommendation and the normal Okta confidential-app default, so a
default Okta app works without weakening its configuration) and to
`none` for public clients (PKCE/device). Explicit `none` with a secret,
or basic/post without one, is rejected. The method applies to client
credentials, code exchange, refresh, and device requests. Deliberate
behavior change: confidential clients previously always sent the secret
in the POST body; they now default to Basic (Keycloak accepts both).

No `audience`/`resource` parameters were added: the supported target is
an Okta custom authorization server with the API audience configured
server-side, so client-provided audience parameters are unnecessary;
`add_extra_param` support exists if a concrete provider contract ever
needs them.

## Device polling behavior changes (deliberate, tested)

- The first token poll now happens immediately rather than after one
interval (RFC 8628 allows both).
- Transient failures (HTTP 429, 5xx, `temporarily_unavailable`, network
errors) now retry with exponential backoff capped at 10s instead of
retrying at the fixed interval; polling never spins faster than once per
second even if a server reports a zero interval.

## Security and compatibility

- Issuer and discovered endpoints (and device verification URIs) still
require HTTPS except explicit loopback HTTP, enforced before any crate
URL type is built
- Token HTTP client keeps the hardened redirect policy that refuses
insecure redirect targets; regression test added
- Errors never embed raw response bodies (avoids leaking tokens through
parse failures); all credential types stay redacted in Debug
- Transient conditions (429/5xx/`temporarily_unavailable`) remain
retryable in device polling and hard errors elsewhere; refresh keeps
rotation and reauthentication semantics
- Existing public APIs stay source-compatible except the added
`OAuthConfig.client_auth_method` field

## Tests

Rust: client-auth methods across code
exchange/refresh/client-credentials/device (none/basic/post),
auth-method resolution and validation, transient device retries,
denial/expiry, redirect rejection, malformed-response leak check, PKCE
URL assertions, redaction. Python and Node: enum values, conversion,
unknown-method errors, config defaults.

Manual Okta validation recipe (no automated Okta credentials): create a
custom authorization server with an API audience, one confidential web
app (Basic) for authorization-code, one native app (PKCE, no secret),
one native device app; point `issuer_url` at the custom server, set
`client_auth_method` only for the POST-required case; verify token
acquisition, refresh after expiry, and `x-lancedb-credential-type: oidc`
against a LanceDB deployment. Never commit tenant URLs or secrets.

Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
Jack Ye
2026-09-16 20:11:08 +08:00
committed by GitHub
co-authored by Xuanwo
parent ed100ccc31
commit f3ef21b8ca
19 changed files with 1485 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+21
View File
@@ -5526,6 +5526,7 @@ dependencies = [
"metrics-util",
"moka",
"num-traits",
"oauth2",
"object_store 0.14.1",
"pin-project",
"polars",
@@ -6419,6 +6420,25 @@ version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
[[package]]
name = "oauth2"
version = "5.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d"
dependencies = [
"base64 0.22.1",
"chrono",
"getrandom 0.2.17",
"http 1.5.0",
"rand 0.8.6",
"serde",
"serde_json",
"serde_path_to_error",
"sha2 0.10.9",
"thiserror 1.0.69",
"url",
]
[[package]]
name = "objc2"
version = "0.6.4"
@@ -10584,6 +10604,7 @@ dependencies = [
"idna",
"percent-encoding",
"serde",
"serde_derive",
]
[[package]]
Cargo.toml
+1
View File
@@ -68,6 +68,7 @@ rand = "0.9"
snafu = "0.8"
url = "2"
num-traits = "0.2"
oauth2 = { version = "5.0", default-features = false }
regex = "1.10"
semver = "1.0.25"
serde = "1"
docs/src/js/enumerations/ClientAuthMethod.md
+48
View File
@@ -0,0 +1,48 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / ClientAuthMethod
# Enumeration: ClientAuthMethod
How the client authenticates to the OAuth token endpoint.
The method applies to every OAuth request that carries client
authentication: client-credentials, authorization-code exchange,
refresh-token, and device-authorization requests. The Azure managed
identity flow ignores this option.
## Enumeration Members
### ClientSecretBasic
```ts
ClientSecretBasic: "client_secret_basic";
```
HTTP Basic authentication. This is the RFC 6749 recommended method and
the normal default for confidential clients, including default Okta
applications. Requires `clientSecret`.
***
### ClientSecretPost
```ts
ClientSecretPost: "client_secret_post";
```
Credentials in the request body, for providers configured to require it.
Requires `clientSecret`.
***
### None
```ts
None: "none";
```
No client authentication, for public clients using PKCE or the device
flow. Cannot be combined with `clientSecret`.
docs/src/js/globals.md
+1
View File
@@ -11,6 +11,7 @@
## Enumerations
- [ClientAuthMethod](enumerations/ClientAuthMethod.md)
- [FullTextQueryType](enumerations/FullTextQueryType.md)
- [OAuthFlowType](enumerations/OAuthFlowType.md)
- [Occur](enumerations/Occur.md)
docs/src/js/interfaces/NativeOAuthConfig.md
+13
View File
@@ -35,6 +35,19 @@ Port for the authorization_code loopback callback server.
***
### clientAuthMethod?
```ts
optional clientAuthMethod: string;
```
How the client authenticates to the token endpoint: "none",
"client_secret_basic", or "client_secret_post". Defaults to
"client_secret_basic" when a client secret is set, and "none" for
public clients.
***
### clientId
```ts
docs/src/js/interfaces/OAuthConfig.md
+14
View File
@@ -84,6 +84,20 @@ Port for the AuthorizationCode loopback callback server (default: 8400).
***
### clientAuthMethod?
```ts
optional clientAuthMethod: ClientAuthMethod;
```
How the client authenticates to the token endpoint (default: auto).
With a `clientSecret` the default is `ClientAuthMethod.ClientSecretBasic`,
which matches the RFC 6749 recommendation and the default configuration
of Okta confidential applications; without a secret the client is public
and no client authentication is sent.
***
### clientId
```ts
nodejs/__test__/remote.test.ts
+37
View File
@@ -5,9 +5,12 @@ import * as http from "http";
import { RequestListener } from "http";
import packageJson = require("../package.json");
import {
ClientAuthMethod,
ClientConfig,
Connection,
ConnectionOptions,
OAuthConfig,
OAuthFlowType,
TlsConfig,
connect,
} from "../lancedb";
@@ -438,6 +441,40 @@ describe("remote connection", () => {
]);
});
describe("OAuthConfig", () => {
it("should expose client auth method values", () => {
expect(ClientAuthMethod.None).toBe("none");
expect(ClientAuthMethod.ClientSecretBasic).toBe("client_secret_basic");
expect(ClientAuthMethod.ClientSecretPost).toBe("client_secret_post");
});
it("should accept a confidential client with basic auth", () => {
const config: OAuthConfig = {
issuerUrl: "https://issuer.example.com",
clientId: "client-id",
clientSecret: "secret",
scopes: ["openid"],
flow: OAuthFlowType.AuthorizationCode,
clientAuthMethod: ClientAuthMethod.ClientSecretBasic,
};
expect(config.clientAuthMethod).toBe(ClientAuthMethod.ClientSecretBasic);
});
it("should accept a public PKCE client without auth method or secret", () => {
const config: OAuthConfig = {
issuerUrl: "https://issuer.example.com",
clientId: "client-id",
scopes: ["openid"],
flow: OAuthFlowType.AuthorizationCode,
usePkce: true,
};
expect(config.clientSecret).toBeUndefined();
expect(config.clientAuthMethod).toBeUndefined();
});
});
describe("TlsConfig", () => {
it("should create TlsConfig with all fields", () => {
const tlsConfig: TlsConfig = {
nodejs/lancedb/index.ts
+1
View File
@@ -172,6 +172,7 @@ export {
} from "./header";
export {
ClientAuthMethod,
OAuthConfig,
OAuthFlowType,
OAuthSession,
nodejs/lancedb/oauth.ts
+36
View File
@@ -47,6 +47,33 @@ export interface TokenCacheOptions {
lockTimeoutSecs?: number;
}
/**
* How the client authenticates to the OAuth token endpoint.
*
* The method applies to every OAuth request that carries client
* authentication: client-credentials, authorization-code exchange,
* refresh-token, and device-authorization requests. The Azure managed
* identity flow ignores this option.
*/
export enum ClientAuthMethod {
/**
* No client authentication, for public clients using PKCE or the device
* flow. Cannot be combined with `clientSecret`.
*/
None = "none",
/**
* HTTP Basic authentication. This is the RFC 6749 recommended method and
* the normal default for confidential clients, including default Okta
* applications. Requires `clientSecret`.
*/
ClientSecretBasic = "client_secret_basic",
/**
* Credentials in the request body, for providers configured to require it.
* Requires `clientSecret`.
*/
ClientSecretPost = "client_secret_post",
}
/**
* OAuth configuration for LanceDB authentication.
*
@@ -140,6 +167,15 @@ export interface OAuthConfig {
/** Client secret (required for ClientCredentials). */
clientSecret?: string;
/**
* How the client authenticates to the token endpoint (default: auto).
* With a `clientSecret` the default is `ClientAuthMethod.ClientSecretBasic`,
* which matches the RFC 6749 recommendation and the default configuration
* of Okta confidential applications; without a secret the client is public
* and no client authentication is sent.
*/
clientAuthMethod?: ClientAuthMethod;
/** Loopback redirect URI for AuthorizationCode. */
redirectUri?: string;
nodejs/src/remote.rs
+85
View File
@@ -197,6 +197,11 @@ pub struct OAuthConfig {
pub flow: Option<String>,
/// Client secret (required for client_credentials).
pub client_secret: Option<String>,
/// How the client authenticates to the token endpoint: "none",
/// "client_secret_basic", or "client_secret_post". Defaults to
/// "client_secret_basic" when a client secret is set, and "none" for
/// public clients.
pub client_auth_method: Option<String>,
/// Loopback redirect URI for authorization_code.
pub redirect_uri: Option<String>,
/// Port for the authorization_code loopback callback server.
@@ -227,6 +232,7 @@ impl std::fmt::Debug for OAuthConfig {
"client_secret",
&self.client_secret.as_deref().map(|_| "<redacted>"),
)
.field("client_auth_method", &self.client_auth_method)
.field("redirect_uri", &self.redirect_uri)
.field("callback_port", &self.callback_port)
.field("use_pkce", &self.use_pkce)
@@ -270,10 +276,27 @@ impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
}
};
let client_auth_method = match config.client_auth_method.as_deref() {
Some("none") => Some(lancedb::remote::oauth::ClientAuthMethod::None),
Some("client_secret_basic") => {
Some(lancedb::remote::oauth::ClientAuthMethod::ClientSecretBasic)
}
Some("client_secret_post") => {
Some(lancedb::remote::oauth::ClientAuthMethod::ClientSecretPost)
}
None => None,
Some(other) => {
return Err(Error::InvalidInput {
message: format!("Unknown OAuth client auth method: {other}"),
});
}
};
Ok(Self {
issuer_url: config.issuer_url,
client_id: config.client_id,
client_secret: config.client_secret,
client_auth_method,
scopes: config.scopes,
resource: config.resource,
audience: config.audience,
@@ -427,6 +450,7 @@ mod tests {
scopes: vec!["scope".to_string()],
flow: Some("typo".to_string()),
client_secret: None,
client_auth_method: None,
redirect_uri: None,
callback_port: None,
use_pkce: None,
@@ -453,6 +477,7 @@ mod tests {
scopes: vec!["scope".to_string()],
flow: Some("client_credentials".to_string()),
client_secret: Some("super-secret".to_string()),
client_auth_method: None,
redirect_uri: None,
callback_port: None,
use_pkce: None,
@@ -476,6 +501,7 @@ mod tests {
scopes: vec!["openid".to_string()],
flow: Some("authorization_code".to_string()),
client_secret: Some("secret".to_string()),
client_auth_method: None,
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
callback_port: Some(9000),
use_pkce: Some(false),
@@ -508,6 +534,7 @@ mod tests {
scopes: vec!["openid".to_string()],
flow: Some("device_code".to_string()),
client_secret: None,
client_auth_method: None,
redirect_uri: None,
callback_port: None,
use_pkce: None,
@@ -524,4 +551,62 @@ mod tests {
lancedb::remote::oauth::OAuthFlow::DeviceCode
));
}
#[test]
fn test_client_auth_method_conversion() {
use lancedb::remote::oauth::ClientAuthMethod;
for (value, expected) in [
("none", ClientAuthMethod::None),
("client_secret_basic", ClientAuthMethod::ClientSecretBasic),
("client_secret_post", ClientAuthMethod::ClientSecretPost),
] {
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,
client_auth_method: Some(value.to_string()),
resource: None,
audience: 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_eq!(converted.client_auth_method, Some(expected));
}
}
#[test]
fn test_unknown_client_auth_method_returns_invalid_input() {
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,
client_auth_method: Some("typo".to_string()),
resource: None,
audience: 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();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "Unknown OAuth client auth method: typo"
));
}
}
python/python/lancedb/remote/__init__.py
+8 -1
View File
@@ -9,7 +9,13 @@ from typing import List, Optional
from lancedb import __version__
from .header import HeaderProvider
from .oauth import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
from .oauth import (
ClientAuthMethod,
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 +28,7 @@ __all__ = [
"HeaderProvider",
"OAuthConfig",
"OAuthFlowType",
"ClientAuthMethod",
"OAuthSession",
"TokenCacheOptions",
]
python/python/lancedb/remote/oauth.py
+34 -2
View File
@@ -22,6 +22,29 @@ class OAuthFlowType(str, Enum):
"""Azure Managed Identity via IMDS."""
class ClientAuthMethod(str, Enum):
"""How the client authenticates to the OAuth token endpoint.
The method applies to every OAuth request that carries client
authentication: client-credentials, authorization-code exchange,
refresh-token, and device-authorization requests. The Azure managed
identity flow ignores this option.
"""
NONE = "none"
"""No client authentication, for public clients using PKCE or the device
flow. Cannot be combined with ``client_secret``."""
CLIENT_SECRET_BASIC = "client_secret_basic"
"""HTTP Basic authentication. This is the RFC 6749 recommended method and
the normal default for confidential clients, including default Okta
applications. Requires ``client_secret``."""
CLIENT_SECRET_POST = "client_secret_post"
"""Credentials in the request body, for providers configured to require
it. Requires ``client_secret``."""
@dataclass
class TokenCacheOptions:
"""Options for the persistent OAuth token cache.
@@ -77,6 +100,13 @@ class OAuthConfig:
Authentication flow to use. Default: CLIENT_CREDENTIALS.
client_secret : Optional[str]
Client secret (required for CLIENT_CREDENTIALS).
client_auth_method : Optional[ClientAuthMethod]
How the client authenticates to the token endpoint (default: auto).
With a ``client_secret`` the default is
``ClientAuthMethod.CLIENT_SECRET_BASIC``, which matches the RFC 6749
recommendation and the default configuration of Okta confidential
applications; without a secret the client is public and no client
authentication is sent.
redirect_uri : Optional[str]
Loopback redirect URI for AUTHORIZATION_CODE. The default is
``http://127.0.0.1:{callback_port}/callback``.
@@ -140,8 +170,9 @@ class OAuthConfig:
... flow=OAuthFlowType.AUTHORIZATION_CODE,
... )
Device Authorization with a persistent cache, so later processes reuse
the session without a new device prompt:
Device Authorization, with a persistent cache so later processes reuse
the session without a new device prompt. The verification URL and user
code are written to standard error before polling begins:
>>> config = OAuthConfig(
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
@@ -157,6 +188,7 @@ class OAuthConfig:
scopes: List[str]
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
client_secret: Optional[str] = field(default=None, repr=False)
client_auth_method: Optional[ClientAuthMethod] = None
redirect_uri: Optional[str] = None
callback_port: Optional[int] = None
use_pkce: bool = True
python/src/oauth.rs
+48 -1
View File
@@ -9,7 +9,7 @@ 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::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow};
use lancedb::remote::oauth::{AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow};
use lancedb::remote::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
/// Python-side persistent token cache options, extracted via FromPyObject.
@@ -42,6 +42,7 @@ pub struct PyOAuthConfig {
pub audience: Option<String>,
pub flow: String,
pub client_secret: Option<String>,
pub client_auth_method: Option<String>,
pub redirect_uri: Option<String>,
pub callback_port: Option<u16>,
pub use_pkce: bool,
@@ -77,10 +78,23 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
}
};
let client_auth_method = match py.client_auth_method.as_deref() {
Some("none") => Some(ClientAuthMethod::None),
Some("client_secret_basic") => Some(ClientAuthMethod::ClientSecretBasic),
Some("client_secret_post") => Some(ClientAuthMethod::ClientSecretPost),
None => None,
Some(other) => {
return Err(Error::InvalidInput {
message: format!("Unknown OAuth client auth method: {other}"),
});
}
};
Ok(Self {
issuer_url: py.issuer_url,
client_id: py.client_id,
client_secret: py.client_secret,
client_auth_method,
scopes: py.scopes,
resource: py.resource,
audience: py.audience,
@@ -255,6 +269,7 @@ mod tests {
scopes: vec!["scope".to_string()],
flow: "device_code".to_string(),
client_secret: None,
client_auth_method: None,
redirect_uri: None,
callback_port: None,
use_pkce: true,
@@ -315,6 +330,38 @@ mod tests {
assert!(matches!(converted.flow, OAuthFlow::DeviceCode));
}
#[test]
fn test_client_auth_method_conversion() {
for (value, expected) in [
("none", ClientAuthMethod::None),
("client_secret_basic", ClientAuthMethod::ClientSecretBasic),
("client_secret_post", ClientAuthMethod::ClientSecretPost),
] {
let config = PyOAuthConfig {
client_auth_method: Some(value.to_string()),
..base_config()
};
let converted = OAuthConfig::try_from(config).unwrap();
assert_eq!(converted.client_auth_method, Some(expected));
}
}
#[test]
fn test_unknown_client_auth_method_returns_invalid_input() {
let config = PyOAuthConfig {
client_auth_method: Some("typo".to_string()),
..base_config()
};
let err = OAuthConfig::try_from(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "Unknown OAuth client auth method: typo"
));
}
#[test]
fn test_token_cache_conversion() {
let config = PyOAuthConfig {
python/tests/test_oauth.py
+28
View File
@@ -63,6 +63,34 @@ def test_device_code_flow_value():
assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code"
def test_client_auth_method_values():
oauth = _load_oauth_module()
assert oauth.ClientAuthMethod.NONE.value == "none"
assert oauth.ClientAuthMethod.CLIENT_SECRET_BASIC.value == "client_secret_basic"
assert oauth.ClientAuthMethod.CLIENT_SECRET_POST.value == "client_secret_post"
def test_client_auth_method_defaults_to_none():
oauth = _load_oauth_module()
config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
client_auth_method=oauth.ClientAuthMethod.CLIENT_SECRET_POST,
)
assert config.client_auth_method is oauth.ClientAuthMethod.CLIENT_SECRET_POST
default_config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["openid"],
)
assert default_config.client_auth_method is None
def test_token_cache_options_default_to_memory_only():
oauth = _load_oauth_module()
rust/lancedb/Cargo.toml
+2
View File
@@ -82,6 +82,7 @@ reqwest = { version = "0.12.0", default-features = false, features = [
], optional = true }
tonic = { workspace = true, optional = true }
http = { version = "1", optional = true } # Matching what is in reqwest
oauth2 = { workspace = true, optional = true }
urlencoding = { version = "2", optional = true }
base64 = { version = "0.22", optional = true }
fs4 = { version = "0.13", optional = true }
@@ -159,6 +160,7 @@ remote = [
"dep:prost",
"dep:reqwest",
"dep:http",
"dep:oauth2",
"dep:tonic",
"dep:urlencoding",
"dep:base64",
rust/lancedb/src/connection.rs
+2
View File
@@ -1546,6 +1546,7 @@ mod tests {
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: crate::remote::OAuthFlow::ClientCredentials,
client_auth_method: None,
refresh_buffer_secs: None,
resource: None,
audience: None,
@@ -1591,6 +1592,7 @@ mod tests {
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: crate::remote::OAuthFlow::ClientCredentials,
client_auth_method: None,
refresh_buffer_secs: None,
resource: None,
audience: None,
rust/lancedb/src/remote.rs
+3 -1
View File
@@ -32,5 +32,7 @@ fn extract_job_id(body: &str) -> Option<String> {
pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig};
pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder};
pub use oauth::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow, OAuthHeaderProvider};
pub use oauth::{
AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow, OAuthHeaderProvider,
};
pub use token_cache::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
rust/lancedb/src/remote/oauth.rs
+1084 -494
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/remote/token_cache.rs
+19 -11
View File
@@ -44,6 +44,7 @@
//! client_secret: None,
//! scopes: vec!["openid".to_string()],
//! flow: OAuthFlow::DeviceCode,
//! client_auth_method: None,
//! refresh_buffer_secs: None,
//! resource: Some("https://api.example.com".to_string()),
//! audience: None,
@@ -420,7 +421,10 @@ impl TokenCache {
/// Build a record from a token response, or `None` when the response
/// carries no refresh token (nothing may be persisted).
fn record_from_response(&self, response: &TokenResponse) -> Option<CachedTokenRecord> {
let refresh_token = response.refresh_token.clone()?;
let refresh_token = response
.refresh_token
.as_ref()
.map(|token| token.secret().clone())?;
Some(CachedTokenRecord {
version: CACHE_RECORD_VERSION,
issuer_url: self.key.issuer_url.clone(),
@@ -842,6 +846,7 @@ pub struct SessionLogout {
/// client_secret: None,
/// scopes: vec!["openid".to_string()],
/// flow: OAuthFlow::DeviceCode,
/// client_auth_method: None,
/// refresh_buffer_secs: None,
/// resource: None,
/// audience: None,
@@ -983,6 +988,8 @@ mod tests {
use crate::remote::HeaderProvider;
use crate::remote::oauth::OAuthHeaderProvider;
use oauth2::basic::BasicTokenType;
use oauth2::{AccessToken, RefreshToken};
use serial_test::serial;
/// Temp directory that satisfies the cache hardening checks. CI runners
@@ -1003,6 +1010,7 @@ mod tests {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: None,
client_auth_method: None,
scopes: vec!["openid".to_string()],
flow: OAuthFlow::DeviceCode,
refresh_buffer_secs: None,
@@ -1479,8 +1487,8 @@ mod tests {
assert!(keys.insert(cache.key.file_stem.clone()));
let record = cache
.record_from_response(&TokenResponse {
access_token: "unused".into(),
refresh_token: Some("seed-refresh".into()),
access_token: AccessToken::new("unused".into()),
refresh_token: Some(RefreshToken::new("seed-refresh".into())),
expires_in: Some(3600),
token_type: None,
})
@@ -1634,10 +1642,10 @@ mod tests {
)
.unwrap();
let response = TokenResponse {
access_token: "access-token".to_string(),
refresh_token: Some("refresh-token".to_string()),
access_token: AccessToken::new("access-token".to_string()),
refresh_token: Some(RefreshToken::new("refresh-token".to_string())),
expires_in: Some(3600),
token_type: Some("Bearer".to_string()),
token_type: Some(BasicTokenType::Bearer),
};
let record = cache.record_from_response(&response).unwrap();
let debug = format!("{record:?}");
@@ -1678,8 +1686,8 @@ mod tests {
.unwrap();
let record = cache
.record_from_response(&TokenResponse {
access_token: "a".to_string(),
refresh_token: Some("r".to_string()),
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
@@ -1740,8 +1748,8 @@ mod tests {
// A complete, well-formed record with an unknown schema version.
let record = cache
.record_from_response(&TokenResponse {
access_token: "a".to_string(),
refresh_token: Some("r".to_string()),
access_token: AccessToken::new("a".to_string()),
refresh_token: Some(RefreshToken::new("r".to_string())),
expires_in: None,
token_type: None,
})
@@ -1860,7 +1868,7 @@ mod tests {
)
.unwrap();
let response = TokenResponse {
access_token: "access-token".to_string(),
access_token: AccessToken::new("access-token".to_string()),
refresh_token: None,
expires_in: Some(3600),
token_type: None,