OAuth: activate filtered ClickHouse roles per request from a JWT claim (#140)

cmd/altinity-mcp/main.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"reflect"
+	"regexp"
 	"strings"
 	"sync"
 	"syscall"
@@ -1175,6 +1176,16 @@ func validateOAuthRuntimeConfig(cfg config.Config) error {
 		}
 	}
 
+	// Per-request role activation: role_filter is OPTIONAL. When set it narrows
+	// which role_claim roles are activated and must be a valid regex; when unset
+	// every role the claim carries is activated (the IdP curates the set, and CH
+	// re-validates the token and enforces grants).
+	if rf := strings.TrimSpace(cfg.Server.OAuth.RoleFilter); rf != "" {
+		if _, err := regexp.Compile(rf); err != nil {
+			return fmt.Errorf("oauth: role_filter is not a valid regex: %w", err)
+		}
+	}
+
 	return nil
 }
 

cmd/altinity-mcp/main_test.go

@@ -3409,6 +3409,42 @@ func TestValidateOAuthRuntimeConfig(t *testing.T) {
 		err := validateOAuthRuntimeConfig(cfg)
 		require.ErrorContains(t, err, "token_url")
 	})
+
+	t.Run("role_claim_without_filter_ok", func(t *testing.T) {
+		t.Parallel()
+		o := resourceServerBase()
+		o.RoleClaim = "https://clickhouse/roles" // role_filter intentionally empty
+		cfg := config.Config{
+			Server:     config.ServerConfig{OAuth: o},
+			ClickHouse: config.ClickHouseConfig{Protocol: config.HTTPProtocol},
+		}
+		require.NoError(t, validateOAuthRuntimeConfig(cfg))
+	})
+
+	t.Run("role_filter_valid_regex_ok", func(t *testing.T) {
+		t.Parallel()
+		o := resourceServerBase()
+		o.RoleClaim = "https://clickhouse/roles"
+		o.RoleFilter = "^anon_"
+		cfg := config.Config{
+			Server:     config.ServerConfig{OAuth: o},
+			ClickHouse: config.ClickHouseConfig{Protocol: config.HTTPProtocol},
+		}
+		require.NoError(t, validateOAuthRuntimeConfig(cfg))
+	})
+
+	t.Run("role_filter_invalid_regex_rejected", func(t *testing.T) {
+		t.Parallel()
+		o := resourceServerBase()
+		o.RoleClaim = "https://clickhouse/roles"
+		o.RoleFilter = "(["
+		cfg := config.Config{
+			Server:     config.ServerConfig{OAuth: o},
+			ClickHouse: config.ClickHouseConfig{Protocol: config.HTTPProtocol},
+		}
+		err := validateOAuthRuntimeConfig(cfg)
+		require.ErrorContains(t, err, "role_filter is not a valid regex")
+	})
 }
 
 func TestWarnOAuthMisconfiguration(t *testing.T) {

docs/oauth_authorization.md

@@ -272,6 +272,17 @@ server:
     # Scopes required in every incoming bearer JWT
     required_scopes: []
 
+    # Per-request ClickHouse role activation. role_claim names a JWT claim
+    # holding a JSON array of role names, activated per request via HTTP role=
+    # params. role_filter is OPTIONAL: when set, only the matching subset is
+    # activated (narrowing); when empty, all roles in the claim are activated.
+    # Either way an empty resolved set fails closed (request denied, no fallback
+    # to the full grant). Leave role_claim empty to disable the feature. Anchor
+    # role_filter (^…/…$) — it is a partial match, so an unanchored "anon" would
+    # also match "not_anon_real".
+    role_claim: ""        # e.g. "https://clickhouse/roles"
+    role_filter: ""       # optional; e.g. "^anon_" (empty = activate all claim roles)
+
     # Token lifetimes (broker mode)
     access_token_ttl_seconds: 3600
     refresh_token_ttl_seconds: 2592000   # 30 d
@@ -300,6 +311,8 @@ server:
 | `public_resource_url` | Externally visible MCP endpoint URL. **Required** behind a reverse proxy. |
 | `public_auth_server_url` | Externally visible OAuth AS URL. **Required** behind a reverse proxy when `broker: true`. |
 | `upstream_offline_access` | Request `offline_access` upstream so the IdP consent screen offers long-lived sessions. Default `false`. |
+| `role_claim` | JWT claim holding a JSON array of ClickHouse role names to activate per request (e.g. `https://clickhouse/roles`). Empty disables. Read from the validated token's namespaced/custom claims. |
+| `role_filter` | **Optional** regex narrowing which `role_claim` roles are activated (e.g. `_mcp$`). When empty, **all** roles in the claim are activated (the IdP curates the set; CH re-validates the token and enforces grants). Only ever narrows; an empty *resolved* set (claim absent/empty, or filter matched nothing) fails closed — request denied, no fallback to the full grant. HTTP protocol only; applies in both Bearer and Basic/sidecar paths. **Anchor it** (`^…`/`…$`) — it's a partial match, so an unanchored `anon` would also match `not_anon_real`. |
 
 ## Provider-specific setup
 

go.mod

@@ -20,7 +20,7 @@ require golang.org/x/crypto v0.52.0 // indirect
 
 require (
 	github.com/ClickHouse/ch-go v0.71.0 // indirect
-	github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260527145654-bdefa859fd1b
+	github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260615173123-b30b0552560b
 	github.com/andybalholm/brotli v1.2.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect

go.sum

@@ -4,8 +4,8 @@ github.com/ClickHouse/ch-go v0.71.0 h1:bUdZ/EZj/LcVHsMqaRUP2holqygrPWQKeMjc6nZoy
 github.com/ClickHouse/ch-go v0.71.0/go.mod h1:NwbNc+7jaqfY58dmdDUbG4Jl22vThgx1cYjBw0vtgXw=
 github.com/ClickHouse/clickhouse-go/v2 v2.46.0 h1:s3eRy+hYmu5uzotB6ZhDofgHu8kDgGN/fpmjxRkqSpk=
 github.com/ClickHouse/clickhouse-go/v2 v2.46.0/go.mod h1:giJfUVlMkcfUEPVfRpt51zZaGEx9i17gCos8gBl392c=
-github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260527145654-bdefa859fd1b h1:b6XHGgSAedsoRet5sahr4YABFC2tG+0Imd+8JGx16SY=
-github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260527145654-bdefa859fd1b/go.mod h1:yLgv0x586vIzPB5JaA6DkqQsSe4PLRT3PhU2iHO7qsI=
+github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260615173123-b30b0552560b h1:2CHdIO/kNhgq077KkhReHLonr/HliNHit8wmKq6wDcU=
+github.com/altinity/go-mcp-oauth-sdk v0.1.1-0.20260615173123-b30b0552560b/go.mod h1:yLgv0x586vIzPB5JaA6DkqQsSe4PLRT3PhU2iHO7qsI=
 github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
 github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=

helm/altinity-mcp/values.yaml

@@ -216,6 +216,18 @@ config:
       scopes: []
       # -- Required scopes for access (token must have all of these)
       required_scopes: []
+      # -- JWT claim holding a JSON array of ClickHouse role names to activate
+      # per request (e.g. "https://clickhouse/roles"). When set, only the
+      # role_filter-matching subset is activated via HTTP `role=` params,
+      # narrowing the user's active roles. Empty disables the feature.
+      role_claim: ""
+      # -- Optional regex narrowing which role_claim roles are activated
+      # (e.g. "^anon_"). When empty, all roles in the claim are activated (the
+      # IdP curates the set; CH re-validates the token and enforces grants). An
+      # empty resolved set fails closed (request denied), never falling back to
+      # the full grant. Anchor it (^…/…$): it is a partial match, so an
+      # unanchored "anon" would also match "not_anon_real".
+      role_filter: ""
     # Dynamic tools generated from ClickHouse views
     dynamic_tools: []
     #  - regexp: "db\\..*"

pkg/clickhouse/client.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"net/http"
 	"os"
 	"regexp"
 	"strings"
@@ -130,7 +131,7 @@ func (c *Client) connect() error {
 		Password: c.config.Password,
 	}
 
-	conn, openErr := clickhouse.Open(&clickhouse.Options{
+	opts := &clickhouse.Options{
 		Addr:            []string{fmt.Sprintf("%s:%d", c.config.Host, c.config.Port)},
 		Auth:            auth,
 		TLS:             tlsConfig,
@@ -143,7 +144,22 @@ func (c *Client) connect() error {
 		MaxIdleConns:    5,
 		ConnMaxLifetime: time.Hour,
 		DialStrategy:    dialWithoutQueryDeadline,
-	})
+	}
+
+	// Per-request role activation: append repeated `role=` query params to every
+	// HTTP request. The driver's Settings map is single-valued and serialized
+	// with url.Values.Set, so it can't express `?role=a&role=b`; wrapping the
+	// transport is the supported way. TransportFunc is HTTP-only in the driver
+	// (native TCP ignores it) — OAuth mode runs over HTTP, and the caller
+	// refuses role activation on non-HTTP protocols.
+	if len(c.config.Roles) > 0 {
+		roles := append([]string(nil), c.config.Roles...)
+		opts.TransportFunc = func(t *http.Transport) (http.RoundTripper, error) {
+			return &roleRoundTripper{wrapped: t, roles: roles}, nil
+		}
+	}
+
+	conn, openErr := clickhouse.Open(opts)
 
 	if openErr != nil {
 		log.Error().
@@ -183,6 +199,26 @@ func dialWithoutQueryDeadline(ctx context.Context, connID int, opt *clickhouse.O
 	return clickhouse.DefaultDialStrategy(context.WithoutCancel(ctx), connID, opt, dial)
 }
 
+// roleRoundTripper appends a `role=` query param per configured role to every
+// outgoing HTTP request, activating exactly those roles for the request
+// (ClickHouse honours repeated role params, e.g. ?role=a&role=b). It wraps the
+// driver's default transport; the request is cloned so the shared *http.Request
+// is never mutated (the client may retry).
+type roleRoundTripper struct {
+	wrapped http.RoundTripper
+	roles   []string
+}
+
+func (rt *roleRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	r2 := req.Clone(req.Context())
+	q := r2.URL.Query()
+	for _, role := range rt.roles {
+		q.Add("role", role)
+	}
+	r2.URL.RawQuery = q.Encode()
+	return rt.wrapped.RoundTrip(r2)
+}
+
 func prepareHTTPAuthForClickHouse(cfg config.ClickHouseConfig) (map[string]string, clickhouse.GetJWTFunc) {
 	if len(cfg.HttpHeaders) == 0 {
 		return nil, nil

pkg/clickhouse/client_roles_test.go

@@ -0,0 +1,142 @@
+package clickhouse
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"testing"
+
+	"github.com/altinity/altinity-mcp/internal/testutil/embeddedch"
+	"github.com/stretchr/testify/require"
+)
+
+// captureRT records the RawQuery of the request it receives and returns an
+// empty 200 response. It stands in for the driver's real transport.
+type captureRT struct{ gotQuery string }
+
+func (c *captureRT) RoundTrip(req *http.Request) (*http.Response, error) {
+	c.gotQuery = req.URL.RawQuery
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(strings.NewReader("")),
+		Header:     make(http.Header),
+		Request:    req,
+	}, nil
+}
+
+// TestRoleRoundTripper pins the multi-role wire encoding: each configured role
+// is appended as its own repeated `role=` query param, existing params are
+// preserved, and the caller's request is not mutated (so retries are safe).
+func TestRoleRoundTripper(t *testing.T) {
+	t.Parallel()
+	cap := &captureRT{}
+	rt := &roleRoundTripper{wrapped: cap, roles: []string{"sandbox_mcp", "readonly_mcp"}}
+
+	req, err := http.NewRequest(http.MethodPost, "http://ch:8123/?database=default", strings.NewReader("SELECT 1"))
+	require.NoError(t, err)
+	origQuery := req.URL.RawQuery
+
+	resp, err := rt.RoundTrip(req)
+	require.NoError(t, err)
+	require.NoError(t, resp.Body.Close())
+
+	// The shared request object must be untouched.
+	require.Equal(t, origQuery, req.URL.RawQuery)
+
+	vals, err := url.ParseQuery(cap.gotQuery)
+	require.NoError(t, err)
+	require.Equal(t, []string{"sandbox_mcp", "readonly_mcp"}, vals["role"], "both roles sent as repeated params")
+	require.Equal(t, "default", vals.Get("database"), "pre-existing query params preserved")
+}
+
+// roleTestUsersXML enables SQL access management on the default user so the
+// test can CREATE ROLE / CREATE USER / GRANT. embedded-clickhouse ships no
+// users.xml of its own, so we provide a minimal one.
+const roleTestUsersXML = `<?xml version="1.0"?>
+<clickhouse>
+    <users>
+        <default>
+            <password></password>
+            <networks><ip>::/0</ip></networks>
+            <profile>default</profile>
+            <quota>default</quota>
+            <access_management>1</access_management>
+        </default>
+    </users>
+    <profiles><default/></profiles>
+    <quotas><default/></quotas>
+</clickhouse>
+`
+
+// TestClientRoleActivation proves end-to-end that config.Roles activates only
+// the named roles for a request: currentRoles() returns exactly the activated
+// set, and a query needing a non-activated role's privilege is denied.
+func TestClientRoleActivation(t *testing.T) {
+	t.Parallel()
+	// SQL CREATE ROLE/USER needs a writeable access storage; point
+	// access_control_path at a temp dir (users.xml alone is read-only).
+	accessDropIn := fmt.Sprintf(
+		"<clickhouse><access_control_path>%s/</access_control_path></clickhouse>",
+		t.TempDir(),
+	)
+	cfg := embeddedch.Setup(t,
+		embeddedch.WithUsersXML(roleTestUsersXML),
+		embeddedch.WithConfigDropIn(accessDropIn),
+	) // HTTP protocol by default
+	ctx := context.Background()
+
+	admin, err := NewClient(ctx, *cfg)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, admin.Close()) }()
+
+	exec := func(q string) {
+		t.Helper()
+		res, err := admin.ExecuteQuery(ctx, q)
+		require.NoError(t, err, q)
+		require.Empty(t, res.Error, q)
+	}
+	exec("CREATE DATABASE IF NOT EXISTS roletest")
+	exec("CREATE TABLE IF NOT EXISTS roletest.public (x UInt64) ENGINE = Memory")
+	exec("CREATE TABLE IF NOT EXISTS roletest.secret (x UInt64) ENGINE = Memory")
+	exec("CREATE ROLE IF NOT EXISTS r_public_mcp")
+	exec("CREATE ROLE IF NOT EXISTS r_secret_real")
+	exec("GRANT SELECT ON roletest.public TO r_public_mcp")
+	exec("GRANT SELECT ON roletest.secret TO r_secret_real")
+	exec("CREATE USER IF NOT EXISTS ruser IDENTIFIED WITH no_password HOST ANY")
+	exec("GRANT r_public_mcp, r_secret_real TO ruser")
+
+	// Base config for the restricted user, activating only the sandbox role.
+	userCfg := *cfg
+	userCfg.Username = "ruser"
+	userCfg.Password = ""
+	userCfg.Roles = []string{"r_public_mcp"}
+
+	cl, err := NewClient(ctx, userCfg)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, cl.Close()) }()
+
+	t.Run("narrows_active_roles", func(t *testing.T) {
+		res, err := cl.ExecuteQuery(ctx, "SELECT arrayStringConcat(arraySort(currentRoles()), ',') AS roles")
+		require.NoError(t, err)
+		require.Empty(t, res.Error)
+		require.Len(t, res.Rows, 1)
+		require.Equal(t, "r_public_mcp", fmt.Sprint(res.Rows[0][0]),
+			"only the activated role is current, not the full granted set")
+	})
+
+	t.Run("activated_role_privilege_allowed", func(t *testing.T) {
+		res, err := cl.ExecuteQuery(ctx, "SELECT count() FROM roletest.public")
+		require.NoError(t, err)
+		require.Empty(t, res.Error)
+	})
+
+	t.Run("non_activated_role_privilege_denied", func(t *testing.T) {
+		_, err := cl.ExecuteQuery(ctx, "SELECT count() FROM roletest.secret")
+		require.Error(t, err, "secret table needs r_secret_real, which was not activated")
+		require.Contains(t, strings.ToLower(err.Error()), "privileg",
+			"expected an ACCESS_DENIED / not-enough-privileges error")
+	})
+}

pkg/config/config.go

@@ -53,6 +53,11 @@ type ClickHouseConfig struct {
 	MaxResultBytes int               `json:"max_result_bytes,omitempty" yaml:"max_result_bytes,omitempty" flag:"clickhouse-max-result-bytes" env:"CLICKHOUSE_MAX_RESULT_BYTES" desc:"Per-request approximate byte cap on result body (0=default 50000, <0=disable)"`
 	HttpHeaders    map[string]string `json:"http_headers" yaml:"http_headers" flag:"clickhouse-http-headers" env:"CLICKHOUSE_HTTP_HEADERS" desc:"HTTP Headers for ClickHouse"`
 	ExtraSettings  map[string]string `json:"extra_settings,omitempty" yaml:"extra_settings,omitempty" desc:"Per-request ClickHouse settings injected by tool_input_settings"`
+	// Roles is the per-request set of ClickHouse roles to activate for this
+	// request (sent as HTTP `role=` query params). Populated at request time
+	// from the OAuth role_claim/role_filter; never set from CLI/file/env. Only
+	// narrows the user's active roles. HTTP protocol only.
+	Roles []string `json:"-" yaml:"-"`
 	// MaxQueryLength caps the size in bytes of a single SQL query string sent by a client.
 	// Default 10 MB when 0. Set to a negative number to disable the check.
 	MaxQueryLength int `json:"max_query_length,omitempty" yaml:"max_query_length,omitempty" flag:"clickhouse-max-query-length" env:"CLICKHOUSE_MAX_QUERY_LENGTH" desc:"Max bytes of SQL query string accepted from clients (0=default 10MB, <0=disabled)"`

pkg/server/server.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"regexp"
 	"strings"
 	"sync"
 
@@ -33,6 +34,12 @@ type ClickHouseJWEServer struct {
 	// keyed by "host:port". Populated on the first OAuth request to an endpoint;
 	// cleared on config reload so operator changes take effect without restart.
 	chOAuthMethodCache sync.Map
+	// roleFilterRe is the compiled oauth.role_filter, used to narrow the roles
+	// activated per ClickHouse request. Compiled up-front in
+	// NewClickHouseMCPServer when role_claim is set; the roleFilter() getter
+	// lazily builds it for struct-literal test servers. nil when the feature
+	// is unconfigured.
+	roleFilterRe *regexp.Regexp
 }
 
 // ToolHandlerFunc is a function type for tool handlers
@@ -75,6 +82,19 @@ func NewClickHouseMCPServer(cfg config.Config, version string) *ClickHouseJWESer
 		oauthVerifier:  oauth.NewVerifier(cfg.Server.OAuth),
 		blockedClauses: NormalizeBlockedClauses(cfg.Server.BlockedQueryClauses),
 	}
+	// Compile the role filter up-front when role activation is configured. An
+	// empty role_filter compiles to a match-all regex (filtering is optional —
+	// every role_claim role is activated). validateOAuthRuntimeConfig already
+	// rejected an invalid pattern, so a compile error here is not expected;
+	// leave roleFilterRe nil if it ever occurs so role activation fails closed.
+	if strings.TrimSpace(cfg.Server.OAuth.RoleClaim) != "" {
+		pat := strings.TrimSpace(cfg.Server.OAuth.RoleFilter)
+		if re, err := regexp.Compile(pat); err == nil {
+			chJweServer.roleFilterRe = re
+		} else {
+			log.Error().Err(err).Str("role_filter", pat).Msg("oauth: failed to compile role_filter; per-request role activation will fail closed")
+		}
+	}
 
 	// Register tools, resources, and prompts.
 	// Pass pointer to the server's Config so RegisterTools can store converted