Feature/add trino source header #90
Conversation
WalkthroughAdds a TrinoSource field to TrinoConfig (loaded from TRINO_SOURCE), updates the Trino client to optionally send an X-Trino-Source header when TrinoSource is non-empty, and adds unit tests for config defaults and header behavior. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant App as Caller
participant Client as TrinoClient
participant Cfg as TrinoConfig
participant DB as Trino Driver
App->>Client: Execute(ctx, query)
Client->>Cfg: check TrinoSource
alt TrinoSource non-empty
Note over Client,DB: attach X-Trino-Source header
Client->>DB: QueryContext(ctx, query, Named("X-Trino-Source", TrinoSource))
else TrinoSource empty
Client->>DB: QueryContext(ctx, query)
end
DB-->>Client: rows / error
alt error
Client-->>App: wrapped error
else success
Client->>Client: process rows (defer close)
Client-->>App: results
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
internal/config/config.go (1)
107-109: Trim and sanitize TRINO_SOURCE to avoid header formatting pitfalls.Leading/trailing spaces or accidental newlines can produce invalid headers.
Apply:
-// Get Trino Source from env/config (no default) -trinoSource := getEnv("TRINO_SOURCE", "") +// Get Trino Source from env/config (no default) and sanitize +trinoSource := strings.TrimSpace(getEnv("TRINO_SOURCE", "")) +// Disallow CR/LF to prevent malformed headers +trinoSource = strings.Map(func(r rune) rune { + if r == '\r' || r == '\n' { + return -1 + } + return r +}, trinoSource)internal/config/config_test.go (1)
8-52: Use t.Setenv to avoid env leakage between subtests; add a whitespace-sanitization case.
t.Setenvsimplifies cleanup and makes tests safer when run in parallel; also add a case with surrounding spaces/newlines if you adopt trimming.-// Clean up environment -os.Unsetenv("TRINO_SOURCE") -// Set environment variable if not UNSET -if tt.envValue != "UNSET" { - os.Setenv("TRINO_SOURCE", tt.envValue) - defer os.Unsetenv("TRINO_SOURCE") -} +// Set env (auto-cleaned) unless UNSET sentinel +if tt.envValue != "UNSET" { + t.Setenv("TRINO_SOURCE", tt.envValue) +} else { + os.Unsetenv("TRINO_SOURCE") +}Optionally add:
{name: "TRINO_SOURCE with spaces/newline", envValue: " app\n", expectedSource: "app"},internal/trino/client.go (1)
129-139: Consider accepting a context to propagate caller deadlines.Expose
ExecuteQueryContext(ctx, query)and haveExecuteQuerycall it withcontext.Background(). This respects upstream timeouts/cancellations while still applying your per-client timeout.func (c *Client) ExecuteQueryContext(parent context.Context, query string) ([]map[string]interface{}, error) { ctx, cancel := context.WithTimeout(parent, c.timeout) defer cancel() // ... } func (c *Client) ExecuteQuery(query string) ([]map[string]interface{}, error) { return c.ExecuteQueryContext(context.Background(), query) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
internal/config/config.go(3 hunks)internal/config/config_test.go(1 hunks)internal/trino/client.go(1 hunks)internal/trino/client_test.go(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
internal/config/config.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/config/config.go: Use environment-based configuration with validation for all connection/security settings
Force SSL when TRINO_SCHEME=https regardless of TRINO_SSL
Validate TRINO_QUERY_TIMEOUT > 0, otherwise fallback to 30s and log a warning
Files:
internal/config/config.go
internal/trino/client.go
📄 CodeRabbit inference engine (CLAUDE.md)
internal/trino/client.go: Enforce read-only SQL via isReadOnlyQuery(): allow SELECT/SHOW/DESCRIBE/EXPLAIN/WITH; block INSERT/UPDATE/DELETE/CREATE/DROP/ALTER unless TRINO_ALLOW_WRITE_QUERIES=true (with warning)
Implement database connection pooling (10 max open, 5 max idle, 5m max lifetime)
Use context-based timeouts for query execution and handle deadline propagation
Files:
internal/trino/client.go
🧠 Learnings (1)
📚 Learning: 2025-09-01T14:40:22.612Z
Learnt from: CR
PR: tuannvm/mcp-trino#0
File: CLAUDE.md:0-0
Timestamp: 2025-09-01T14:40:22.612Z
Learning: Applies to internal/trino/client.go : Use context-based timeouts for query execution and handle deadline propagation
Applied to files:
internal/trino/client.go
🧬 Code graph analysis (1)
internal/config/config_test.go (1)
internal/config/config.go (1)
NewTrinoConfig(43-132)
🔇 Additional comments (3)
internal/config/config.go (2)
38-40: Field addition looks good and is backward-compatible.Public
TrinoSourceon the config struct is the right place; no default keeps behavior unchanged when unset.
122-122: Plumbing of TrinoSource into the returned config is correct.internal/config/config_test.go (1)
54-84: Defaults test reads well and covers the new field.
| func TestTrinoSourceHeader(t *testing.T) { | ||
| tests := []struct { | ||
| name string | ||
| trinoSource string | ||
| expectHeader bool | ||
| }{ | ||
| { | ||
| name: "TrinoSource configured", | ||
| trinoSource: "test-application", | ||
| expectHeader: true, | ||
| }, | ||
| { | ||
| name: "TrinoSource empty", | ||
| trinoSource: "", | ||
| expectHeader: false, | ||
| }, | ||
| } | ||
|
|
||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| // This test verifies that our logic for conditionally sending | ||
| // the X-Trino-Source header is correct | ||
| if tt.trinoSource != "" && !tt.expectHeader { | ||
| t.Error("Logic error: non-empty TrinoSource should expect header") | ||
| } | ||
| if tt.trinoSource == "" && tt.expectHeader { | ||
| t.Error("Logic error: empty TrinoSource should not expect header") | ||
| } | ||
| }) | ||
| } | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Test doesn’t assert the actual header behavior; it only checks table logic.
This won’t catch regressions. Prefer asserting the DSN includes source= or refactor DSN construction into a helper for direct testing.
Example approach:
// in client.go (non-test): factor out
func buildDSN(cfg *config.TrinoConfig) string { /* compose DSN incl. source if set */ }
// in test:
cfg := &config.TrinoConfig{ /* ... TrinoSource: "test-application" ... */ }
dsn := trino.BuildDSN(cfg)
if !strings.Contains(dsn, "source=test-application") { t.Fatalf("missing source in DSN: %s", dsn) }Alternatively, spin up a test HTTP server and point the driver at it to assert X-Trino-Source is present, but the DSN test is simpler and sufficient here.
🌐 Web query:
In the Go driver github.com/trinodb/trino-go-client, which DSN parameter maps to the X-Trino-Source header?
💡 Result:
The DSN parameter is source — it sets the X-Trino-Source header. [1][2]
Sources:
[1] github.com/trinodb/trino-go-client README (DSN parameters: source).
[2] Trino client protocol docs (X-Trino-Source header description).
Extract DSN construction into a helper and test the source DSN parameter
Current test only validates boolean logic and won’t catch regressions in header inclusion. Add a buildDSN(cfg *config.TrinoConfig) function (e.g. in client.go) and write tests asserting the DSN string contains source=test-application when TrinoSource is set and omits it when empty.
| // Execute the query with X-Trino-Source header if configured | ||
| var rows *sql.Rows | ||
| var err error | ||
| if c.config.TrinoSource != "" { | ||
| rows, err = c.db.QueryContext(ctx, query, sql.Named("X-Trino-Source", c.config.TrinoSource)) | ||
| } else { | ||
| rows, err = c.db.QueryContext(ctx, query) | ||
| } | ||
|
|
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Incorrect way to set X-Trino-Source; this passes a query argument, not an HTTP header.
sql.Named("X-Trino-Source", ...) is treated as a bound parameter for placeholders, and will not set a header. It may also cause “unexpected arguments” errors for queries without placeholders.
Apply this diff to revert the per-query arg injection:
-// Execute the query with X-Trino-Source header if configured
-var rows *sql.Rows
-var err error
-if c.config.TrinoSource != "" {
- rows, err = c.db.QueryContext(ctx, query, sql.Named("X-Trino-Source", c.config.TrinoSource))
-} else {
- rows, err = c.db.QueryContext(ctx, query)
-}
+rows, err := c.db.QueryContext(ctx, query)Then set the Source once at connection time via DSN so the driver emits X-Trino-Source automatically:
// after constructing dsn in NewClient:
if cfg.TrinoSource != "" {
dsn = dsn + "&source=" + url.QueryEscape(cfg.TrinoSource)
}This aligns with the Trino Go driver’s config model where source maps to X-Trino-Source.
🌐 Web query:
Does the Trino Go driver (github.com/trinodb/trino-go-client) map the DSN query param "source" to the X-Trino-Source request header?
💡 Result:
Yes. The trino-go-client DSN "source" parameter is the client's source and is sent to the server as the X-Trino-Source request header. [1][2][3]
Sources:
- trino-go-client README (documents the DSN "source" parameter). [1]
- Trino client protocol docs (defines X-Trino-Source header). [2]
- trino-go-client GitHub issue referencing how DSN params (e.g., session_properties) are mapped into X-Trino-* headers in trino.go. [3]
Use DSN source param instead of per-query bound parameter to set X-Trino-Source header
sql.Named("X-Trino-Source", …) is treated as a SQL argument, not an HTTP header, and can trigger “unexpected arguments” errors. The Trino Go driver maps the DSN’s source query parameter to the X-Trino-Source header, so remove the per-query injection and configure source on the DSN in NewClient:
-// Execute the query with X-Trino-Source header if configured
-var rows *sql.Rows
-var err error
-if c.config.TrinoSource != "" {
- rows, err = c.db.QueryContext(ctx, query, sql.Named("X-Trino-Source", c.config.TrinoSource))
-} else {
- rows, err = c.db.QueryContext(ctx, query)
-}
+rows, err := c.db.QueryContext(ctx, query)Then, after constructing the DSN in NewClient:
if cfg.TrinoSource != "" {
dsn += "&source=" + url.QueryEscape(cfg.TrinoSource)
}🤖 Prompt for AI Agents
In internal/trino/client.go around lines 140 to 148, the code passes
sql.Named("X-Trino-Source", c.config.TrinoSource) as a query argument which the
driver treats as a SQL parameter (causing unexpected-argument errors); remove
the per-query injection and instead set the Trino X-Trino-Source header via the
DSN `source` query parameter when creating the client. Update NewClient where
you build the DSN to append `&source=` + url.QueryEscape(cfg.TrinoSource) when
cfg.TrinoSource is non-empty, and then delete the sql.Named(...) usage in the
query execution path so calls use c.db.QueryContext(ctx, query) only.
- Add TRINO_SOURCE environment variable configuration - Support sending X-Trino-Source header in all Trino queries - Helps with audit trails and query tracking in enterprise environments - Header is sent only when TRINO_SOURCE is configured (optional) Signed-off-by: Joao Leite <[email protected]>
- Add tests for TrinoConfig with TRINO_SOURCE environment variable - Add tests for X-Trino-Source header logic - Verify default values and environment variable handling - All tests pass, maintaining backward compatibility Signed-off-by: Joao Leite <[email protected]>
joaoleite
force-pushed
the
feature/add-trino-source-header
branch
from
0e6d112 to
6051798
Compare
1 participant
Summary
Add support for configuring and sending the
X-Trino-Sourceheader in Trino queries through theTRINO_SOURCEenvironment variable.Changes
TrinoSourcefield toTrinoConfigstructTRINO_SOURCEenvironment variable support (optional, no default value)ExecuteQueryto sendX-Trino-Sourceheader when configuredTRINO_SOURCEis explicitly setUse Case
Some Trino installations require the
X-Trino-Sourceheader for audit trails and query tracking in enterprise environments. This change enables users to configure this header without code modifications.Testing
Test Coverage
TRINO_SOURCEenvironment variableExample Usage