feat: add persistent OAuth token cache and session APIs (#4182)

Stacked on #4173 (diff includes it until that merges; will rebase after). Addresses the token-cache part of [Colin's review](https://github.com/lancedb/lancedb/pull/4173#issuecomment-5674048100). Adds an explicit, opt-in persistent OAuth token cache shared by Rust, Python, and Node clients, plus `login` / `status` / `logout` session APIs, so short-lived processes (CLIs, scripts, notebooks) reuse one session instead of restarting a browser or device flow on every start. - **Opt-in and minimal**: existing callers stay memory-only and lazy. Only refresh tokens are persisted (never access tokens, never client secrets), so there are no local token-expiry decisions to get wrong when clocks move. Each process start performs one silent refresh grant. - **Hardened file backend**: private directory (`0700`), per-record files (`0600`), owner validation, symlink rejection, and atomic `rename` replacement. Corrupt, truncated, unknown-version, or permission-invalid records fail with actionable errors naming the file. Native keyring backends were evaluated (keyring crate routes Linux through D-Bus/zbus: heavy deps, headless/CI flakiness) and are deferred; the file store is the explicit opt-in, not a downgrade from a keyring. - **Cache key**: SHA-256 of the canonical identity (issuer, client ID, sorted/de-duplicated scopes, flow, public/confidential), so no secret appears in a filename and distinct identities never collide. Versioned record schema (`version: 1`). One record per identity: last login wins, documented. - **Cross-process rotation locking**: per-key `fs4` file lock (`flock` / `LockFileEx`) around the refresh critical section — acquire, reread the durable record, refresh exactly once, atomically store the rotated refresh token, release. The OS releases locks on process death, so crashes cannot strand stale locks. Only confirmed `invalid_grant`/`invalid_token` deletes a record and reauthenticates; transport, 5xx, 429, and parse failures retain it. - **Session APIs**: `OAuthSession::login/status/logout` in Rust, `lancedb.remote.OAuthSession` (async) in Python, `OAuthSession` class in Node. `status` returns non-secret metadata only. `logout` removes only the local credential — provider revocation (RFC 7009) is a deliberate follow-up, and local logout never terminates browser SSO. Azure managed identity is rejected for persistence (machine identity stays in memory); client credentials have nothing refreshable to persist and stay memory-only. - No CLI binary exists in this repo, so this ships library APIs plus doc examples in all three languages. Tests: Rust unit + mock-IdP integration (cache-key canonicalization/separation, record versioning/corruption/truncation/symlink/owner/perms, lock serialization + release, two concurrent providers proving no `invalid_grant` and correct rotation, transient-failure retention, `invalid_grant` delete + reauthenticate, login/status/logout lifecycle, client-credentials no-op, IMDS rejection, secret redaction); Python lifecycle + a true two-subprocess cross-process reuse test (second process refreshes once, never hits the device endpoint); Node lifecycle + device-flow login test. Local builds were skipped in development; CI validates all bindings. --------- Co-authored-by: Xuanwo <github@xuanwo.io>
2026-09-22 13:05:48 +00:00 · 2026-09-16 02:01:10 +08:00
parent 575286922b
commit 2f88b71c21
25 changed files with 5464 additions and 97 deletions
@@ -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>()?;
@@ -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));
+    }
 }