feat: Implement Trino user impersonation and query attribution #132
Conversation
This commit introduces user impersonation for Trino, allowing queries to be executed on behalf of authenticated OAuth users. This enhances auditability and enables Trino's native access control to enforce user-specific permissions. Key features include: - **User Impersonation**: MCP now propagates the authenticated OAuth user's identity to Trino via the `X-Trino-User` header. This ensures that Trino logs and access control reflect the actual user executing the query, rather than a generic service account. - **Configurable Principal Field**: Users can configure which JWT claim (username, email, or subject) is used as the impersonated user in Trino. - **Query Attribution**: MCP now consistently sets the `X-Trino-Source` header, identifying `mcp-trino` and its version as the source of queries, improving monitoring and debugging. - **Comprehensive Documentation**: A new `impersonation.md` guide provides detailed setup instructions, configuration options, security considerations, and troubleshooting for Trino impersonation. Closes: tuannvm#118
WalkthroughAdds Trino user impersonation: new TrinoConfig fields and version-aware constructor, context helpers to carry an impersonated username, headerRoundTripper injecting X-Trino-User/X-Trino-Source, handlers and client switched to context-aware calls, new docs and unit tests for impersonation helpers. Changes
Sequence Diagram(s)sequenceDiagram
participant User as OAuth User
participant MCP as MCP Server
participant Config as TrinoConfig
participant TrinoClient as Trino Client
participant Trino as Trino Server
User->>MCP: Authenticate (JWT)
MCP->>Config: Read TRINO_ENABLE_IMPERSONATION / TRINO_IMPERSONATION_FIELD
Config-->>MCP: Return enable flag & field
rect rgb(230,245,230)
MCP->>MCP: Extract username from JWT claim (configured field)
MCP->>MCP: ctx = WithImpersonatedUser(ctx, username)
end
MCP->>TrinoClient: ExecuteQueryWithContext(ctx, query)
TrinoClient->>TrinoClient: GetImpersonatedUser(ctx)
rect rgb(230,245,245)
TrinoClient->>TrinoClient: headerRoundTripper adds X-Trino-User & X-Trino-Source
end
TrinoClient->>Trino: HTTP request with headers
Trino-->>TrinoClient: Query response
TrinoClient-->>MCP: Return results
MCP-->>User: Return results
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🧰 Additional context used📓 Path-based instructions (2)internal/mcp/handlers.go📄 CodeRabbit inference engine (CLAUDE.md)
Files:
internal/trino/client.go📄 CodeRabbit inference engine (CLAUDE.md)
Files:
🧠 Learnings (6)📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
📚 Learning: 2025-10-20T06:10:42.267ZApplied to files:
🧬 Code graph analysis (2)internal/mcp/handlers.go (2)
internal/trino/client.go (1)
🔇 Additional comments (12)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/mcp/handlers.go (1)
92-265: Consider applying impersonation context to all handler methods.Currently, only
ExecuteQueryprepares and passes impersonation context. However, other handlers (ListCatalogs,ListSchemas,ListTables,GetTableSchema,ExplainQuery) execute Trino queries internally but don't receive or propagate the impersonation context.This means that queries executed by these tools will run as the service account rather than the authenticated user, which may not align with the impersonation feature's goals for consistent audit trails and access control.
Consider updating these handlers to also prepare and pass impersonation context:
func (h *TrinoHandlers) ListCatalogs(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if h.Config.EnableImpersonation { + ctx = h.prepareImpersonationContext(ctx) + } - catalogs, err := h.TrinoClient.ListCatalogs() + catalogs, err := h.TrinoClient.ListCatalogsWithContext(ctx) if err != nil { log.Printf("Error listing catalogs: %v", err)Note: This would also require adding context-aware versions of the methods in the Trino client (e.g.,
ListCatalogsWithContext,ListSchemasWithContext, etc.).
🧹 Nitpick comments (4)
internal/config/config.go (2)
34-39: OIDC and redirect URI config remain consistent and safeThe clarified comments and the struct fields (
OIDCIssuer,OIDCAudience,OIDCClientID,OIDCClientSecret,OAuthRedirectURIs) align with the existing env-driven behavior and avoid logging any secrets while still supporting backward-compatibleOAUTH_REDIRECT_URI. The struct initialization keeps everything wired correctly.If you ever touch this again, you could consider centralizing the redirect-URI deprecation warning in a tiny helper to avoid scattering similar logic, but it’s fine as-is.
Also applies to: 181-203
45-51: Impersonation config and validation look solid; consider minor UX tweakThe impersonation implementation is consistent:
- Env flags:
TRINO_ENABLE_IMPERSONATIONandTRINO_IMPERSONATION_FIELD(lowercased).- Strict validation of
TRINO_IMPERSONATION_FIELDagainstusername|email|subjectwith a clear error message.- Logging clearly indicates when impersonation is enabled/disabled and warns when it’s enabled without OAuth (good for avoiding silent misconfig).
- Values are correctly exposed via
EnableImpersonationandImpersonationFieldinTrinoConfig.One minor UX improvement you might consider later: log a warning when
TRINO_ENABLE_IMPERSONATIONis non-boolean (parse error) instead of silently treating it asfalse, similar to how you already handleTRINO_QUERY_TIMEOUT. This would make mis-typed envs easier to debug. As per coding guidelines.Also applies to: 110-113, 160-175, 204-205
README.md (1)
104-105: Impersonation documentation matches implementation; optional note on TRINO_SOURCEThe README updates correctly describe:
- Enabling impersonation with
TRINO_ENABLE_IMPERSONATION=true.- Selecting the principal via
TRINO_IMPERSONATION_FIELDwith the same allowed values as the code (username,subject).- The fact that impersonation requires OAuth and uses
X-Trino-User, matching the config and client behavior.At some point you might also mention the
TRINO_SOURCEenv var here (query attribution viaX-Trino-Source), since it’s now configurable and useful for observability, but it’s not strictly required for this PR.Also applies to: 160-172
docs/impersonation.md (1)
336-337: Add language specifiers to code blocks.Several code blocks are missing language specifiers, which would improve syntax highlighting and readability.
Apply these changes:
### Check MCP Logs Look for: -``` +```text INFO: Trino user impersonation enabled (TRINO_ENABLE_IMPERSONATION=true) INFO: Impersonation principal field: email MCP: Preparing impersonation context for email: [email protected] Trino: Setting X-Trino-User header to: [email protected]```diff ### Check Trino Logs Queries should show actual users: -``` +```text [email protected] query=SELECT * FROM table [email protected] query=SELECT * FROM table```diff Instead of: -``` +```text user=mcp_service_account query=SELECT * FROM table```diff ### Request Flow -``` +```text User OAuth Login ↓ JWT Token GeneratedAlso applies to: 346-347, 352-353, 420-421
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
README.md(2 hunks)cmd/main.go(1 hunks)docs/impersonation.md(1 hunks)internal/config/config.go(3 hunks)internal/mcp/handlers.go(2 hunks)internal/mcp/server.go(1 hunks)internal/trino/client.go(4 hunks)internal/trino/impersonation_test.go(1 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
cmd/main.go
📄 CodeRabbit inference engine (CLAUDE.md)
cmd/main.go: Main entry point must handle server initialization, transport selection (STDIO vs HTTP with SSE compatibility), graceful shutdown with signal handling, CORS support for web clients, and version/build metadata management
Expose HTTP transport endpoints: StreamableHTTP at /mcp, SSE-compatible endpoint at /sse for backward compatibility, and status endpoint at GET /
Files:
cmd/main.go
internal/mcp/handlers.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/mcp/handlers.go: All MCP tools must return JSON-formatted responses and standardize tool results
Validate tool parameters and implement consistent error handling and logging for debugging
Files:
internal/mcp/handlers.go
internal/config/config.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/config/config.go: Configuration must be environment-based with validation, including gathering OAuth settings (with validation delegated to oauth-mcp-proxy)
Enforce HTTPS scheme to set SSL=true regardless of TRINO_SSL and fall back to a 30s default for invalid TRINO_QUERY_TIMEOUT values (with a warning)
Use OAUTH_ENABLED as the single source of truth for OAuth activation and support OAUTH_MODE (native/proxy) and OAUTH_PROVIDER (hmac/okta/google/azure)
Files:
internal/config/config.go
internal/mcp/server.go
📄 CodeRabbit inference engine (CLAUDE.md)
Register OAuth 2.1 middleware via oauth-mcp-proxy for both native and proxy modes with supported providers (HMAC, Okta, Google, Azure AD)
Files:
internal/mcp/server.go
internal/trino/client.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/trino/client.go: Implement read-only SQL enforcement via isReadOnlyQuery(): allow SELECT/SHOW/DESCRIBE/EXPLAIN/WITH; block INSERT/UPDATE/DELETE/CREATE/DROP/ALTER by default; support override via TRINO_ALLOW_WRITE_QUERIES with a warning
Apply context-based query timeouts using TRINO_QUERY_TIMEOUT (>0)
Configure DB connection pooling to 10 max open, 5 max idle, and 5m max lifetime
Process and format query results consistently before returning to callers
Files:
internal/trino/client.go
🧠 Learnings (8)
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Configure DB connection pooling to 10 max open, 5 max idle, and 5m max lifetime
Applied to files:
cmd/main.gointernal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/mcp/handlers.go : Validate tool parameters and implement consistent error handling and logging for debugging
Applied to files:
internal/mcp/handlers.gointernal/mcp/server.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Process and format query results consistently before returning to callers
Applied to files:
internal/mcp/handlers.gointernal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/mcp/handlers.go : All MCP tools must return JSON-formatted responses and standardize tool results
Applied to files:
internal/mcp/handlers.gointernal/mcp/server.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Apply context-based query timeouts using TRINO_QUERY_TIMEOUT (>0)
Applied to files:
internal/mcp/handlers.gointernal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/config/config.go : Configuration must be environment-based with validation, including gathering OAuth settings (with validation delegated to oauth-mcp-proxy)
Applied to files:
internal/config/config.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/mcp/server.go : Register OAuth 2.1 middleware via oauth-mcp-proxy for both native and proxy modes with supported providers (HMAC, Okta, Google, Azure AD)
Applied to files:
internal/mcp/server.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Implement read-only SQL enforcement via isReadOnlyQuery(): allow SELECT/SHOW/DESCRIBE/EXPLAIN/WITH; block INSERT/UPDATE/DELETE/CREATE/DROP/ALTER by default; support override via TRINO_ALLOW_WRITE_QUERIES with a warning
Applied to files:
internal/trino/client.go
🧬 Code graph analysis (5)
cmd/main.go (1)
internal/config/config.go (1)
NewTrinoConfigWithVersion(59-208)
internal/mcp/handlers.go (2)
internal/trino/client.go (2)
Client(51-56)WithImpersonatedUser(120-122)internal/config/config.go (1)
TrinoConfig(13-51)
internal/mcp/server.go (1)
internal/mcp/handlers.go (1)
NewTrinoHandlers(23-28)
internal/trino/client.go (1)
internal/config/config.go (1)
TrinoConfig(13-51)
internal/trino/impersonation_test.go (1)
internal/trino/client.go (2)
WithImpersonatedUser(120-122)GetImpersonatedUser(125-128)
🪛 LanguageTool
docs/impersonation.md
[style] ~8-~8: It’s more common nowadays to write this noun as one word.
Context: ... Audit trails - Queries show actual user names in Trino logs - ✅ Access control - ...
(RECOMMENDED_COMPOUNDS)
🪛 markdownlint-cli2 (0.18.1)
docs/impersonation.md
336-336: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
346-346: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
352-352: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
420-420: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🔇 Additional comments (12)
cmd/main.go (1)
24-29: Version-aware config wiring looks correctUsing
config.NewTrinoConfigWithVersion(Version)cleanly injects the build-time version into Trino configuration (forTRINO_SOURCEdefaults) without changing startup error handling. No issues from this change.internal/config/config.go (1)
55-56: Version-aware constructor and TrinoSource defaulting are well designed
NewTrinoConfig()delegating toNewTrinoConfigWithVersion("dev")keeps the old API while allowing the main binary to pass the real version. TheTRINO_SOURCEhandling:
- Defaults to
mcp-trino/<version>.- Honors a custom
TRINO_SOURCEwhen set.- Treats an explicitly empty
TRINO_SOURCEas “reset to default”.This is a sensible, predictable scheme, and the explicit log line for query attribution is helpful for operators.
Also applies to: 58-59, 114-120, 177-179, 206-206
internal/mcp/server.go (1)
41-63: Handler construction now correctly injects TrinoConfigUsing
NewTrinoHandlers(trinoClient, trinoConfig)instead of a literal struct ensures handlers see the fullTrinoConfig(including impersonation and query-attribution settings) while keeping OAuth middleware registration unchanged. This aligns well with the new configuration-driven behavior.internal/trino/impersonation_test.go (1)
1-74: Impersonation context tests give good coverageThe tests thoroughly cover the helper behavior:
- Valid, empty, and special-character usernames.
- Unset context behavior (
ok == false).- Preservation of impersonation data through
context.WithTimeout, matching how query execution will derive contexts.No issues here; this is a solid regression suite for the context helpers.
docs/impersonation.md (1)
1-469: Excellent documentation for the impersonation feature!The documentation is comprehensive, well-structured, and provides clear guidance for:
- Configuration and setup
- Provider-specific considerations (Okta, Google, Azure AD)
- Trino access control integration
- Troubleshooting common issues
- Security considerations and implementation details
The examples and step-by-step instructions will help users effectively configure and use the impersonation feature.
internal/mcp/handlers.go (3)
11-12: LGTM! Clean dependency injection pattern.The addition of the
Configfield toTrinoHandlersand the updated constructor properly wire the configuration into the handlers, enabling access to impersonation settings at runtime.Also applies to: 19-20, 23-28
30-51: LGTM! Well-implemented impersonation context preparation.The method correctly:
- Extracts the OAuth user from context
- Selects the appropriate principal field (email, subject, or username)
- Logs impersonation activity for audit trails
- Handles the case where no user is found or principal is empty
- Returns an enriched context with the impersonated user
The defensive check on line 45 prevents adding empty principals to the context, which is good practice.
54-89: LGTM! Proper integration of impersonation into query execution.The changes correctly:
- Conditionally prepare impersonation context only when
EnableImpersonationis true (lines 55-57)- Pass the context to
ExecuteQueryWithContext(line 73), which enables the Trino client to read impersonation data and set theX-Trino-Userheader- Maintain existing error handling and response formatting
The implementation aligns with the documented impersonation flow and coding guidelines.
Based on learnings
internal/trino/client.go (4)
19-24: LGTM! Standard context key pattern.Using a custom type for context keys prevents collisions with other packages that might use string keys in the same context.
26-48: LGTM! Correct RoundTripper implementation.The implementation properly:
- Clones the request to avoid mutating the original (line 33)
- Sets the
X-Trino-Sourceheader for query attribution when configured- Conditionally sets the
X-Trino-Userheader only when impersonation is enabled and a user is present in the context- Delegates to the base transport for the actual HTTP request
This ensures headers are correctly propagated to Trino for impersonation and query attribution.
119-128: LGTM! Standard context helper functions.
WithImpersonatedUserandGetImpersonatedUserfollow the standard Go pattern for storing and retrieving values from context.
274-354: LGTM! Proper context-aware query execution.The implementation correctly:
- Delegates
ExecuteQuerytoExecuteQueryWithContextfor backward compatibility (lines 275-277)- Creates a timeout context from the provided context, preserving impersonation data (line 291)
- Logs impersonation activity when enabled for audit trails (lines 295-299)
- Uses
QueryContextto execute queries with the impersonation-aware context (line 302)This aligns with the coding guidelines for applying context-based query timeouts while maintaining impersonation information throughout the execution path.
Based on learnings
| // Client is a wrapper around Trino client | ||
| type Client struct { | ||
| db *sql.DB | ||
| config *config.TrinoConfig | ||
| timeout time.Duration | ||
| db *sql.DB | ||
| connector driver.Connector | ||
| config *config.TrinoConfig | ||
| timeout time.Duration | ||
| } |
There was a problem hiding this comment.
Unused connector field in Client struct.
The connector field (line 53) is declared in the Client struct but is never initialized in NewClient (lines 107-111 only set db, config, and timeout). This appears to be dead code.
If the field is not needed, remove it:
type Client struct {
db *sql.DB
- connector driver.Connector
config *config.TrinoConfig
timeout time.Duration
}If it's intended for future use, consider adding a comment explaining its purpose.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Client is a wrapper around Trino client | |
| type Client struct { | |
| db *sql.DB | |
| config *config.TrinoConfig | |
| timeout time.Duration | |
| db *sql.DB | |
| connector driver.Connector | |
| config *config.TrinoConfig | |
| timeout time.Duration | |
| } | |
| // Client is a wrapper around Trino client | |
| type Client struct { | |
| db *sql.DB | |
| config *config.TrinoConfig | |
| timeout time.Duration | |
| } |
🤖 Prompt for AI Agents
In internal/trino/client.go around lines 50 to 56, the Client struct declares an
unused connector field; remove the dead field if it's not required, or if it's
intended for future use either initialize it in NewClient (set connector when
creating the client) or add a clear comment explaining its planned purpose and
why it's left uninitialized; make the minimal change (delete the field or add
initialization/comment) and run `go vet`/`go build` to ensure no references
remain.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
internal/trino/client.go (1)
50-56: Remove the unusedconnectorfield.The
connectorfield (line 53) is declared but never initialized inNewClient(lines 107-111 only setdb,config, andtimeout). This issue was previously flagged and remains unresolved.Apply this diff:
type Client struct { db *sql.DB - connector driver.Connector config *config.TrinoConfig timeout time.Duration }
🧹 Nitpick comments (2)
internal/trino/client.go (1)
26-48: Verify base RoundTripper is not nil before use.The headerRoundTripper correctly clones the request and sets X-Trino-Source and X-Trino-User headers based on configuration and context. However, if
t.baseis nil,RoundTripwill panic at runtime.Apply this defensive check:
func (t *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if t.base == nil { + return nil, fmt.Errorf("base RoundTripper is nil") + } req = req.Clone(req.Context()) // Set X-Trino-Source header for query attributioninternal/mcp/handlers.go (1)
55-57: Repetitive impersonation pattern is acceptable but consider middleware.The pattern
if h.Config.EnableImpersonation { ctx = h.prepareImpersonationContext(ctx) }is repeated in all six handler methods. While this is functional and clear, it could potentially be extracted to middleware to reduce duplication and ensure consistency.The current approach is acceptable for this PR. If you want to refactor, consider creating a wrapper middleware function:
func (h *TrinoHandlers) withImpersonation(handler func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error)) func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { if h.Config.EnableImpersonation { ctx = h.prepareImpersonationContext(ctx) } return handler(ctx, request) } }Then wrap each handler during registration. However, this adds complexity and the current approach is more explicit.
Also applies to: 93-95, 116-118, 152-154, 191-193, 240-242
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
internal/mcp/handlers.go(11 hunks)internal/trino/client.go(11 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
internal/mcp/handlers.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/mcp/handlers.go: All MCP tools must return JSON-formatted responses and standardize tool results
Validate tool parameters and implement consistent error handling and logging for debugging
Files:
internal/mcp/handlers.go
internal/trino/client.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/trino/client.go: Implement read-only SQL enforcement via isReadOnlyQuery(): allow SELECT/SHOW/DESCRIBE/EXPLAIN/WITH; block INSERT/UPDATE/DELETE/CREATE/DROP/ALTER by default; support override via TRINO_ALLOW_WRITE_QUERIES with a warning
Apply context-based query timeouts using TRINO_QUERY_TIMEOUT (>0)
Configure DB connection pooling to 10 max open, 5 max idle, and 5m max lifetime
Process and format query results consistently before returning to callers
Files:
internal/trino/client.go
🧠 Learnings (6)
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/mcp/handlers.go : Validate tool parameters and implement consistent error handling and logging for debugging
Applied to files:
internal/mcp/handlers.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Process and format query results consistently before returning to callers
Applied to files:
internal/mcp/handlers.gointernal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/mcp/handlers.go : All MCP tools must return JSON-formatted responses and standardize tool results
Applied to files:
internal/mcp/handlers.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Apply context-based query timeouts using TRINO_QUERY_TIMEOUT (>0)
Applied to files:
internal/mcp/handlers.gointernal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Configure DB connection pooling to 10 max open, 5 max idle, and 5m max lifetime
Applied to files:
internal/trino/client.go
📚 Learning: 2025-10-20T06:10:42.267Z
Learnt from: CR
Repo: tuannvm/mcp-trino PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-10-20T06:10:42.267Z
Learning: Applies to internal/trino/client.go : Implement read-only SQL enforcement via isReadOnlyQuery(): allow SELECT/SHOW/DESCRIBE/EXPLAIN/WITH; block INSERT/UPDATE/DELETE/CREATE/DROP/ALTER by default; support override via TRINO_ALLOW_WRITE_QUERIES with a warning
Applied to files:
internal/trino/client.go
🧬 Code graph analysis (2)
internal/mcp/handlers.go (2)
internal/trino/client.go (2)
Client(51-56)WithImpersonatedUser(120-122)internal/config/config.go (1)
TrinoConfig(13-51)
internal/trino/client.go (1)
internal/config/config.go (1)
TrinoConfig(13-51)
🔇 Additional comments (7)
internal/trino/client.go (5)
19-24: LGTM! Clean context key definition.The unexported context key type and constant follow Go best practices for context values, preventing key collisions.
119-128: LGTM! Context helpers follow Go patterns.The
WithImpersonatedUserandGetImpersonatedUserfunctions provide a clean API for managing impersonation data in contexts, consistent with standard Go context patterns.
275-277: LGTM! Clean delegation pattern.The existing
ExecuteQuerymethod correctly delegates to the new context-aware variant withcontext.Background(), maintaining backward compatibility.
357-359: LGTM! Consistent delegation across all public methods.All public List*/Get*/Explain methods correctly delegate to their WithContext counterparts using
context.Background(), maintaining backward compatibility while enabling context-aware calls.Also applies to: 384-386, 416-418, 451-453, 494-496
60-83: Add protection against multiple RegisterCustomClient calls — library design discourages repeated registrations.While RegisterCustomClient's implementation is technically idempotent (it simply overwrites the map entry without error), the library design indicates repeated calls are not intended: the API returns an error type and provides a separate DeregisterCustomClient function, signaling that registrations should be managed explicitly rather than assumed idempotent. If NewClient is called multiple times (e.g., in tests, multiple instances, or restarts), guard the registration using sync.Once to follow library conventions and prevent unintended client replacement:
var registerOnce sync.Once registerOnce.Do(func() { trino.RegisterCustomClient("mcp-trino", httpClient) })Alternatively, document that NewClient is designed for single instantiation only.
⛔ Skipped due to learnings
Learnt from: CR Repo: tuannvm/mcp-trino PR: 0 File: CLAUDE.md:0-0 Timestamp: 2025-10-20T06:10:42.267Z Learning: Applies to internal/trino/client.go : Configure DB connection pooling to 10 max open, 5 max idle, and 5m max lifetimeinternal/mcp/handlers.go (2)
17-28: LGTM! Clean constructor update for impersonation support.The addition of the
Configfield and updatedNewTrinoHandlersconstructor properly support the impersonation feature by providing access to configuration settings.
73-73: LGTM! All Trino client calls use context-aware variants.All handlers correctly invoke the
WithContextmethods on the Trino client, properly propagating the impersonation context through the request lifecycle. This aligns with the PR objectives for user impersonation.Also applies to: 97-97, 133-133, 172-172, 221-221, 265-265
|
thanks @srstrickland |
2 participants
This commit introduces user impersonation for Trino, allowing queries to be executed on behalf of authenticated OAuth users. This enhances auditability and enables Trino's native access control to enforce user-specific permissions.
Key features include:
X-Trino-Userheader. This ensures that Trino logs and access control reflect the actual user executing the query, rather than a generic service account.X-Trino-Sourceheader, identifyingmcp-trinoand its version as the source of queries, improving monitoring and debugging.impersonation.mdguide provides detailed setup instructions, configuration options, security considerations, and troubleshooting for Trino impersonation.Closes: #118
Summary by CodeRabbit
New Features
Documentation
Tests