Extensibility¶

mcp-trino provides three extension points for customizing behavior: middleware, interceptors, and transformers.

Overview¶

Extension Type	When	Purpose
Middleware	Before/after tool execution	Authentication, logging, rate limiting
Interceptors	Before SQL execution	Query validation, rewriting, audit
Transformers	After query execution	Result modification, redaction, enrichment

flowchart LR
    REQ["Request"] --> MW["Middleware"]
    MW --> INT["Interceptors"]
    INT --> SQL["SQL Execution"]
    SQL --> TF["Transformers"]
    TF --> MW2["Middleware"]
    MW2 --> RES["Response"]

Middleware¶

Middleware intercepts tool calls before and after execution.

Interface¶

type ToolMiddleware interface {
    Before(ctx *ToolContext) error
    After(ctx *ToolContext, result *mcp.CallToolResult, err error)
}

ToolContext¶

The context carries request information through the chain:

type ToolContext struct {
    ToolName   string
    Arguments  map[string]interface{}
    StartTime  time.Time
    RequestID  string
    values     sync.Map
}

// Store and retrieve custom values
ctx.Set("user_id", userID)
userID, ok := ctx.Get("user_id")

Example: Authentication¶

type AuthMiddleware struct {
    validator TokenValidator
}

func (m *AuthMiddleware) Before(ctx *ToolContext) error {
    token, ok := ctx.Get("auth_token")
    if !ok {
        return errors.New("authentication required")
    }

    user, err := m.validator.Validate(token.(string))
    if err != nil {
        return fmt.Errorf("invalid token: %w", err)
    }

    ctx.Set("user", user)
    return nil
}

func (m *AuthMiddleware) After(ctx *ToolContext, result *mcp.CallToolResult, err error) {
    user, _ := ctx.Get("user")
    log.Printf("User %s accessed %s", user, ctx.ToolName)
}

Example: Rate Limiting¶

type RateLimitMiddleware struct {
    limiter *rate.Limiter
}

func NewRateLimitMiddleware(rps float64, burst int) *RateLimitMiddleware {
    return &RateLimitMiddleware{
        limiter: rate.NewLimiter(rate.Limit(rps), burst),
    }
}

func (m *RateLimitMiddleware) Before(ctx *ToolContext) error {
    if !m.limiter.Allow() {
        return errors.New("rate limit exceeded")
    }
    return nil
}

func (m *RateLimitMiddleware) After(ctx *ToolContext, result *mcp.CallToolResult, err error) {}

Example: Request Logging¶

type RequestLogMiddleware struct {
    output io.Writer
}

func (m *RequestLogMiddleware) Before(ctx *ToolContext) error {
    json.NewEncoder(m.output).Encode(map[string]interface{}{
        "event":      "request_start",
        "request_id": ctx.RequestID,
        "tool":       ctx.ToolName,
        "timestamp":  ctx.StartTime.Format(time.RFC3339),
    })
    return nil
}

func (m *RequestLogMiddleware) After(ctx *ToolContext, result *mcp.CallToolResult, err error) {
    status := "success"
    if err != nil {
        status = "error"
    }
    json.NewEncoder(m.output).Encode(map[string]interface{}{
        "event":       "request_end",
        "request_id":  ctx.RequestID,
        "tool":        ctx.ToolName,
        "status":      status,
        "duration_ms": time.Since(ctx.StartTime).Milliseconds(),
    })
}

Middleware Chain¶

Middleware executes in order added. After() runs in reverse:

toolkit.Use(loggingMiddleware)    // 1st Before, Last After
toolkit.Use(authMiddleware)       // 2nd Before, 2nd-to-last After
toolkit.Use(rateLimitMiddleware)  // 3rd Before, 3rd-to-last After

If Before() returns an error, the chain stops and the tool doesn't execute.

Interceptors¶

Interceptors transform SQL before execution.

Interface¶

type QueryInterceptor interface {
    Intercept(ctx context.Context, sql string) (string, error)
}

Example: SQL Validation¶

type SQLValidatorInterceptor struct {
    blockedPatterns []*regexp.Regexp
}

func NewSQLValidatorInterceptor(patterns []string) *SQLValidatorInterceptor {
    compiled := make([]*regexp.Regexp, len(patterns))
    for i, p := range patterns {
        compiled[i] = regexp.MustCompile(p)
    }
    return &SQLValidatorInterceptor{blockedPatterns: compiled}
}

func (i *SQLValidatorInterceptor) Intercept(ctx context.Context, sql string) (string, error) {
    normalized := strings.ToUpper(sql)
    for _, pattern := range i.blockedPatterns {
        if pattern.MatchString(normalized) {
            return "", fmt.Errorf("SQL pattern blocked by security policy")
        }
    }
    return sql, nil
}

Example: Tenant Filter¶

type TenantFilterInterceptor struct {
    column string
}

func (i *TenantFilterInterceptor) Intercept(ctx context.Context, sql string) (string, error) {
    tenant := ctx.Value("tenant")
    if tenant == nil {
        return "", errors.New("tenant not found in context")
    }

    // Wrap query with tenant filter
    return fmt.Sprintf(
        "SELECT * FROM (%s) WHERE %s = '%s'",
        sql, i.column, tenant,
    ), nil
}

Example: Audit Logging¶

type AuditLogInterceptor struct {
    output io.Writer
}

func (i *AuditLogInterceptor) Intercept(ctx context.Context, sql string) (string, error) {
    json.NewEncoder(i.output).Encode(map[string]interface{}{
        "timestamp": time.Now().Format(time.RFC3339),
        "event":     "query_executed",
        "sql":       sql,
        "user":      ctx.Value("user"),
    })
    return sql, nil // Pass through unchanged
}

Example: Schema Rewriting¶

type SchemaRewriteInterceptor struct {
    from, to string
}

func (i *SchemaRewriteInterceptor) Intercept(ctx context.Context, sql string) (string, error) {
    pattern := regexp.MustCompile(`\b` + i.from + `\.`)
    return pattern.ReplaceAllString(sql, i.to+"."), nil
}

Interceptor Chain¶

Each interceptor receives the output of the previous one:

toolkit.AddInterceptor(auditInterceptor)      // 1st - logs original
toolkit.AddInterceptor(validatorInterceptor)  // 2nd - validates
toolkit.AddInterceptor(rewriteInterceptor)    // 3rd - rewrites

Return an error to block execution.

Transformers¶

Transformers modify query results before returning to the client.

Interface¶

type ResultTransformer interface {
    Transform(ctx context.Context, result *QueryResult) (*QueryResult, error)
}

type QueryResult struct {
    Columns  []string
    Rows     [][]interface{}
    RowCount int
    Metadata map[string]interface{}
}

Example: Data Redaction¶

type RedactionTransformer struct {
    sensitiveColumns map[string]bool
    redactedValue    string
}

func NewRedactionTransformer(columns []string) *RedactionTransformer {
    sensitive := make(map[string]bool)
    for _, col := range columns {
        sensitive[strings.ToLower(col)] = true
    }
    return &RedactionTransformer{
        sensitiveColumns: sensitive,
        redactedValue:    "***REDACTED***",
    }
}

func (t *RedactionTransformer) Transform(ctx context.Context, result *QueryResult) (*QueryResult, error) {
    // Find sensitive column indices
    sensitiveIndices := make(map[int]bool)
    for i, col := range result.Columns {
        if t.sensitiveColumns[strings.ToLower(col)] {
            sensitiveIndices[i] = true
        }
    }

    // Redact values
    for _, row := range result.Rows {
        for i := range row {
            if sensitiveIndices[i] {
                row[i] = t.redactedValue
            }
        }
    }

    return result, nil
}

Example: Metadata Enrichment¶

type MetadataTransformer struct{}

func (t *MetadataTransformer) Transform(ctx context.Context, result *QueryResult) (*QueryResult, error) {
    if result.Metadata == nil {
        result.Metadata = make(map[string]interface{})
    }

    result.Metadata["row_count"] = result.RowCount
    result.Metadata["column_count"] = len(result.Columns)
    result.Metadata["timestamp"] = time.Now().Format(time.RFC3339)

    return result, nil
}

Example: Date Formatting¶

type DateFormatTransformer struct {
    outputFormat string
    dateColumns  map[string]bool
}

func (t *DateFormatTransformer) Transform(ctx context.Context, result *QueryResult) (*QueryResult, error) {
    dateIndices := make(map[int]bool)
    for i, col := range result.Columns {
        if t.dateColumns[strings.ToLower(col)] {
            dateIndices[i] = true
        }
    }

    for _, row := range result.Rows {
        for i := range row {
            if dateIndices[i] {
                if ts, ok := row[i].(time.Time); ok {
                    row[i] = ts.Format(t.outputFormat)
                }
            }
        }
    }

    return result, nil
}

Transformer Chain¶

Each transformer receives the output of the previous one:

toolkit.AddTransformer(redactionTransformer)   // 1st - redact sensitive
toolkit.AddTransformer(dateFormatTransformer)  // 2nd - format dates
toolkit.AddTransformer(metadataTransformer)    // 3rd - add metadata

Semantic Providers¶

Add organizational context to tool output by integrating with metadata catalogs.

import (
    "github.com/txn2/mcp-trino/pkg/semantic"
    "github.com/txn2/mcp-trino/pkg/semantic/providers/datahub"
    "github.com/txn2/mcp-trino/pkg/tools"
)

// Create DataHub provider
provider, _ := datahub.New(datahub.FromEnv())
defer provider.Close()

// Add to toolkit with caching
toolkit := tools.NewToolkit(trinoClient, cfg,
    tools.WithSemanticProvider(provider),
    tools.WithSemanticCache(semantic.DefaultCacheConfig()),
)

When configured, trino_describe_table enriches output with:

Table and column descriptions
Ownership information
Tags and domain classifications
Sensitivity markers (PII, sensitive data)

See the Semantic Layer documentation for provider setup, caching, and building custom providers.

Tool Descriptions¶

Override the description that AI agents see for each tool. Useful for adding domain-specific context so the agent understands what data is behind each tool.

Toolkit-Level Overrides¶

Set descriptions for multiple tools at construction time with WithDescriptions:

toolkit := tools.NewToolkit(trinoClient, cfg,
    tools.WithDescriptions(map[tools.ToolName]string{
        tools.ToolQuery:         "Query the retail analytics warehouse",
        tools.ToolDescribeTable: "Describe tables in the retail warehouse",
    }),
)
toolkit.RegisterAll(server)

Per-Registration Overrides¶

Set a description for a single tool at registration time with WithDescription via RegisterWith:

toolkit.RegisterWith(server, tools.ToolQuery,
    tools.WithDescription("Query the retail analytics warehouse"),
)

Priority Chain¶

When multiple sources provide a description, the highest-priority source wins:

per-registration (WithDescription)
        ↓ fallback
toolkit-level (WithDescriptions)
        ↓ fallback
file config (toolkit.descriptions)
        ↓ fallback
built-in default

Use tools.DefaultDescription(tools.ToolQuery) to read the built-in default for any tool.

Tool Annotations¶

Declare behavioral hints on tools so that AI agents understand side effects without executing them. Annotations follow the MCP specification and include readOnlyHint, destructiveHint, idempotentHint, and openWorldHint.

All built-in tools ship with sensible default annotations. Schema-browsing tools are marked read-only and idempotent; trino_query is marked non-destructive but not read-only (since the SQL could be anything).

Default Annotations¶

Tool	ReadOnly	Destructive	Idempotent	OpenWorld
`trino_query`	false	false	false	false
`trino_explain`	true	—	true	false
`trino_browse`	true	—	true	false
`trino_describe_table`	true	—	true	false
`trino_list_connections`	true	—	true	false

Read the built-in default for any tool:

ann := tools.DefaultAnnotations(tools.ToolQuery)
// ann.DestructiveHint => *false

Toolkit-Level Overrides¶

Set annotations for multiple tools at construction time with WithAnnotations:

toolkit := tools.NewToolkit(trinoClient, cfg,
    tools.WithAnnotations(map[tools.ToolName]*mcp.ToolAnnotations{
        tools.ToolQuery: {
            ReadOnlyHint:   true,  // e.g., if a read-only interceptor is active
            IdempotentHint: true,
        },
    }),
)
toolkit.RegisterAll(server)

Per-Registration Overrides¶

Set annotations for a single tool at registration time with WithAnnotation via RegisterWith:

toolkit.RegisterWith(server, tools.ToolQuery,
    tools.WithAnnotation(&mcp.ToolAnnotations{
        ReadOnlyHint:   true,
        IdempotentHint: true,
    }),
)

Priority Chain¶

When multiple sources provide annotations, the highest-priority source wins:

per-registration (WithAnnotation)
        ↓ fallback
toolkit-level (WithAnnotations)
        ↓ fallback
built-in default

Structured Outputs¶

All tool handlers return typed output structs alongside the human-readable TextContent result. This enables downstream consumers (middleware, platform wrappers, programmatic clients) to access structured data without parsing text.

Output Types¶

Tool	Output Type	Key Fields
`trino_query`	`QueryOutput`	`columns`, `rows`, `row_count`, `stats`
`trino_explain`	`ExplainOutput`	`plan`, `type`
`trino_browse`	`BrowseOutput`	`level`, `catalog`, `schema`, `items`, `count`, `pattern`
`trino_describe_table`	`DescribeTableOutput`	`catalog`, `schema`, `table`, `columns`, `column_count`
`trino_list_connections`	`ListConnectionsOutput`	`connections`, `count`

Accessing Structured Output¶

Tool handlers return (*mcp.CallToolResult, *OutputType, error). The second return value carries the structured output:

// In middleware or a wrapping handler, the typed output is available
// as the second return value from the handler function.
//
// Example: QueryOutput
// {
//   "columns": [{"name": "id", "type": "bigint"}, ...],
//   "rows": [{"id": 1, "name": "alice"}, ...],
//   "row_count": 2,
//   "stats": {"row_count": 2, "truncated": false, "duration_ms": 42}
// }

Advertised Output Schemas¶

Every tool advertises an explicit JSON Schema (2020-12) for its structured output. The schemas are declared by hand rather than inferred from the Go output structs, and they are deliberately open: no top-level required list and no additionalProperties: false.

That matters when mcp-trino is composed into a larger server. A host that adds keys to structuredContent after the handler returns — an error envelope, a call reference, a trace ID — produces a result that still validates against what the tool advertised, so a schema-validating client does not reject it. Properties backed by a Go slice admit null as well as an array, because a nil slice marshals to null.

Read the built-in default for any tool:

schema := tools.DefaultOutputSchema(tools.ToolQuery)
// map[string]any{"type": "object", "properties": {...}}

Toolkit-Level Overrides¶

Set output schemas for multiple tools at construction time with WithOutputSchemas:

toolkit := tools.NewToolkit(trinoClient, cfg,
    tools.WithOutputSchemas(map[tools.ToolName]any{
        tools.ToolQuery: map[string]any{
            "type": "object",
            "properties": map[string]any{
                "rows":       map[string]any{"type": []string{"array", "null"}},
                "trace_id":   map[string]any{"type": "string"},
            },
        },
    }),
)

Per-Registration Overrides¶

Set the output schema for a single tool at registration time with WithOutputSchema via RegisterWith:

toolkit.RegisterWith(server, tools.ToolQuery,
    tools.WithOutputSchema(map[string]any{"type": "object"}),
)

A schema may be any value that JSON-marshals to a JSON Schema object with "type": "object" — a map[string]any, a json.RawMessage, or a *jsonschema.Schema. The MCP SDK validates every result against the schema the tool advertises, so an override narrower than what the handler emits turns those calls into protocol errors that discard the result.

Error results are the trap here, not successful ones. When a handler returns an error result with no typed output, the SDK validates the zero value of the output struct — whose slice fields are all null. An override that declares "rows" as a bare "array" therefore breaks every failed call, deterministically, while successful calls keep working. Declare slice-backed properties as []string{"array", "null"}.

Priority Chain¶

per-registration (WithOutputSchema)
        ↓ fallback
toolkit-level (WithOutputSchemas)
        ↓ fallback
built-in default (DefaultOutputSchema)

Built-in Extensions¶

mcp-trino includes ready-to-use extensions:

import "github.com/txn2/mcp-trino/pkg/extensions"

// Middleware
toolkit.Use(extensions.NewLoggingMiddleware(os.Stderr))
toolkit.Use(extensions.NewReadOnlyMiddleware())
toolkit.Use(extensions.NewMetricsMiddleware(collector))

// Interceptors
toolkit.AddInterceptor(extensions.NewQueryLogInterceptor(os.Stderr))
toolkit.AddInterceptor(extensions.NewReadOnlyInterceptor())

// Transformers
toolkit.AddTransformer(extensions.NewMetadataTransformer())
toolkit.AddTransformer(extensions.NewErrorHelpTransformer())

Complete Example¶

Production server with all extension types:

package main

import (
    "log"
    "os"

    "github.com/modelcontextprotocol/go-sdk/server"

    "github.com/txn2/mcp-trino/pkg/client"
    "github.com/txn2/mcp-trino/pkg/extensions"
    "github.com/txn2/mcp-trino/pkg/tools"
)

func main() {
    trinoClient, _ := client.New(client.FromEnv())
    defer trinoClient.Close()

    toolkit := tools.NewToolkit(trinoClient, tools.DefaultConfig())

    // Middleware: auth, logging, rate limiting
    toolkit.Use(&AuthMiddleware{validator: tokenValidator})
    toolkit.Use(extensions.NewLoggingMiddleware(os.Stderr))
    toolkit.Use(&RateLimitMiddleware{limiter: rateLimiter})
    toolkit.Use(extensions.NewReadOnlyMiddleware())

    // Interceptors: audit, validation, tenant filter
    toolkit.AddInterceptor(&AuditLogInterceptor{output: auditFile})
    toolkit.AddInterceptor(NewSQLValidatorInterceptor(blockedPatterns))
    toolkit.AddInterceptor(&TenantFilterInterceptor{column: "tenant_id"})

    // Transformers: redaction, metadata
    toolkit.AddTransformer(NewRedactionTransformer([]string{"ssn", "credit_card"}))
    toolkit.AddTransformer(extensions.NewMetadataTransformer())

    mcpServer := server.NewMCPServer("enterprise-server", "1.0.0")
    toolkit.RegisterAll(mcpServer)

    if err := server.ServeStdio(mcpServer); err != nil {
        log.Fatal(err)
    }
}

Testing Extensions¶

Testing Middleware¶

func TestAuthMiddleware(t *testing.T) {
    mw := &AuthMiddleware{validator: mockValidator}

    ctx := &ToolContext{}
    ctx.Set("auth_token", "valid-token")

    err := mw.Before(ctx)
    assert.NoError(t, err)

    user, ok := ctx.Get("user")
    assert.True(t, ok)
    assert.Equal(t, "test-user", user)
}

Testing Interceptors¶

func TestSQLValidator(t *testing.T) {
    interceptor := NewSQLValidatorInterceptor([]string{`DROP\s+TABLE`})

    tests := []struct {
        sql     string
        wantErr bool
    }{
        {"SELECT * FROM users", false},
        {"DROP TABLE users", true},
    }

    for _, tt := range tests {
        _, err := interceptor.Intercept(context.Background(), tt.sql)
        if tt.wantErr {
            assert.Error(t, err)
        } else {
            assert.NoError(t, err)
        }
    }
}

Testing Transformers¶

func TestRedactionTransformer(t *testing.T) {
    transformer := NewRedactionTransformer([]string{"password"})

    result := &QueryResult{
        Columns: []string{"id", "name", "password"},
        Rows: [][]interface{}{
            {1, "Alice", "secret123"},
        },
    }

    transformed, err := transformer.Transform(context.Background(), result)
    assert.NoError(t, err)
    assert.Equal(t, "***REDACTED***", transformed.Rows[0][2])
}