Skip to content

feat: per-request OAuth pass-through for Snowflake connections#650

Open
namabile wants to merge 1 commit intomalloydata:mainfrom
namabile:feat/oauth-passthrough
Open

feat: per-request OAuth pass-through for Snowflake connections#650
namabile wants to merge 1 commit intomalloydata:mainfrom
namabile:feat/oauth-passthrough

Conversation

@namabile
Copy link
Copy Markdown
Contributor

@namabile namabile commented Apr 1, 2026

Summary

Adds opt-in per-request OAuth token pass-through for Snowflake connections, enabling queries to execute under the authenticated user's identity instead of a shared service account.

Resolves #629

Motivation

Publisher's shared connection model blocks adoption in enterprise environments where Snowflake row-level security, column masking, and audit trails depend on CURRENT_USER. Every major BI tool supports per-user identity pass-through — this brings that capability to Publisher.

What changed

  • request_context.ts (new) — AsyncLocalStorage context to carry the X-Database-Token header through the request lifecycle
  • server.ts — MCP POST handler extracts the token header and wraps execution in request context
  • connection.ts — New oauthPassthrough flag on Snowflake config; getRequestConnections() creates per-request OAuth connections when a token is present
  • model.ts — Query execution uses per-request connections when available; OAuth connections are closed in finally to prevent session leaks
  • api-doc.yaml — Documents the oauthPassthrough config field

+573 / -133 lines across 8 files (includes tests)

Design decisions

  • Opt-in per connection — Only Snowflake connections with oauthPassthrough: true are affected. All other connections and requests without the header work exactly as before.
  • No auth middleware — Publisher doesn't validate the token. That's the gateway's responsibility (consistent with the MCP auth spec recommendation for gateway-mediated token exchange).
  • Short-lived connections — Per-request connections are created and closed within the query lifecycle to avoid Snowflake session accumulation.
  • Snowflake-first — The snowflake-sdk natively supports authenticator: "OAUTH". BigQuery/Trino support can follow the same pattern later.

Test plan

  • Unit tests for AsyncLocalStorage request context (request_context.spec.ts)
  • Unit tests for OAuth config storage and getRequestConnections() map swapping (connection.spec.ts)
  • Unit tests for Model constructor changes and query execution with OAuth connections (model.spec.ts)
  • Python SDK regenerated — no drift from api-doc.yaml changes
  • Integration test with a real Snowflake OAuth token (requires Snowflake OAuth client setup)
  • Verify backward compatibility: requests without X-Database-Token header use shared connections unchanged

Signed-off-by: namabile nick@ultrathinksolutions.com

Add support for per-request OAuth token injection on Snowflake
connections. When a connection has `oauthPassthrough: true`, queries
execute under the authenticated user's Snowflake role instead of
using static credentials.

How it works:
- MCP clients send an OAuth token via the `X-Database-Token` HTTP header
- The token propagates through AsyncLocalStorage (request_context.ts)
  so downstream code can access it without parameter threading
- At query time, `getQueryResults()` creates a short-lived Snowflake
  connection using `authenticator: "OAUTH"` with the per-request token
- Per-request connections are closed in a `finally` block to prevent
  Snowflake session leaks

Key changes:
- api-doc.yaml: Add `oauthPassthrough` boolean to SnowflakeConnection
- request_context.ts: AsyncLocalStorage module for per-request tokens
- server.ts: Extract X-Database-Token header, run in request context
- connection.ts: OAuth config registry, per-request connection factory
- model.ts: Create per-request Runtime with OAuth connections in
  getQueryResults(), with cleanup in finally block

Includes unit tests for all new functions.

Signed-off-by: namabile <nick@ultrathinksolutions.com>
@namabile namabile force-pushed the feat/oauth-passthrough branch from 02b90e8 to 8878068 Compare April 1, 2026 21:23
@namabile
Copy link
Copy Markdown
Contributor Author

namabile commented Apr 1, 2026

Integration testing note

This feature has been validated end-to-end in a production-like environment with:

  • Okta as the OIDC identity provider (custom authorization server with Snowflake audience)
  • Snowflake External OAuth configured to accept Okta-issued JWTs
  • MCP gateway performing OIDC validation and forwarding the user's OAuth token via X-Database-Token header
  • Snowflake SNOWFLAKE_SAMPLE_DATA.TPCH_SF1 as the test dataset

Verified behaviors:

  • Queries execute as the authenticated user's Snowflake role (confirmed via CURRENT_USER() and CURRENT_ROLE())
  • Snowflake RLS and column masking policies apply correctly per-user
  • Audit logs show the actual user identity, not the service account
  • Fallback to shared connection works when no token header is present (backward compatible)
  • Per-request connections close cleanly after query completion (no session accumulation)
  • All 54 unit tests pass

The CI failure on "Connection Integration Tests" is expected — fork PRs don't have access to the upstream repo's BQ_PRESTO_TRINO_KEY secret. This is not related to our changes.

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.

feat: Per-request database connections via OAuth token for multi-user deployments

1 participant