Issues #22, #23, #24 - QueryStats Improvements (#26)

* Add query ID and warnings to `QueryResult`, refactor `QueryStats` duration handling

- Added `QueryID` and `Warnings` fields to `QueryResult`.
- Refactored `QueryStats` to use `DurationMs` (in milliseconds) instead of `time.Duration`.
- Updated tests to verify new fields and JSON serialization.
- Introduced `queryProgressUpdater` to capture query ID during execution.
- Added fallback logic to handle unsupported progress callbacks.

* Update `QueryStats` tests to verify `QueryID` addition and `DurationMs` refactor

* Refactor tests: replace `time.Duration` with `DurationMs` in `QueryStats`

* Update `QueryResult` to use `DurationMs` instead of `Milliseconds` for execution time formatting

* Refactor tests: replace `Duration` with `DurationMs` in `QueryStats` usage

* Remove Warnings field from QueryResult and related tests

- Removed the `Warnings` field from `QueryResult` as it is not exposed in trino-go-client v0.333.0.
- Updated tests to reflect this change.
- Added a comment noting the Trino REST API behavior and linked to the relevant GitHub issue.

* Add tests for `queryProgressUpdater` including concurrent access handling

- Introduced unit tests for `queryProgressUpdater` to verify correct handling of query ID updates.
- Added a concurrency test to ensure thread safety during concurrent reads and writes.

* Add codecov configuration to enforce coverage targets and enable PR comments

- Introduced `codecov.yml` with project and patch coverage targets (82% and 80% respectively).
- Configured Codecov to provide PR comments for coverage changes.

Author: cjimti

codecov.yml

@@ -0,0 +1,15 @@
+coverage:
+  status:
+    project:
+      default:
+        target: 82%
+        threshold: 2%
+    patch:
+      default:
+        target: 80%
+        threshold: 5%
+
+comment:
+  layout: "reach,diff,flags,files"
+  behavior: default
+  require_changes: true

pkg/client/client.go

@@ -6,9 +6,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
-	_ "github.com/trinodb/trino-go-client/trino" // Trino database driver
+	"github.com/trinodb/trino-go-client/trino"
 )
 
 // Client is a wrapper around the Trino database connection.
@@ -68,6 +69,8 @@ type QueryResult struct {
 	Columns []ColumnInfo     `json:"columns"`
 	Rows    []map[string]any `json:"rows"`
 	Stats   QueryStats       `json:"stats"`
+	// Note: Trino REST API returns warnings, but trino-go-client v0.333.0
+	// does not expose them. See: https://github.com/trinodb/trino-go-client/issues
 }
 
 // ColumnInfo describes a column in the result set.
@@ -79,10 +82,11 @@ type ColumnInfo struct {
 
 // QueryStats contains execution statistics.
 type QueryStats struct {
-	RowCount     int           `json:"row_count"`
-	Duration     time.Duration `json:"duration_ms"`
-	Truncated    bool          `json:"truncated"`
-	LimitApplied int           `json:"limit_applied,omitempty"`
+	RowCount     int    `json:"row_count"`
+	DurationMs   int64  `json:"duration_ms"`
+	Truncated    bool   `json:"truncated"`
+	LimitApplied int    `json:"limit_applied,omitempty"`
+	QueryID      string `json:"query_id,omitempty"`
 }
 
 // QueryOptions configures query execution.
@@ -107,6 +111,28 @@ func DefaultQueryOptions() QueryOptions {
 	}
 }
 
+// queryProgressUpdater captures the query ID from Trino progress updates.
+type queryProgressUpdater struct {
+	mu      sync.Mutex
+	queryID string
+}
+
+// Update implements the trino.ProgressUpdater interface.
+func (p *queryProgressUpdater) Update(info trino.QueryProgressInfo) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if info.QueryId != "" {
+		p.queryID = info.QueryId
+	}
+}
+
+// QueryID returns the captured query ID.
+func (p *queryProgressUpdater) QueryID() string {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.queryID
+}
+
 // Query executes a SQL query and returns the results.
 func (c *Client) Query(ctx context.Context, sqlQuery string, opts QueryOptions) (*QueryResult, error) {
 	start := time.Now()
@@ -125,10 +151,28 @@ func (c *Client) Query(ctx context.Context, sqlQuery string, opts QueryOptions)
 		limit = 1000
 	}
 
-	// Execute query
-	rows, err := c.db.QueryContext(ctx, sqlQuery)
+	// Set up progress updater to capture query ID
+	progressUpdater := &queryProgressUpdater{}
+
+	// Execute query with progress callback to capture query ID.
+	// The progress callback is a Trino-specific feature. If the driver doesn't
+	// support it (e.g., when using sqlmock for testing), fall back to a simple query.
+	rows, err := c.db.QueryContext(ctx, sqlQuery,
+		sql.Named("X-Trino-Progress-Callback", trino.ProgressUpdater(progressUpdater)),
+		sql.Named("X-Trino-Progress-Callback-Period", 100*time.Millisecond),
+	)
 	if err != nil {
-		return nil, fmt.Errorf("query failed: %w", err)
+		// Check if the error is due to unsupported argument type (e.g., when using sqlmock).
+		// In that case, retry without the progress callback.
+		if strings.Contains(err.Error(), "unsupported type") {
+			rows, err = c.db.QueryContext(ctx, sqlQuery)
+			if err != nil {
+				return nil, fmt.Errorf("query failed: %w", err)
+			}
+			progressUpdater = nil // Clear so we don't try to get QueryID
+		} else {
+			return nil, fmt.Errorf("query failed: %w", err)
+		}
 	}
 	defer func() { _ = rows.Close() }()
 
@@ -187,12 +231,16 @@ func (c *Client) Query(ctx context.Context, sqlQuery string, opts QueryOptions)
 		return nil, fmt.Errorf("row iteration error: %w", err)
 	}
 
-	result.Stats = QueryStats{
+	stats := QueryStats{
 		RowCount:     rowCount,
-		Duration:     time.Since(start),
+		DurationMs:   time.Since(start).Milliseconds(),
 		Truncated:    truncated,
 		LimitApplied: limit,
 	}
+	if progressUpdater != nil {
+		stats.QueryID = progressUpdater.QueryID()
+	}
+	result.Stats = stats
 
 	return result, nil
 }

pkg/client/client_test.go

@@ -3,8 +3,11 @@ package client
 import (
 	"bytes"
 	"encoding/json"
+	"sync"
 	"testing"
 	"time"
+
+	"github.com/trinodb/trino-go-client/trino"
 )
 
 func TestNew_InvalidConfig(t *testing.T) {
@@ -202,9 +205,10 @@ func TestQueryResult_JSON(t *testing.T) {
 		},
 		Stats: QueryStats{
 			RowCount:     2,
-			Duration:     100 * time.Millisecond,
+			DurationMs:   100,
 			Truncated:    false,
 			LimitApplied: 1000,
+			QueryID:      "20240115_123456_00001_abcde",
 		},
 	}
 
@@ -227,6 +231,9 @@ func TestQueryResult_JSON(t *testing.T) {
 	if decoded.Stats.RowCount != 2 {
 		t.Errorf("expected RowCount 2, got %d", decoded.Stats.RowCount)
 	}
+	if decoded.Stats.QueryID != "20240115_123456_00001_abcde" {
+		t.Errorf("expected QueryID '20240115_123456_00001_abcde', got %q", decoded.Stats.QueryID)
+	}
 }
 
 func TestColumnInfo_JSON(t *testing.T) {
@@ -330,16 +337,23 @@ func TestColumnDef_JSON(t *testing.T) {
 func TestQueryStats_JSON(t *testing.T) {
 	stats := QueryStats{
 		RowCount:     100,
-		Duration:     500 * time.Millisecond,
+		DurationMs:   500,
 		Truncated:    true,
 		LimitApplied: 100,
+		QueryID:      "test_query_id",
 	}
 
 	data, err := json.Marshal(stats)
 	if err != nil {
 		t.Fatalf("failed to marshal QueryStats: %v", err)
 	}
 
+	// Verify the JSON contains the correct field name and value
+	expectedJSON := `{"row_count":100,"duration_ms":500,"truncated":true,"limit_applied":100,"query_id":"test_query_id"}`
+	if string(data) != expectedJSON {
+		t.Errorf("expected JSON %s, got %s", expectedJSON, string(data))
+	}
+
 	var decoded QueryStats
 	if err := json.Unmarshal(data, &decoded); err != nil {
 		t.Fatalf("failed to unmarshal QueryStats: %v", err)
@@ -348,9 +362,15 @@ func TestQueryStats_JSON(t *testing.T) {
 	if decoded.RowCount != 100 {
 		t.Errorf("expected RowCount 100, got %d", decoded.RowCount)
 	}
+	if decoded.DurationMs != 500 {
+		t.Errorf("expected DurationMs 500, got %d", decoded.DurationMs)
+	}
 	if !decoded.Truncated {
 		t.Error("expected Truncated to be true")
 	}
+	if decoded.QueryID != "test_query_id" {
+		t.Errorf("expected QueryID 'test_query_id', got %q", decoded.QueryID)
+	}
 }
 
 func TestQueryOptions_Defaults(t *testing.T) {
@@ -435,3 +455,81 @@ func TestQuoteIdentifier(t *testing.T) {
 		})
 	}
 }
+
+func TestQueryProgressUpdater_Update(t *testing.T) {
+	tests := []struct {
+		name            string
+		updates         []trino.QueryProgressInfo
+		expectedQueryID string
+	}{
+		{
+			name:            "single update with query ID",
+			updates:         []trino.QueryProgressInfo{{QueryId: "20240115_123456_00001_abcde"}},
+			expectedQueryID: "20240115_123456_00001_abcde",
+		},
+		{
+			name:            "empty query ID ignored",
+			updates:         []trino.QueryProgressInfo{{QueryId: ""}},
+			expectedQueryID: "",
+		},
+		{
+			name: "multiple updates keeps last non-empty",
+			updates: []trino.QueryProgressInfo{
+				{QueryId: "first_query_id"},
+				{QueryId: "second_query_id"},
+			},
+			expectedQueryID: "second_query_id",
+		},
+		{
+			name: "empty update after valid ID keeps valid",
+			updates: []trino.QueryProgressInfo{
+				{QueryId: "valid_query_id"},
+				{QueryId: ""},
+			},
+			expectedQueryID: "valid_query_id",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			updater := &queryProgressUpdater{}
+
+			for _, update := range tt.updates {
+				updater.Update(update)
+			}
+
+			if got := updater.QueryID(); got != tt.expectedQueryID {
+				t.Errorf("QueryID() = %q, want %q", got, tt.expectedQueryID)
+			}
+		})
+	}
+}
+
+func TestQueryProgressUpdater_ConcurrentAccess(t *testing.T) {
+	updater := &queryProgressUpdater{}
+	var wg sync.WaitGroup
+
+	// Spawn multiple goroutines to test thread safety
+	for i := 0; i < 100; i++ {
+		wg.Add(2)
+
+		// Writer goroutine
+		go func() {
+			defer wg.Done()
+			updater.Update(trino.QueryProgressInfo{QueryId: "query_id"})
+		}()
+
+		// Reader goroutine
+		go func() {
+			defer wg.Done()
+			_ = updater.QueryID()
+		}()
+	}
+
+	wg.Wait()
+
+	// Should have a valid query ID after all updates
+	if got := updater.QueryID(); got != "query_id" {
+		t.Errorf("QueryID() = %q, want %q", got, "query_id")
+	}
+}

pkg/client/integration_test.go

@@ -260,23 +260,27 @@ func TestColumnDef_Fields(t *testing.T) {
 func TestQueryStats_Fields(t *testing.T) {
 	stats := QueryStats{
 		RowCount:     1000,
-		Duration:     500 * time.Millisecond,
+		DurationMs:   500,
 		Truncated:    true,
 		LimitApplied: 1000,
+		QueryID:      "test_query_123",
 	}
 
 	if stats.RowCount != 1000 {
 		t.Errorf("expected RowCount 1000, got %d", stats.RowCount)
 	}
-	if stats.Duration != 500*time.Millisecond {
-		t.Errorf("expected Duration 500ms, got %v", stats.Duration)
+	if stats.DurationMs != 500 {
+		t.Errorf("expected DurationMs 500, got %d", stats.DurationMs)
 	}
 	if !stats.Truncated {
 		t.Error("expected Truncated to be true")
 	}
 	if stats.LimitApplied != 1000 {
 		t.Errorf("expected LimitApplied 1000, got %d", stats.LimitApplied)
 	}
+	if stats.QueryID != "test_query_123" {
+		t.Errorf("expected QueryID 'test_query_123', got %q", stats.QueryID)
+	}
 }
 
 // TestColumnInfo_Fields tests the ColumnInfo struct.

pkg/tools/mock_client_test.go

@@ -2,7 +2,6 @@ package tools
 
 import (
 	"context"
-	"time"
 
 	"github.com/txn2/mcp-trino/pkg/client"
 )
@@ -59,7 +58,7 @@ func NewMockTrinoClient() *MockTrinoClient {
 				},
 				Stats: client.QueryStats{
 					RowCount:     2,
-					Duration:     100 * time.Millisecond,
+					DurationMs:   100,
 					Truncated:    false,
 					LimitApplied: 1000,
 				},

pkg/tools/query.go

@@ -165,7 +165,7 @@ func formatCSV(result *client.QueryResult) string {
 	if result.Stats.Truncated {
 		output += fmt.Sprintf(" (truncated at limit %d)", result.Stats.LimitApplied)
 	}
-	output += fmt.Sprintf(", executed in %dms", result.Stats.Duration.Milliseconds())
+	output += fmt.Sprintf(", executed in %dms", result.Stats.DurationMs)
 
 	return output
 }
@@ -209,7 +209,7 @@ func formatMarkdown(result *client.QueryResult) string {
 	if result.Stats.Truncated {
 		output += fmt.Sprintf(" (truncated at limit %d)", result.Stats.LimitApplied)
 	}
-	output += fmt.Sprintf(", executed in %dms*", result.Stats.Duration.Milliseconds())
+	output += fmt.Sprintf(", executed in %dms*", result.Stats.DurationMs)
 
 	return output
 }