Skip to content

Feature/add trino source header #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Feature/add trino source header #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a85813e
feat: add support for X-Trino-Source header
joaoleite Sep 4, 2025
6051798
test: add unit tests for TRINO_SOURCE configuration
joaoleite Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions internal/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ type TrinoConfig struct {
OIDCClientID string // OIDC client ID
OIDCClientSecret string // OIDC client secret
OAuthRedirectURI string // Fixed OAuth redirect URI (overrides dynamic callback)
// Custom Trino Source header
TrinoSource string // Value for X-Trino-Source header
}

// NewTrinoConfig creates a new TrinoConfig with values from environment variables or defaults
Expand Down Expand Up @@ -102,6 +104,9 @@ func NewTrinoConfig() (*TrinoConfig, error) {
}
}

// Get Trino Source from env/config (no default)
trinoSource := getEnv("TRINO_SOURCE", "")

return &TrinoConfig{
Host: getEnv("TRINO_HOST", "localhost"),
Port: port,
Expand All @@ -114,6 +119,7 @@ func NewTrinoConfig() (*TrinoConfig, error) {
SSLInsecure: sslInsecure,
AllowWriteQueries: allowWriteQueries,
QueryTimeout: queryTimeout,
TrinoSource: trinoSource,
OAuthEnabled: oauthEnabled,
OAuthProvider: oauthProvider,
JWTSecret: jwtSecret,
Expand Down
84 changes: 84 additions & 0 deletions internal/config/config_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
package config

import (
"os"
"testing"
)

func TestNewTrinoConfig_TrinoSource(t *testing.T) {
tests := []struct {
name string
envValue string
expectedSource string
}{
{
name: "TRINO_SOURCE set to custom value",
envValue: "dataeng-trino-api",
expectedSource: "dataeng-trino-api",
},
{
name: "TRINO_SOURCE set to empty string",
envValue: "",
expectedSource: "",
},
{
name: "TRINO_SOURCE not set",
envValue: "UNSET",
expectedSource: "",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// 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")
}

config, err := NewTrinoConfig()
if err != nil {
t.Fatalf("NewTrinoConfig() failed: %v", err)
}

if config.TrinoSource != tt.expectedSource {
t.Errorf("TrinoSource = %q, want %q", config.TrinoSource, tt.expectedSource)
}
})
}
}

func TestNewTrinoConfig_DefaultValues(t *testing.T) {
// Clean up environment
for _, env := range []string{"TRINO_HOST", "TRINO_PORT", "TRINO_USER", "TRINO_CATALOG", "TRINO_SCHEMA", "TRINO_SOURCE"} {
os.Unsetenv(env)
}

config, err := NewTrinoConfig()
if err != nil {
t.Fatalf("NewTrinoConfig() failed: %v", err)
}

// Check default values
if config.Host != "localhost" {
t.Errorf("Host = %q, want %q", config.Host, "localhost")
}
if config.Port != 8080 {
t.Errorf("Port = %d, want %d", config.Port, 8080)
}
if config.User != "trino" {
t.Errorf("User = %q, want %q", config.User, "trino")
}
if config.Catalog != "memory" {
t.Errorf("Catalog = %q, want %q", config.Catalog, "memory")
}
if config.Schema != "default" {
t.Errorf("Schema = %q, want %q", config.Schema, "default")
}
if config.TrinoSource != "" {
t.Errorf("TrinoSource = %q, want empty string", config.TrinoSource)
}
}
11 changes: 9 additions & 2 deletions internal/trino/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,8 +137,15 @@ func (c *Client) ExecuteQuery(query string) ([]map[string]interface{}, error) {
ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
defer cancel()

// Execute the query
rows, err := c.db.QueryContext(ctx, query)
// 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)
}

Comment on lines +140 to +148
Copy link

@coderabbitai coderabbitai bot Sep 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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.

if err != nil {
return nil, fmt.Errorf("query execution failed: %w", err)
}
Expand Down
32 changes: 32 additions & 0 deletions internal/trino/client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,3 +179,35 @@ func TestIsReadOnlyQuery(t *testing.T) {
})
}
}

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")
}
})
}
}
Comment on lines +183 to +213
Copy link

@coderabbitai coderabbitai bot Sep 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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.