Skip to content

Add Databricks Vector Search MCP token exchange with per-workspace token isolation#27

Merged
yosrixp merged 1 commit intomainfrom
databricks-vsearch
Apr 16, 2026
Merged

Add Databricks Vector Search MCP token exchange with per-workspace token isolation#27
yosrixp merged 1 commit intomainfrom
databricks-vsearch

Conversation

@yosrixp
Copy link
Copy Markdown
Collaborator

@yosrixp yosrixp commented Apr 16, 2026

Summary

  • Add Databricks Vector Search as a new MCP provider that shares the same codebase as Databricks SQL (no class duplication). The single TokenExchangeServiceDatabricksImpl is scoped as @Dependent and configured via Instance<> injection in TokenExchangeServiceProducer to create two distinct instances -- one for SQL (scope=sql) and one for Vector Search (scope=vector-search) -- each with its own DatabricksTokenExchangeConfig.
  • Generalize Databricks internals: rename DatabricksSqlWorkspaceResolver to DatabricksWorkspaceResolver (handles multi-segment paths like /<workspace>/<catalog>/<schema>/mcp), DatabricksSqlTokenClient to DatabricksTokenClient, and extract a common DatabricksTokenExchangeConfig interface extended by both SQL and Vector Search config mappings.
  • Fix /userinfo "Okta token record not found" error that affects all providers (Databricks SQL/Vector Search, Glean, Google Monitoring/Logging, Splunk) whose exchanged access token outlives the Okta row TTL. When the Okta id_token row is missing or stale, /userinfo now attempts an upstream Okta refresh via UpstreamRefreshService.refreshUpstream() with the same distributed lock and optimistic versioning as the /token refresh_token grant, then stores the refreshed tokens and continues.
  • Add resource mappings and token-exchange configuration for Databricks Vector Search across all environments (local, stage, prod-us-east-1, prod-us-west-2) and Helm configmap template.

Test plan

  • Verify UserInfoResourceTest -- 5 new tests covering upstream refresh on Okta row missing, id_token null, lock failure, and UpstreamRefreshException, with metric verification (mop_token_exchange_step_total{exchange_step=upstream_refresh})
  • Verify DatabricksWorkspaceResolverTest -- new cases for vector search multi-segment URL patterns
  • Verify TokenExchangeServiceDatabricksImplTest -- vector search token exchange with scope=vector-search
  • Verify TokenExchangeServiceProducerTest -- both databricks-sql and databricks-vector-search resolve to distinct instances
  • Verify ExchangedTokenUserinfoStoreProviderResolverTest -- vector search audience resolution
  • Verify ConfigServiceTest -- glob-to-regex matching for databricks-vector-search/*/*/*/mcp wildcard resource URIs
  • Verify OauthProviderLabelTest -- databricks-vector-search normalizes correctly
  • End-to-end: call /token with a Databricks Vector Search resource, then /userinfo with the returned access token -- confirm userinfo returns Okta claims with mcp_resource_idp = databricks-vector-search-<hostname>
  • End-to-end: let Okta row expire, call /userinfo again -- confirm it auto-refreshes and returns 200 instead of 401

…ken isolation

Signed-off-by: yosrixp <yosrixp@yahoo.com>
@yosrixp yosrixp merged commit f7ce392 into main Apr 16, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants