feat(oidc): support resource and audience (#4193)

Support configuring resource and audience for OAuth authorization, token
exchange, and refresh requests.

python/python/lancedb/_lancedb.pyi

@@ -284,6 +284,10 @@ class JobInfo:
     def created_at_millis(self) -> int: ...
 
 class SessionStatus:
+    @property
+    def resource(self) -> Optional[str]: ...
+    @property
+    def audience(self) -> Optional[str]: ...
     @property
     def refreshable(self) -> bool: ...
     @property

python/python/lancedb/remote/oauth.py

@@ -86,6 +86,13 @@ class OAuthConfig:
         Protect AUTHORIZATION_CODE with S256 PKCE (default: True).
     managed_identity_client_id : Optional[str]
         Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
+    resource : Optional[str]
+        Resource indicator (RFC 8707), forwarded verbatim to authorization and
+        token endpoints, including refresh requests. Must be an absolute URI
+        without a fragment. Not supported for Azure managed identity.
+    audience : Optional[str]
+        Provider-specific audience, forwarded to authorization and token
+        endpoints, including refresh requests. Not supported for 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
@@ -106,6 +113,12 @@ class OAuthConfig:
     ...     scopes=["api://lancedb-api/.default"],
     ... )
 
+    Providers that require an explicit target can use ``resource`` and/or
+    ``audience`` (these are forwarded unchanged):
+
+    >>> config.resource = "https://api.example.com"
+    >>> config.audience = "lancedb-api"
+
     Azure Managed Identity:
 
     >>> config = OAuthConfig(
@@ -150,6 +163,8 @@ class OAuthConfig:
     managed_identity_client_id: Optional[str] = None
     refresh_buffer_secs: Optional[int] = None
     token_cache: Optional[TokenCacheOptions] = None
+    resource: Optional[str] = None
+    audience: Optional[str] = None
 
 
 class OAuthSession:

python/src/oauth.rs

@@ -36,6 +36,10 @@ pub struct PyOAuthConfig {
     pub issuer_url: String,
     pub client_id: String,
     pub scopes: Vec<String>,
+    /// Optional resource indicator for authorization and token requests.
+    pub resource: Option<String>,
+    /// Optional provider-specific audience for authorization and token requests.
+    pub audience: Option<String>,
     pub flow: String,
     pub client_secret: Option<String>,
     pub redirect_uri: Option<String>,
@@ -78,6 +82,8 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
             client_id: py.client_id,
             client_secret: py.client_secret,
             scopes: py.scopes,
+            resource: py.resource,
+            audience: py.audience,
             flow,
             refresh_buffer_secs: py.refresh_buffer_secs,
             token_cache: py.token_cache.map(TokenCacheOptions::from),
@@ -119,6 +125,18 @@ impl PySessionStatus {
         self.inner.scopes.clone()
     }
 
+    /// Resource indicator used to obtain the cached session.
+    #[getter]
+    pub fn resource(&self) -> Option<String> {
+        self.inner.resource.clone()
+    }
+
+    /// Provider-specific audience used to obtain the cached session.
+    #[getter]
+    pub fn audience(&self) -> Option<String> {
+        self.inner.audience.clone()
+    }
+
     /// Flow that produced the cached session.
     #[getter]
     pub fn flow(&self) -> String {
@@ -242,6 +260,8 @@ mod tests {
             use_pkce: true,
             managed_identity_client_id: None,
             refresh_buffer_secs: None,
+            resource: None,
+            audience: None,
             token_cache: None,
         }
     }
@@ -269,6 +289,8 @@ mod tests {
             redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
             callback_port: Some(9000),
             use_pkce: false,
+            resource: Some("urn:resource".into()),
+            audience: Some("audience".into()),
             ..base_config()
         };
 
@@ -282,6 +304,8 @@ mod tests {
         );
         assert_eq!(options.callback_port, Some(9000));
         assert!(!options.use_pkce);
+        assert_eq!(converted.resource.as_deref(), Some("urn:resource"));
+        assert_eq!(converted.audience.as_deref(), Some("audience"));
     }
 
     #[test]
@@ -294,6 +318,8 @@ mod tests {
     #[test]
     fn test_token_cache_conversion() {
         let config = PyOAuthConfig {
+            resource: None,
+            audience: None,
             token_cache: Some(PyTokenCacheOptions {
                 cache_dir: Some("/tmp/oauth-cache".to_string()),
                 lock_timeout_secs: Some(5),

python/tests/test_oauth.py

@@ -72,6 +72,8 @@ def test_token_cache_options_default_to_memory_only():
         scopes=["openid"],
     )
     assert config.token_cache is None
+    assert config.resource is None
+    assert config.audience is None
 
     options = oauth.TokenCacheOptions()
     assert options.cache_dir is None
@@ -121,6 +123,7 @@ class _MockIdpState:
         self.invalid_grant_rejections = 0
         self.access_tokens_issued = 0
         self.current_refresh = None
+        self.requests = []
 
 
 class _MockIdpHandler(BaseHTTPRequestHandler):
@@ -156,6 +159,8 @@ class _MockIdpHandler(BaseHTTPRequestHandler):
         length = int(self.headers.get("Content-Length", 0))
         body = self.rfile.read(length).decode()
         params = urllib.parse.parse_qs(body)
+        with self.state.lock:
+            self.state.requests.append(params)
 
         if self.path == "/device":
             with self.state.lock:
@@ -212,11 +217,11 @@ def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]:
     return state, server
 
 
-def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path):
+def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path, target: dict):
     env = dict(os.environ)
     env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true"
     result = subprocess.run(
-        [sys.executable, str(script), issuer_url, str(cache_dir)],
+        [sys.executable, str(script), issuer_url, str(cache_dir), json.dumps(target)],
         capture_output=True,
         text=True,
         timeout=120,
@@ -230,6 +235,7 @@ def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path):
 
 LOGIN_SCRIPT = """
 import asyncio
+import json
 import sys
 
 from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
@@ -241,15 +247,19 @@ config = OAuthConfig(
     scopes=["openid"],
     flow=OAuthFlowType.DEVICE_CODE,
     token_cache=TokenCacheOptions(cache_dir=cache_dir),
+    **json.loads(sys.argv[3]),
 )
 session = OAuthSession(config)
 status = asyncio.run(session.login())
 assert status.refreshable, "login must cache a refresh token"
+assert status.resource == config.resource
+assert status.audience == config.audience
 print("LOGIN-OK")
 """
 
 REUSE_SCRIPT = """
 import asyncio
+import json
 import sys
 
 import lancedb
@@ -262,6 +272,7 @@ config = OAuthConfig(
     scopes=["openid"],
     flow=OAuthFlowType.DEVICE_CODE,
     token_cache=TokenCacheOptions(cache_dir=cache_dir),
+    **json.loads(sys.argv[3]),
 )
 
 session = OAuthSession(config)
@@ -292,7 +303,17 @@ print("REUSE-OK")
 """
 
 
-def test_cross_process_session_reuse_without_new_prompt(tmp_path):
+@pytest.mark.parametrize(
+    "target",
+    [
+        {},
+        {
+            "resource": "https://api.example.com/a?x=1&y=two",
+            "audience": "audience + & / ü",
+        },
+    ],
+)
+def test_cross_process_session_reuse_without_new_prompt(tmp_path, target):
     pytest.importorskip("lancedb")
     state, server = _start_mock_idp()
     try:
@@ -303,11 +324,11 @@ def test_cross_process_session_reuse_without_new_prompt(tmp_path):
         reuse_script.write_text(REUSE_SCRIPT)
         cache_dir = tmp_path / "oauth-cache"
 
-        result = _run_subprocess(login_script, issuer_url, cache_dir)
+        result = _run_subprocess(login_script, issuer_url, cache_dir, target)
         assert "LOGIN-OK" in result.stdout
         assert state.device_authorizations == 1
 
-        result = _run_subprocess(reuse_script, issuer_url, cache_dir)
+        result = _run_subprocess(reuse_script, issuer_url, cache_dir, target)
         assert "REUSE-OK" in result.stdout
         assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout
 
@@ -317,11 +338,14 @@ def test_cross_process_session_reuse_without_new_prompt(tmp_path):
         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 len(state.requests) == 3
+        for params in state.requests:
+            for key in ("resource", "audience"):
+                assert params.get(key) == ([target[key]] if key in target else None)
+        config = _device_config(_remote_oauth(), issuer_url, cache_dir)
+        config.resource = target.get("resource")
+        config.audience = target.get("audience")
+        logout = asyncio.run(_remote_oauth().OAuthSession(config).logout())
         assert logout.removed is True
     finally:
         server.shutdown()