Propagate request context from tool handlers to database operations

[sgaunet]

Problem

MCP tool handlers in main.go receive mcp.CallToolRequest objects that contain a context, but this context is not propagated to database operations. Instead, all database operations use context.Background().

Current implementation:

// In main.go tool handler (around line 247)
func(request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    query := args["query"].(string)
    
    // Request context is ignored!
    result, err := appInstance.ExecuteQuery(query)
    // ...
}

// In client.go
func (c *PostgreSQLClientImpl) ExecuteQuery(query string) ([]map[string]any, error) {
    // Uses context.Background() instead of request context
    rows, err := c.db.QueryContext(context.Background(), query)
    // ...
}

Issues

1. No timeout control: Client can't set deadlines for operations
2. No cancellation: Can't cancel in-flight requests
3. Lost context values: Any context metadata from the request is ignored
4. Resource leaks: Long-running queries can't be interrupted

Impact Scenarios

Scenario 1: Long-Running Query

Current behavior:

• User sends query that takes 5 minutes
• No way to cancel it
• Resources held indefinitely

Desired behavior:

• Request context has deadline
• Query respects deadline
• Cancels automatically if client disconnects

Scenario 2: Client Disconnect

Current behavior:

• Client disconnects
• Query continues running on server
• Wastes database resources

Desired behavior:

• Context cancels when client disconnects
• Query is interrupted
• Resources released immediately

Proposed Solution

This issue is related to Issue #20 (Add context timeout support), but focuses on propagating the MCP request context rather than creating new contexts.

Step 1: Update Client Interface to Accept Context

// internal/app/interfaces.go
type PostgreSQLClient interface {
    Connect(ctx context.Context, connectionString string) error
    Ping(ctx context.Context) error
    Close() error
    ListDatabases(ctx context.Context) ([]string, error)
    ListSchemas(ctx context.Context) ([]string, error)
    ListTables(ctx context.Context, schema string) ([]*TableInfo, error)
    DescribeTable(ctx context.Context, schema, table string) ([]*ColumnInfo, error)
    ExecuteQuery(ctx context.Context, query string) ([]map[string]any, error)
    ListIndexes(ctx context.Context, schema, table string) ([]*IndexInfo, error)
    ExplainQuery(ctx context.Context, query string) (string, error)
    GetTableStats(ctx context.Context, schema, table string) (*TableStats, error)
}

Step 2: Update App Layer to Propagate Context

// internal/app/app.go
func (a *App) ExecuteQuery(ctx context.Context, query string) ([]map[string]any, error) {
    if err := a.ensureConnection(ctx); err != nil {
        return nil, fmt.Errorf("failed to execute query: %w", err)
    }
    
    result, err := a.client.ExecuteQuery(ctx, query)
    if err != nil {
        return nil, fmt.Errorf("failed to execute query: %w", err)
    }
    return result, nil
}

Step 3: Update Tool Handlers to Pass Request Context

// main.go
func setupExecuteQueryTool(appInstance *app.App) *mcp.Tool {
    return mcp.NewTool(
        "execute_query",
        "Execute a SQL query",
        inputSchema,
        func(request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
            query := args["query"].(string)
            
            // Pass request context to app layer
            result, err := appInstance.ExecuteQuery(request.Context(), query)
            if err != nil {
                return mcp.NewToolResultError(err.Error())
            }
            // ...
        },
    )
}

Step 4: Update Client Implementation

// internal/app/client.go
func (c *PostgreSQLClientImpl) ExecuteQuery(ctx context.Context, query string) ([]map[string]any, error) {
    if err := c.validateQuery(query); err != nil {
        return nil, err
    }
    
    // Use provided context instead of Background
    rows, err := c.db.QueryContext(ctx, query)
    if err != nil {
        return nil, fmt.Errorf("query execution failed: %w", err)
    }
    defer rows.Close()
    
    return c.processRows(rows)
}

Benefits

1. Timeout support: Requests can have deadlines
2. Cancellation: Operations can be cancelled
3. Resource efficiency: No wasted database connections
4. Observability: Context values can carry tracing information
5. Client control: Clients can set their own timeout policies

Testing Strategy

Unit Tests

func TestExecuteQueryWithTimeout(t *testing.T) {
    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()
    
    // Simulate slow query
    _, err := app.ExecuteQuery(ctx, "SELECT pg_sleep(10)")
    
    assert.Error(t, err)
    assert.True(t, errors.Is(err, context.DeadlineExceeded))
}

Integration Tests

func TestExecuteQueryWithCancellation(t *testing.T) {
    ctx, cancel := context.WithCancel(context.Background())
    
    // Start query in goroutine
    errCh := make(chan error)
    go func() {
        _, err := app.ExecuteQuery(ctx, "SELECT pg_sleep(10)")
        errCh <- err
    }()
    
    // Cancel after short delay
    time.Sleep(100 * time.Millisecond)
    cancel()
    
    err := <-errCh
    assert.Error(t, err)
    assert.True(t, errors.Is(err, context.Canceled))
}

Migration Checklist

This is a breaking change to internal APIs:

• Update PostgreSQLClient interface with context parameters
• Update PostgreSQLClientImpl to accept and use context
• Update all App methods to accept and propagate context
• Update all tool handlers to pass request.Context()
• Update ensureConnection to accept context
• Update all tests to pass context
• Add tests for context cancellation
• Add tests for context deadline
• Document context handling in README
• Consider adding default timeout if request context has none

Related Issues

• Issue Add context timeout support for database operations #20: Add context timeout support for database operations (foundational)
• Issue Add graceful shutdown handling for MCP server #38: Add graceful shutdown handling (uses context for cleanup)

Impact

• Severity: MEDIUM
• Type: Enhancement
• Location: main.go, internal/app/app.go, internal/app/client.go, internal/app/interfaces.go
• Benefits: Better resource management, cancellation support, timeout control

Implementation Order

1. First implement Issue Add context timeout support for database operations #20 (add context parameters)
2. Then implement this issue (propagate request context)
3. This ensures a clean migration path