From 0debef0fb2731d0f8110513cc8583481940d949c Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Mon, 29 Sep 2025 21:14:29 -0400 Subject: [PATCH 01/13] Initial basic implementation of 'tiger_db_execute_query' tool --- go.mod | 4 +- internal/tiger/mcp/db.go | 237 ++++++++++++++++++++++++++++ internal/tiger/mcp/server.go | 13 +- internal/tiger/mcp/service_tools.go | 40 ++--- 4 files changed, 267 insertions(+), 27 deletions(-) create mode 100644 internal/tiger/mcp/db.go diff --git a/go.mod b/go.mod index 1c073180..d6a4d887 100644 --- a/go.mod +++ b/go.mod @@ -4,12 +4,14 @@ go 1.25.1 require ( github.com/charmbracelet/bubbletea v1.3.6 + github.com/google/jsonschema-go v0.2.3 github.com/jackc/pgx/v5 v5.7.5 github.com/modelcontextprotocol/go-sdk v0.5.0 github.com/oapi-codegen/oapi-codegen/v2 v2.5.0 github.com/oapi-codegen/runtime v1.1.2 github.com/olekukonko/tablewriter v1.0.9 github.com/spf13/cobra v1.9.1 + github.com/spf13/pflag v1.0.6 github.com/spf13/viper v1.20.1 github.com/zalando/go-keyring v0.2.6 go.uber.org/mock v0.5.2 @@ -38,7 +40,6 @@ require ( github.com/go-openapi/swag v0.23.0 // indirect github.com/go-viper/mapstructure/v2 v2.2.1 // indirect github.com/godbus/dbus/v5 v5.1.0 // indirect - github.com/google/jsonschema-go v0.2.3 // indirect github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect @@ -67,7 +68,6 @@ require ( github.com/speakeasy-api/openapi-overlay v0.10.2 // indirect github.com/spf13/afero v1.12.0 // indirect github.com/spf13/cast v1.7.1 // indirect - github.com/spf13/pflag v1.0.6 // indirect github.com/subosito/gotenv v1.6.0 // indirect github.com/vmware-labs/yaml-jsonpath v0.3.2 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db.go new file mode 100644 index 00000000..237d39e0 --- /dev/null +++ b/internal/tiger/mcp/db.go @@ -0,0 +1,237 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/jackc/pgx/v5" + "github.com/modelcontextprotocol/go-sdk/mcp" + "go.uber.org/zap" + + "github.com/timescale/tiger-cli/internal/tiger/api" + "github.com/timescale/tiger-cli/internal/tiger/logging" + "github.com/timescale/tiger-cli/internal/tiger/util" +) + +// DBExecuteQueryInput represents input for tiger_db_execute_query +type DBExecuteQueryInput struct { + ServiceID string `json:"service_id,omitempty"` + Query string `json:"query"` + Timeout *int `json:"timeout,omitempty"` +} + +func (DBExecuteQueryInput) Schema() *jsonschema.Schema { + schema := util.Must(jsonschema.For[DBExecuteQueryInput](nil)) + + schema.Properties["service_id"].Description = "Service ID to execute query on (uses default if not provided)" + + schema.Properties["query"].Description = "SQL query to execute" + + schema.Properties["timeout"].Description = "Query timeout in seconds (default: 30)" + schema.Properties["timeout"].Minimum = util.Ptr(0.0) + schema.Properties["timeout"].Default = util.Must(json.Marshal(30)) + schema.Properties["timeout"].Examples = []any{10, 30, 60} + + return schema +} + +// DBExecuteQueryOutput represents output for tiger_db_execute_query +type DBExecuteQueryOutput struct { + Columns []string `json:"columns"` + Rows [][]any `json:"rows"` + RowCount int `json:"row_count"` + ExecutionTime string `json:"execution_time"` +} + +// registerDatabaseTools registers database operation tools with comprehensive schemas and descriptions +func (s *Server) registerDatabaseTools() { + // tiger_db_execute_query + mcp.AddTool(s.mcpServer, &mcp.Tool{ + Name: "tiger_db_execute_query", + Title: "Execute SQL Query", + Description: `Execute a SQL query against a service database. + +This tool connects to a database service and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Perfect for data exploration, schema inspection, and database operations. + +IMPORTANT: Use with caution - this tool can execute any SQL statement including INSERT, UPDATE, DELETE, and DDL commands. Always review queries before execution. + +Perfect for: +- Querying data from tables and views +- Inspecting database schema +- Testing database connectivity with real queries +- Performing data analysis and exploration`, + InputSchema: DBExecuteQueryInput{}.Schema(), + Annotations: &mcp.ToolAnnotations{ + DestructiveHint: util.Ptr(true), // Can execute destructive SQL + Title: "Execute SQL Query", + }, + }, s.handleDBExecuteQuery) +} + +// handleDBExecuteQuery handles the tiger_db_execute_query MCP tool +func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequest, input DBExecuteQueryInput) (*mcp.CallToolResult, DBExecuteQueryOutput, error) { + // Create fresh API client with current credentials + apiClient, err := s.createAPIClient() + if err != nil { + return nil, DBExecuteQueryOutput{}, err + } + + // Load fresh config and validate project ID is set + cfg, err := s.loadConfigWithProjectID() + if err != nil { + return nil, DBExecuteQueryOutput{}, err + } + + // Get service ID (use default from config if not provided) + serviceID := input.ServiceID + if serviceID == "" { + if cfg.ServiceID == "" { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("service ID is required. Please provide service_id or run 'tiger config set service_id '") + } + serviceID = cfg.ServiceID + } + + // Set default timeout if not provided + timeout := 30 * time.Second + if input.Timeout != nil { + timeout = time.Duration(*input.Timeout) * time.Second + } + + logging.Debug("MCP: Executing database query", + zap.String("project_id", cfg.ProjectID), + zap.String("service_id", serviceID), + zap.Duration("timeout", timeout), + ) + + // Get service details to construct connection string + serviceResp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, cfg.ProjectID, serviceID) + if err != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to get service details: %w", err) + } + + switch serviceResp.StatusCode() { + case 200: + if serviceResp.JSON200 == nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("empty response from API") + } + case 401: + return nil, DBExecuteQueryOutput{}, fmt.Errorf("authentication failed: invalid API key") + case 403: + return nil, DBExecuteQueryOutput{}, fmt.Errorf("permission denied: insufficient access to service") + case 404: + return nil, DBExecuteQueryOutput{}, fmt.Errorf("service '%s' not found in project '%s'", serviceID, cfg.ProjectID) + default: + return nil, DBExecuteQueryOutput{}, fmt.Errorf("API request failed with status %d", serviceResp.StatusCode()) + } + + service := *serviceResp.JSON200 + + // Build connection string with password (use direct connection, default role tsdbadmin) + connString, err := s.buildConnectionString(service, false, "tsdbadmin") + if err != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to build connection string: %w", err) + } + + // Create query context with timeout + queryCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + // Connect to database + conn, err := pgx.Connect(queryCtx, connString) + if err != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to connect to database: %w", err) + } + defer conn.Close(context.Background()) + + // Execute query and measure time + startTime := time.Now() + rows, err := conn.Query(queryCtx, input.Query) + if err != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("query execution failed: %w", err) + } + defer rows.Close() + + // Get column names from field descriptions + fieldDescriptions := rows.FieldDescriptions() + columns := make([]string, len(fieldDescriptions)) + for i, fd := range fieldDescriptions { + columns[i] = string(fd.Name) + } + + // Collect all rows + var resultRows [][]any + for rows.Next() { + // Scan values into generic interface slice + values, err := rows.Values() + if err != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to scan row: %w", err) + } + resultRows = append(resultRows, values) + } + + // Check for errors during iteration + if rows.Err() != nil { + return nil, DBExecuteQueryOutput{}, fmt.Errorf("error during row iteration: %w", rows.Err()) + } + + output := DBExecuteQueryOutput{ + Columns: columns, + Rows: resultRows, + RowCount: len(resultRows), + ExecutionTime: time.Since(startTime).String(), + } + + return nil, output, nil +} + +// buildConnectionString creates a PostgreSQL connection string from service details with password included +func (s *Server) buildConnectionString(service api.Service, pooled bool, role string) (string, error) { + if service.Endpoint == nil { + return "", fmt.Errorf("service endpoint not available") + } + + var endpoint *api.Endpoint + var host string + var port int + + // Use pooler endpoint if requested and available, otherwise use direct endpoint + if pooled && service.ConnectionPooler != nil && service.ConnectionPooler.Endpoint != nil { + endpoint = service.ConnectionPooler.Endpoint + } else { + // If pooled was requested but no pooler is available, use direct endpoint + endpoint = service.Endpoint + } + + if endpoint.Host == nil { + return "", fmt.Errorf("endpoint host not available") + } + host = *endpoint.Host + + if endpoint.Port != nil { + port = *endpoint.Port + } else { + port = 5432 // Default PostgreSQL port + } + + // Get password from storage + storage := util.GetPasswordStorage() + password, err := storage.Get(service) + if err != nil { + return "", fmt.Errorf("failed to retrieve password: %w", err) + } + + if password == "" { + return "", fmt.Errorf("no password available for service") + } + + // Database is always "tsdb" for TimescaleDB/PostgreSQL services + database := "tsdb" + + // Build connection string with password included + connectionString := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", role, password, host, port, database) + + return connectionString, nil +} diff --git a/internal/tiger/mcp/server.go b/internal/tiger/mcp/server.go index 0e04d0f7..4f04b33a 100644 --- a/internal/tiger/mcp/server.go +++ b/internal/tiger/mcp/server.go @@ -60,6 +60,9 @@ func (s *Server) registerTools(ctx context.Context) { // Service management tools s.registerServiceTools() + // Database operation tools + s.registerDatabaseTools() + // TODO: Register more tool groups // Register remote docs MCP server proxy @@ -85,18 +88,18 @@ func (s *Server) createAPIClient() (*api.ClientWithResponses, error) { return apiClient, nil } -// loadProjectID loads fresh config and returns the current project ID -func (s *Server) loadProjectID() (string, error) { +// loadConfigWithProjectID loads fresh config and validates that project ID is set +func (s *Server) loadConfigWithProjectID() (*config.Config, error) { // Load fresh config cfg, err := config.Load() if err != nil { - return "", fmt.Errorf("failed to load config: %w", err) + return nil, fmt.Errorf("failed to load config: %w", err) } if cfg.ProjectID == "" { - return "", fmt.Errorf("project ID is required. Please run 'tiger auth login' with --project-id") + return nil, fmt.Errorf("project ID is required. Please run 'tiger auth login'") } - return cfg.ProjectID, nil + return cfg, nil } // Close gracefully shuts down the MCP server and all proxy connections diff --git a/internal/tiger/mcp/service_tools.go b/internal/tiger/mcp/service_tools.go index e1b4121c..f5dff163 100644 --- a/internal/tiger/mcp/service_tools.go +++ b/internal/tiger/mcp/service_tools.go @@ -258,19 +258,19 @@ func (s *Server) handleServiceList(ctx context.Context, req *mcp.CallToolRequest return nil, ServiceListOutput{}, err } - // Load fresh project ID from current config - projectID, err := s.loadProjectID() + // Load fresh config and validate project ID is set + cfg, err := s.loadConfigWithProjectID() if err != nil { return nil, ServiceListOutput{}, err } - logging.Debug("MCP: Listing services", zap.String("project_id", projectID)) + logging.Debug("MCP: Listing services", zap.String("project_id", cfg.ProjectID)) // Make API call to list services ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := apiClient.GetProjectsProjectIdServicesWithResponse(ctx, projectID) + resp, err := apiClient.GetProjectsProjectIdServicesWithResponse(ctx, cfg.ProjectID) if err != nil { return nil, ServiceListOutput{}, fmt.Errorf("failed to list services: %w", err) } @@ -310,21 +310,21 @@ func (s *Server) handleServiceShow(ctx context.Context, req *mcp.CallToolRequest return nil, ServiceShowOutput{}, err } - // Load fresh project ID from current config - projectID, err := s.loadProjectID() + // Load fresh config and validate project ID is set + cfg, err := s.loadConfigWithProjectID() if err != nil { return nil, ServiceShowOutput{}, err } logging.Debug("MCP: Showing service details", - zap.String("project_id", projectID), + zap.String("project_id", cfg.ProjectID), zap.String("service_id", input.ServiceID)) // Make API call to get service details ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, projectID, input.ServiceID) + resp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, cfg.ProjectID, input.ServiceID) if err != nil { return nil, ServiceShowOutput{}, fmt.Errorf("failed to get service details: %w", err) } @@ -348,7 +348,7 @@ func (s *Server) handleServiceShow(ctx context.Context, req *mcp.CallToolRequest case 403: return nil, ServiceShowOutput{}, fmt.Errorf("permission denied: insufficient access to service") case 404: - return nil, ServiceShowOutput{}, fmt.Errorf("service '%s' not found in project '%s'", input.ServiceID, projectID) + return nil, ServiceShowOutput{}, fmt.Errorf("service '%s' not found in project '%s'", input.ServiceID, cfg.ProjectID) default: return nil, ServiceShowOutput{}, fmt.Errorf("API request failed with status %d", resp.StatusCode()) } @@ -362,8 +362,8 @@ func (s *Server) handleServiceCreate(ctx context.Context, req *mcp.CallToolReque return nil, ServiceCreateOutput{}, err } - // Load fresh project ID from current config - projectID, err := s.loadProjectID() + // Load fresh config and validate project ID is set + cfg, err := s.loadConfigWithProjectID() if err != nil { return nil, ServiceCreateOutput{}, err } @@ -396,7 +396,7 @@ func (s *Server) handleServiceCreate(ctx context.Context, req *mcp.CallToolReque } logging.Debug("MCP: Creating service", - zap.String("project_id", projectID), + zap.String("project_id", cfg.ProjectID), zap.String("name", input.Name), zap.String("type", input.Type), zap.String("region", input.Region), @@ -419,7 +419,7 @@ func (s *Server) handleServiceCreate(ctx context.Context, req *mcp.CallToolReque ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := apiClient.PostProjectsProjectIdServicesWithResponse(ctx, projectID, serviceCreateReq) + resp, err := apiClient.PostProjectsProjectIdServicesWithResponse(ctx, cfg.ProjectID, serviceCreateReq) if err != nil { return nil, ServiceCreateOutput{}, fmt.Errorf("failed to create service: %w", err) } @@ -465,7 +465,7 @@ func (s *Server) handleServiceCreate(ctx context.Context, req *mcp.CallToolReque timeout = time.Duration(*input.Timeout) * time.Minute } - output.Service, err = s.waitForServiceReady(apiClient, projectID, serviceID, timeout, serviceStatus) + output.Service, err = s.waitForServiceReady(apiClient, cfg.ProjectID, serviceID, timeout, serviceStatus) if err != nil { output.Message = fmt.Sprintf("Error: %s", err.Error()) } else { @@ -493,14 +493,14 @@ func (s *Server) handleServiceUpdatePassword(ctx context.Context, req *mcp.CallT return nil, ServiceUpdatePasswordOutput{}, err } - // Load fresh project ID from current config - projectID, err := s.loadProjectID() + // Load fresh config and validate project ID is set + cfg, err := s.loadConfigWithProjectID() if err != nil { return nil, ServiceUpdatePasswordOutput{}, err } logging.Debug("MCP: Updating service password", - zap.String("project_id", projectID), + zap.String("project_id", cfg.ProjectID), zap.String("service_id", input.ServiceID)) // Prepare password update request @@ -512,7 +512,7 @@ func (s *Server) handleServiceUpdatePassword(ctx context.Context, req *mcp.CallT ctx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - resp, err := apiClient.PostProjectsProjectIdServicesServiceIdUpdatePasswordWithResponse(ctx, projectID, input.ServiceID, updateReq) + resp, err := apiClient.PostProjectsProjectIdServicesServiceIdUpdatePasswordWithResponse(ctx, cfg.ProjectID, input.ServiceID, updateReq) if err != nil { return nil, ServiceUpdatePasswordOutput{}, fmt.Errorf("failed to update service password: %w", err) } @@ -525,7 +525,7 @@ func (s *Server) handleServiceUpdatePassword(ctx context.Context, req *mcp.CallT } // Get service details for password storage (similar to CLI implementation) - serviceResp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, projectID, input.ServiceID) + serviceResp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, cfg.ProjectID, input.ServiceID) if err == nil && serviceResp.StatusCode() == 200 && serviceResp.JSON200 != nil { // Save the new password using the shared util function result, err := util.SavePasswordWithResult(api.Service(*serviceResp.JSON200), input.Password) @@ -544,7 +544,7 @@ func (s *Server) handleServiceUpdatePassword(ctx context.Context, req *mcp.CallT case 403: return nil, ServiceUpdatePasswordOutput{}, fmt.Errorf("permission denied: insufficient access to update service password") case 404: - return nil, ServiceUpdatePasswordOutput{}, fmt.Errorf("service '%s' not found in project '%s'", input.ServiceID, projectID) + return nil, ServiceUpdatePasswordOutput{}, fmt.Errorf("service '%s' not found in project '%s'", input.ServiceID, cfg.ProjectID) case 400: var errorMsg string if resp.JSON400 != nil && resp.JSON400.Message != nil { From c3ddb6c70b3be696ef53756773a2ec41acdd02a5 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Mon, 29 Sep 2025 21:37:28 -0400 Subject: [PATCH 02/13] Consolidate logic for building connection strings --- internal/tiger/cmd/db.go | 92 +----- internal/tiger/cmd/db_test.go | 389 ++-------------------- internal/tiger/mcp/db.go | 57 +--- internal/tiger/util/connection.go | 124 +++++++ internal/tiger/util/connection_test.go | 441 +++++++++++++++++++++++++ 5 files changed, 613 insertions(+), 490 deletions(-) create mode 100644 internal/tiger/util/connection.go create mode 100644 internal/tiger/util/connection_test.go diff --git a/internal/tiger/cmd/db.go b/internal/tiger/cmd/db.go index 09630dca..90043d97 100644 --- a/internal/tiger/cmd/db.go +++ b/internal/tiger/cmd/db.go @@ -60,7 +60,12 @@ Examples: return err } - connectionString, err := buildConnectionString(service, dbConnectionStringPooled, dbConnectionStringRole, dbConnectionStringWithPassword, cmd) + connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: dbConnectionStringPooled, + Role: dbConnectionStringRole, + WithPassword: dbConnectionStringWithPassword, + WarnWriter: cmd.ErrOrStderr(), + }) if err != nil { return fmt.Errorf("failed to build connection string: %w", err) } @@ -133,7 +138,12 @@ Examples: } // Get connection string using existing logic - connectionString, err := buildConnectionString(service, dbConnectPooled, dbConnectRole, false, cmd) + connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: dbConnectPooled, + Role: dbConnectRole, + WithPassword: false, + WarnWriter: cmd.ErrOrStderr(), + }) if err != nil { return fmt.Errorf("failed to build connection string: %w", err) } @@ -192,7 +202,12 @@ Examples: } // Build connection string for testing - connectionString, err := buildConnectionString(service, dbTestConnectionPooled, dbTestConnectionRole, false, cmd) + connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: dbTestConnectionPooled, + Role: dbTestConnectionRole, + WithPassword: false, + WarnWriter: cmd.ErrOrStderr(), + }) if err != nil { return exitWithCode(ExitInvalidParameters, fmt.Errorf("failed to build connection string: %w", err)) } @@ -229,77 +244,6 @@ func buildDbCmd() *cobra.Command { return cmd } -// buildConnectionString creates a PostgreSQL connection string from service details -func buildConnectionString(service api.Service, pooled bool, role string, withPassword bool, cmd *cobra.Command) (string, error) { - if service.Endpoint == nil { - return "", fmt.Errorf("service endpoint not available") - } - - var endpoint *api.Endpoint - var host string - var port int - - // Use pooler endpoint if requested and available, otherwise use direct endpoint - if pooled && service.ConnectionPooler != nil && service.ConnectionPooler.Endpoint != nil { - endpoint = service.ConnectionPooler.Endpoint - } else { - // If pooled was requested but no pooler is available, warn the user - if pooled { - fmt.Fprintf(cmd.ErrOrStderr(), "⚠️ Warning: Connection pooler not available for this service, using direct connection\n") - } - endpoint = service.Endpoint - } - - if endpoint.Host == nil { - return "", fmt.Errorf("endpoint host not available") - } - host = *endpoint.Host - - if endpoint.Port != nil { - port = *endpoint.Port - } else { - port = 5432 // Default PostgreSQL port - } - - // Database is always "tsdb" for TimescaleDB/PostgreSQL services - database := "tsdb" - - // Build connection string in PostgreSQL URI format - var connectionString string - if withPassword { - // Get password from storage if requested - storage := util.GetPasswordStorage() - password, err := storage.Get(service) - if err != nil { - // Provide specific error messages based on storage type - switch storage.(type) { - case *util.NoStorage: - return "", fmt.Errorf("password storage is disabled (--password-storage=none)") - case *util.KeyringStorage: - return "", fmt.Errorf("no password found in keyring for this service") - case *util.PgpassStorage: - return "", fmt.Errorf("no password found in ~/.pgpass for this service") - default: - return "", fmt.Errorf("failed to retrieve password: %w", err) - } - } - - if password == "" { - return "", fmt.Errorf("no password available for service") - } - - // Include password in connection string - connectionString = fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", role, password, host, port, database) - } else { - // Build connection string without password (default behavior) - // Password is handled separately via PGPASSWORD env var or ~/.pgpass file - // This ensures credentials are never visible in process arguments - connectionString = fmt.Sprintf("postgresql://%s@%s:%d/%s?sslmode=require", role, host, port, database) - } - - return connectionString, nil -} - // getServiceDetails is a helper that handles common service lookup logic and returns the service details func getServiceDetails(cmd *cobra.Command, args []string) (api.Service, error) { // Get config diff --git a/internal/tiger/cmd/db_test.go b/internal/tiger/cmd/db_test.go index 2313f97a..b2737e13 100644 --- a/internal/tiger/cmd/db_test.go +++ b/internal/tiger/cmd/db_test.go @@ -126,7 +126,7 @@ func TestDBConnectionString_NoAuth(t *testing.T) { func TestDBConnectionString_PoolerWarning(t *testing.T) { // This test demonstrates that the warning functionality works - // by directly testing the buildConnectionString function + // by directly testing the util.BuildConnectionString function // Service without connection pooler service := api.Service{ @@ -137,13 +137,16 @@ func TestDBConnectionString_PoolerWarning(t *testing.T) { ConnectionPooler: nil, // No pooler available } - // Create a test command to capture stderr - cmd := &cobra.Command{} + // Create a buffer to capture stderr errBuf := new(bytes.Buffer) - cmd.SetErr(errBuf) // Request pooled connection when pooler is not available - connectionString, err := buildConnectionString(service, true, "tsdbadmin", false, cmd) + connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: true, + Role: "tsdbadmin", + WithPassword: false, + WarnWriter: errBuf, + }) if err != nil { t.Fatalf("Unexpected error: %v", err) @@ -782,152 +785,6 @@ func TestIsConnectionRejected(t *testing.T) { } } -func TestBuildConnectionString(t *testing.T) { - testCases := []struct { - name string - service api.Service - pooled bool - role string - expectedString string - expectError bool - expectWarning bool - }{ - { - name: "Basic connection string", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("test-host.tigerdata.com"), - Port: util.Ptr(5432), - }, - }, - pooled: false, - role: "tsdbadmin", - expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", - expectError: false, - }, - { - name: "Connection string with custom role", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("test-host.tigerdata.com"), - Port: util.Ptr(5432), - }, - }, - pooled: false, - role: "readonly", - expectedString: "postgresql://readonly@test-host.tigerdata.com:5432/tsdb?sslmode=require", - expectError: false, - }, - { - name: "Connection string with default port", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("test-host.tigerdata.com"), - Port: nil, // Should use default 5432 - }, - }, - pooled: false, - role: "tsdbadmin", - expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", - expectError: false, - }, - { - name: "Pooled connection string", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("direct-host.tigerdata.com"), - Port: util.Ptr(5432), - }, - ConnectionPooler: &api.ConnectionPooler{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("pooler-host.tigerdata.com"), - Port: util.Ptr(6432), - }, - }, - }, - pooled: true, - role: "tsdbadmin", - expectedString: "postgresql://tsdbadmin@pooler-host.tigerdata.com:6432/tsdb?sslmode=require", - expectError: false, - }, - { - name: "Pooled connection fallback to direct when pooler unavailable", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: util.Ptr("direct-host.tigerdata.com"), - Port: util.Ptr(5432), - }, - ConnectionPooler: nil, // No pooler available - }, - pooled: true, - role: "tsdbadmin", - expectedString: "postgresql://tsdbadmin@direct-host.tigerdata.com:5432/tsdb?sslmode=require", - expectError: false, - expectWarning: true, // Should warn about pooler not available - }, - { - name: "Error when no endpoint available", - service: api.Service{ - Endpoint: nil, - }, - pooled: false, - role: "tsdbadmin", - expectError: true, - }, - { - name: "Error when no host available", - service: api.Service{ - Endpoint: &api.Endpoint{ - Host: nil, - Port: util.Ptr(5432), - }, - }, - pooled: false, - role: "tsdbadmin", - expectError: true, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - // Create a test command to capture stderr output - cmd := &cobra.Command{} - errBuf := new(bytes.Buffer) - cmd.SetErr(errBuf) - - result, err := buildConnectionString(tc.service, tc.pooled, tc.role, false, cmd) - - if tc.expectError { - if err == nil { - t.Errorf("Expected error but got none") - } - return - } - - if err != nil { - t.Errorf("Unexpected error: %v", err) - return - } - - if result != tc.expectedString { - t.Errorf("Expected connection string %q, got %q", tc.expectedString, result) - } - - // Check for warning message - stderrOutput := errBuf.String() - if tc.expectWarning { - if !strings.Contains(stderrOutput, "Warning: Connection pooler not available") { - t.Errorf("Expected warning about pooler not available, but got: %q", stderrOutput) - } - } else { - if stderrOutput != "" { - t.Errorf("Expected no warning, but got: %q", stderrOutput) - } - } - }) - } -} - func TestDBTestConnection_TimeoutParsing(t *testing.T) { testCases := []struct { name string @@ -1033,182 +890,6 @@ func TestDBTestConnection_TimeoutParsing(t *testing.T) { } } -func TestBuildConnectionString_WithPassword_KeyringStorage(t *testing.T) { - // Set keyring as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "keyring") - defer viper.Set("password_storage", originalStorage) - - // Create a test service - serviceID := "test-password-service" - projectID := "test-password-project" - host := "test-host.com" - port := 5432 - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - Endpoint: &api.Endpoint{ - Host: &host, - Port: &port, - }, - } - - // Store a test password in keyring - testPassword := "test-password-keyring-123" - storage := util.GetPasswordStorage() - err := storage.Save(service, testPassword) - if err != nil { - t.Fatalf("Failed to save test password: %v", err) - } - defer storage.Remove(service) // Clean up after test - - // Create a test command - cmd := &cobra.Command{} - - // Call buildConnectionString with withPassword=true - result, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) - - if err != nil { - t.Fatalf("buildConnectionString failed: %v", err) - } - - // Verify that the password is included in the result - expectedResult := fmt.Sprintf("postgresql://tsdbadmin:%s@%s:%d/tsdb?sslmode=require", testPassword, host, port) - if result != expectedResult { - t.Errorf("Expected connection string with password '%s', got '%s'", expectedResult, result) - } - - // Verify the password is actually in the connection string - if !strings.Contains(result, testPassword) { - t.Errorf("Password '%s' not found in connection string: %s", testPassword, result) - } -} - -func TestBuildConnectionString_WithPassword_PgpassStorage(t *testing.T) { - // Set pgpass as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "pgpass") - defer viper.Set("password_storage", originalStorage) - - // Create a test service with endpoint information (required for pgpass) - serviceID := "test-pgpass-service" - projectID := "test-pgpass-project" - host := "test-pgpass-host.com" - port := 5432 - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - Endpoint: &api.Endpoint{ - Host: &host, - Port: &port, - }, - } - - // Store a test password in pgpass - testPassword := "test-password-pgpass-456" - storage := util.GetPasswordStorage() - err := storage.Save(service, testPassword) - if err != nil { - t.Fatalf("Failed to save test password: %v", err) - } - defer storage.Remove(service) // Clean up after test - - // Create a test command - cmd := &cobra.Command{} - - // Call buildConnectionString with withPassword=true - result, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) - - if err != nil { - t.Fatalf("buildConnectionString failed: %v", err) - } - - // Verify that the password is included in the result - expectedResult := fmt.Sprintf("postgresql://tsdbadmin:%s@%s:%d/tsdb?sslmode=require", testPassword, host, port) - if result != expectedResult { - t.Errorf("Expected connection string with password '%s', got '%s'", expectedResult, result) - } - - // Verify the password is actually in the connection string - if !strings.Contains(result, testPassword) { - t.Errorf("Password '%s' not found in connection string: %s", testPassword, result) - } -} - -func TestBuildConnectionString_WithPassword_NoStorage(t *testing.T) { - // Set no storage as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "none") - defer viper.Set("password_storage", originalStorage) - - // Create a test service - serviceID := "test-nostorage-service" - projectID := "test-nostorage-project" - host := "test-host.com" - port := 5432 - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - Endpoint: &api.Endpoint{ - Host: &host, - Port: &port, - }, - } - - // Create a test command - cmd := &cobra.Command{} - - // Call buildConnectionString with withPassword=true - should fail - _, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) - - if err == nil { - t.Fatal("Expected error when password storage is disabled, but got none") - } - - // Verify we get the expected error message - expectedError := "password storage is disabled (--password-storage=none)" - if !strings.Contains(err.Error(), expectedError) { - t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) - } -} - -func TestBuildConnectionString_WithPassword_NoPasswordAvailable(t *testing.T) { - // Set keyring as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "keyring") - defer viper.Set("password_storage", originalStorage) - - // Create a test service (but don't store any password for it) - serviceID := "test-nopassword-service" - projectID := "test-nopassword-project" - host := "test-host.com" - port := 5432 - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - Endpoint: &api.Endpoint{ - Host: &host, - Port: &port, - }, - } - - // Create a test command - cmd := &cobra.Command{} - - // Call buildConnectionString with withPassword=true - should fail - _, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) - - if err == nil { - t.Fatal("Expected error when no password is available, but got none") - } - - // Verify we get the expected error message - expectedError := "no password found in keyring for this service" - if !strings.Contains(err.Error(), expectedError) { - t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) - } -} - func TestDBConnectionString_WithPassword(t *testing.T) { // This test verifies the end-to-end --with-password flag functionality // using direct function testing since full integration would require a real service @@ -1241,11 +922,16 @@ func TestDBConnectionString_WithPassword(t *testing.T) { } defer storage.Remove(service) // Clean up after test - // Test buildConnectionString without password (default behavior) + // Test util.BuildConnectionString without password (default behavior) cmd := &cobra.Command{} - baseConnectionString, err := buildConnectionString(service, false, "tsdbadmin", false, cmd) + baseConnectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: false, + WarnWriter: cmd.ErrOrStderr(), + }) if err != nil { - t.Fatalf("buildConnectionString failed: %v", err) + t.Fatalf("BuildConnectionString failed: %v", err) } expectedBase := fmt.Sprintf("postgresql://tsdbadmin@%s:%d/tsdb?sslmode=require", host, port) @@ -1258,10 +944,15 @@ func TestDBConnectionString_WithPassword(t *testing.T) { t.Errorf("Base connection string should not contain password, but it does: %s", baseConnectionString) } - // Test buildConnectionString with password (simulating --with-password flag) - connectionStringWithPassword, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) + // Test util.BuildConnectionString with password (simulating --with-password flag) + connectionStringWithPassword, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + WarnWriter: cmd.ErrOrStderr(), + }) if err != nil { - t.Fatalf("buildConnectionString with password failed: %v", err) + t.Fatalf("BuildConnectionString with password failed: %v", err) } expectedWithPassword := fmt.Sprintf("postgresql://tsdbadmin:%s@%s:%d/tsdb?sslmode=require", testPassword, host, port) @@ -1274,35 +965,3 @@ func TestDBConnectionString_WithPassword(t *testing.T) { t.Errorf("Connection string with password should contain '%s', but it doesn't: %s", testPassword, connectionStringWithPassword) } } - -func TestBuildConnectionString_WithPassword_InvalidServiceEndpoint(t *testing.T) { - // Set keyring as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "keyring") - defer viper.Set("password_storage", originalStorage) - - // Create a test service without endpoint (invalid) - serviceID := "test-invalid-service" - projectID := "test-invalid-project" - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - Endpoint: nil, // Invalid - no endpoint - } - - // Create a test command - cmd := &cobra.Command{} - - // Call buildConnectionString with withPassword=true - should fail - _, err := buildConnectionString(service, false, "tsdbadmin", true, cmd) - - if err == nil { - t.Fatal("Expected error for invalid service endpoint, but got none") - } - - // Verify we get an endpoint error - expectedError := "service endpoint not available" - if !strings.Contains(err.Error(), expectedError) { - t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) - } -} diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db.go index 237d39e0..82e9dc11 100644 --- a/internal/tiger/mcp/db.go +++ b/internal/tiger/mcp/db.go @@ -11,7 +11,6 @@ import ( "github.com/modelcontextprotocol/go-sdk/mcp" "go.uber.org/zap" - "github.com/timescale/tiger-cli/internal/tiger/api" "github.com/timescale/tiger-cli/internal/tiger/logging" "github.com/timescale/tiger-cli/internal/tiger/util" ) @@ -130,7 +129,12 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ service := *serviceResp.JSON200 // Build connection string with password (use direct connection, default role tsdbadmin) - connString, err := s.buildConnectionString(service, false, "tsdbadmin") + connString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, // MCP always includes password + WarnWriter: nil, // No warnings in MCP context + }) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to build connection string: %w", err) } @@ -186,52 +190,3 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ return nil, output, nil } - -// buildConnectionString creates a PostgreSQL connection string from service details with password included -func (s *Server) buildConnectionString(service api.Service, pooled bool, role string) (string, error) { - if service.Endpoint == nil { - return "", fmt.Errorf("service endpoint not available") - } - - var endpoint *api.Endpoint - var host string - var port int - - // Use pooler endpoint if requested and available, otherwise use direct endpoint - if pooled && service.ConnectionPooler != nil && service.ConnectionPooler.Endpoint != nil { - endpoint = service.ConnectionPooler.Endpoint - } else { - // If pooled was requested but no pooler is available, use direct endpoint - endpoint = service.Endpoint - } - - if endpoint.Host == nil { - return "", fmt.Errorf("endpoint host not available") - } - host = *endpoint.Host - - if endpoint.Port != nil { - port = *endpoint.Port - } else { - port = 5432 // Default PostgreSQL port - } - - // Get password from storage - storage := util.GetPasswordStorage() - password, err := storage.Get(service) - if err != nil { - return "", fmt.Errorf("failed to retrieve password: %w", err) - } - - if password == "" { - return "", fmt.Errorf("no password available for service") - } - - // Database is always "tsdb" for TimescaleDB/PostgreSQL services - database := "tsdb" - - // Build connection string with password included - connectionString := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", role, password, host, port, database) - - return connectionString, nil -} diff --git a/internal/tiger/util/connection.go b/internal/tiger/util/connection.go new file mode 100644 index 00000000..6f785ab6 --- /dev/null +++ b/internal/tiger/util/connection.go @@ -0,0 +1,124 @@ +package util + +import ( + "fmt" + "io" + + "github.com/timescale/tiger-cli/internal/tiger/api" +) + +// ConnectionStringOptions configures how the connection string is built +type ConnectionStringOptions struct { + // Pooled determines whether to use the pooler endpoint (if available) + Pooled bool + + // Role is the database role/username to use (e.g., "tsdbadmin") + Role string + + // WithPassword determines whether to include the password in the connection string + // If false, the connection string will not include a password (for use with PGPASSWORD env var or ~/.pgpass) + WithPassword bool + + // WarnWriter is an optional writer for warning messages (e.g., when pooler is requested but not available) + // If nil, warnings are suppressed + WarnWriter io.Writer +} + +// BuildConnectionString creates a PostgreSQL connection string from service details +// +// The function supports various configuration options through ConnectionStringOptions: +// - Pooled connections (if available on the service) +// - With or without password embedded in the URI +// - Custom database role/username +// - Optional warning output when pooler is unavailable +// +// Examples: +// +// // Simple connection string without password (for use with PGPASSWORD or ~/.pgpass) +// connStr, err := BuildConnectionString(service, ConnectionStringOptions{ +// Role: "tsdbadmin", +// WithPassword: false, +// }) +// +// // Connection string with password embedded +// connStr, err := BuildConnectionString(service, ConnectionStringOptions{ +// Role: "tsdbadmin", +// WithPassword: true, +// }) +// +// // Pooled connection with warnings +// connStr, err := BuildConnectionString(service, ConnectionStringOptions{ +// Pooled: true, +// Role: "tsdbadmin", +// WithPassword: true, +// WarnWriter: os.Stderr, +// }) +func BuildConnectionString(service api.Service, opts ConnectionStringOptions) (string, error) { + if service.Endpoint == nil { + return "", fmt.Errorf("service endpoint not available") + } + + var endpoint *api.Endpoint + var host string + var port int + + // Use pooler endpoint if requested and available, otherwise use direct endpoint + if opts.Pooled && service.ConnectionPooler != nil && service.ConnectionPooler.Endpoint != nil { + endpoint = service.ConnectionPooler.Endpoint + } else { + // If pooled was requested but no pooler is available, warn if writer is provided + if opts.Pooled && opts.WarnWriter != nil { + fmt.Fprintf(opts.WarnWriter, "⚠️ Warning: Connection pooler not available for this service, using direct connection\n") + } + endpoint = service.Endpoint + } + + if endpoint.Host == nil { + return "", fmt.Errorf("endpoint host not available") + } + host = *endpoint.Host + + if endpoint.Port != nil { + port = *endpoint.Port + } else { + port = 5432 // Default PostgreSQL port + } + + // Database is always "tsdb" for TimescaleDB/PostgreSQL services + database := "tsdb" + + // Build connection string in PostgreSQL URI format + var connectionString string + if opts.WithPassword { + // Get password from storage + storage := GetPasswordStorage() + password, err := storage.Get(service) + if err != nil { + // Provide specific error messages based on storage type + switch storage.(type) { + case *NoStorage: + return "", fmt.Errorf("password storage is disabled (--password-storage=none)") + case *KeyringStorage: + return "", fmt.Errorf("no password found in keyring for this service") + case *PgpassStorage: + return "", fmt.Errorf("no password found in ~/.pgpass for this service") + default: + return "", fmt.Errorf("failed to retrieve password: %w", err) + } + } + + if password == "" { + return "", fmt.Errorf("no password available for service") + } + + // Include password in connection string + connectionString = fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", opts.Role, password, host, port, database) + } else { + // Build connection string without password (default behavior) + // Password is handled separately via PGPASSWORD env var or ~/.pgpass file + // This ensures credentials are never visible in process arguments + connectionString = fmt.Sprintf("postgresql://%s@%s:%d/%s?sslmode=require", opts.Role, host, port, database) + } + + return connectionString, nil +} diff --git a/internal/tiger/util/connection_test.go b/internal/tiger/util/connection_test.go new file mode 100644 index 00000000..a9b7a790 --- /dev/null +++ b/internal/tiger/util/connection_test.go @@ -0,0 +1,441 @@ +package util + +import ( + "bytes" + "fmt" + "strings" + "testing" + + "github.com/spf13/viper" + + "github.com/timescale/tiger-cli/internal/tiger/api" +) + +func TestBuildConnectionString_Basic(t *testing.T) { + testCases := []struct { + name string + service api.Service + opts ConnectionStringOptions + expectedString string + expectError bool + expectWarning bool + }{ + { + name: "Basic connection string without password", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("test-host.tigerdata.com"), + Port: Ptr(5432), + }, + }, + opts: ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: false, + }, + expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", + expectError: false, + }, + { + name: "Connection string with custom role", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("test-host.tigerdata.com"), + Port: Ptr(5432), + }, + }, + opts: ConnectionStringOptions{ + Pooled: false, + Role: "readonly", + WithPassword: false, + }, + expectedString: "postgresql://readonly@test-host.tigerdata.com:5432/tsdb?sslmode=require", + expectError: false, + }, + { + name: "Connection string with default port", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("test-host.tigerdata.com"), + Port: nil, // Should use default 5432 + }, + }, + opts: ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: false, + }, + expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", + expectError: false, + }, + { + name: "Pooled connection string", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("direct-host.tigerdata.com"), + Port: Ptr(5432), + }, + ConnectionPooler: &api.ConnectionPooler{ + Endpoint: &api.Endpoint{ + Host: Ptr("pooler-host.tigerdata.com"), + Port: Ptr(6432), + }, + }, + }, + opts: ConnectionStringOptions{ + Pooled: true, + Role: "tsdbadmin", + WithPassword: false, + }, + expectedString: "postgresql://tsdbadmin@pooler-host.tigerdata.com:6432/tsdb?sslmode=require", + expectError: false, + }, + { + name: "Pooled connection fallback to direct when pooler unavailable", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("direct-host.tigerdata.com"), + Port: Ptr(5432), + }, + ConnectionPooler: nil, // No pooler available + }, + opts: ConnectionStringOptions{ + Pooled: true, + Role: "tsdbadmin", + WithPassword: false, + WarnWriter: new(bytes.Buffer), // Enable warnings + }, + expectedString: "postgresql://tsdbadmin@direct-host.tigerdata.com:5432/tsdb?sslmode=require", + expectError: false, + expectWarning: true, // Should warn about pooler not available + }, + { + name: "Error when no endpoint available", + service: api.Service{ + Endpoint: nil, + }, + opts: ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: false, + }, + expectError: true, + }, + { + name: "Error when no host available", + service: api.Service{ + Endpoint: &api.Endpoint{ + Host: nil, + Port: Ptr(5432), + }, + }, + opts: ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: false, + }, + expectError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // If expecting a warning, create a buffer for WarnWriter + var warnBuf *bytes.Buffer + if tc.expectWarning && tc.opts.WarnWriter == nil { + warnBuf = new(bytes.Buffer) + tc.opts.WarnWriter = warnBuf + } else if !tc.expectWarning && tc.opts.WarnWriter != nil { + warnBuf = tc.opts.WarnWriter.(*bytes.Buffer) + } + + result, err := BuildConnectionString(tc.service, tc.opts) + + if tc.expectError { + if err == nil { + t.Errorf("Expected error but got none") + } + return + } + + if err != nil { + t.Errorf("Unexpected error: %v", err) + return + } + + if result != tc.expectedString { + t.Errorf("Expected connection string %q, got %q", tc.expectedString, result) + } + + // Check for warning message + if warnBuf != nil { + stderrOutput := warnBuf.String() + if tc.expectWarning { + if !strings.Contains(stderrOutput, "Warning: Connection pooler not available") { + t.Errorf("Expected warning about pooler not available, but got: %q", stderrOutput) + } + } else { + if stderrOutput != "" { + t.Errorf("Expected no warning, but got: %q", stderrOutput) + } + } + } + }) + } +} + +func TestBuildConnectionString_WithPassword_KeyringStorage(t *testing.T) { + // Set keyring as the password storage method for this test + originalStorage := viper.GetString("password_storage") + viper.Set("password_storage", "keyring") + defer viper.Set("password_storage", originalStorage) + + // Create a test service + serviceID := "test-password-service" + projectID := "test-password-project" + host := "test-host.com" + port := 5432 + service := api.Service{ + ServiceId: &serviceID, + ProjectId: &projectID, + Endpoint: &api.Endpoint{ + Host: &host, + Port: &port, + }, + } + + // Store a test password in keyring + testPassword := "test-password-keyring-123" + storage := GetPasswordStorage() + err := storage.Save(service, testPassword) + if err != nil { + t.Fatalf("Failed to save test password: %v", err) + } + defer storage.Remove(service) // Clean up after test + + // Call BuildConnectionString with withPassword=true + result, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + }) + + if err != nil { + t.Fatalf("BuildConnectionString failed: %v", err) + } + + // Verify that the password is included in the result + expectedResult := fmt.Sprintf("postgresql://tsdbadmin:%s@%s:%d/tsdb?sslmode=require", testPassword, host, port) + if result != expectedResult { + t.Errorf("Expected connection string with password '%s', got '%s'", expectedResult, result) + } + + // Verify the password is actually in the connection string + if !strings.Contains(result, testPassword) { + t.Errorf("Password '%s' not found in connection string: %s", testPassword, result) + } +} + +func TestBuildConnectionString_WithPassword_PgpassStorage(t *testing.T) { + // Set pgpass as the password storage method for this test + originalStorage := viper.GetString("password_storage") + viper.Set("password_storage", "pgpass") + defer viper.Set("password_storage", originalStorage) + + // Create a test service with endpoint information (required for pgpass) + serviceID := "test-pgpass-service" + projectID := "test-pgpass-project" + host := "test-pgpass-host.com" + port := 5432 + service := api.Service{ + ServiceId: &serviceID, + ProjectId: &projectID, + Endpoint: &api.Endpoint{ + Host: &host, + Port: &port, + }, + } + + // Store a test password in pgpass + testPassword := "test-password-pgpass-456" + storage := GetPasswordStorage() + err := storage.Save(service, testPassword) + if err != nil { + t.Fatalf("Failed to save test password: %v", err) + } + defer storage.Remove(service) // Clean up after test + + // Call BuildConnectionString with withPassword=true + result, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + }) + + if err != nil { + t.Fatalf("BuildConnectionString failed: %v", err) + } + + // Verify that the password is included in the result + expectedResult := fmt.Sprintf("postgresql://tsdbadmin:%s@%s:%d/tsdb?sslmode=require", testPassword, host, port) + if result != expectedResult { + t.Errorf("Expected connection string with password '%s', got '%s'", expectedResult, result) + } + + // Verify the password is actually in the connection string + if !strings.Contains(result, testPassword) { + t.Errorf("Password '%s' not found in connection string: %s", testPassword, result) + } +} + +func TestBuildConnectionString_WithPassword_NoStorage(t *testing.T) { + // Set no storage as the password storage method for this test + originalStorage := viper.GetString("password_storage") + viper.Set("password_storage", "none") + defer viper.Set("password_storage", originalStorage) + + // Create a test service + serviceID := "test-nostorage-service" + projectID := "test-nostorage-project" + host := "test-host.com" + port := 5432 + service := api.Service{ + ServiceId: &serviceID, + ProjectId: &projectID, + Endpoint: &api.Endpoint{ + Host: &host, + Port: &port, + }, + } + + // Call BuildConnectionString with withPassword=true - should fail + _, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + }) + + if err == nil { + t.Fatal("Expected error when password storage is disabled, but got none") + } + + // Verify we get the expected error message + expectedError := "password storage is disabled (--password-storage=none)" + if !strings.Contains(err.Error(), expectedError) { + t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) + } +} + +func TestBuildConnectionString_WithPassword_NoPasswordAvailable(t *testing.T) { + // Set keyring as the password storage method for this test + originalStorage := viper.GetString("password_storage") + viper.Set("password_storage", "keyring") + defer viper.Set("password_storage", originalStorage) + + // Create a test service (but don't store any password for it) + serviceID := "test-nopassword-service" + projectID := "test-nopassword-project" + host := "test-host.com" + port := 5432 + service := api.Service{ + ServiceId: &serviceID, + ProjectId: &projectID, + Endpoint: &api.Endpoint{ + Host: &host, + Port: &port, + }, + } + + // Call BuildConnectionString with withPassword=true - should fail + _, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + }) + + if err == nil { + t.Fatal("Expected error when no password is available, but got none") + } + + // Verify we get the expected error message + expectedError := "no password found in keyring for this service" + if !strings.Contains(err.Error(), expectedError) { + t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) + } +} + +func TestBuildConnectionString_WithPassword_InvalidServiceEndpoint(t *testing.T) { + // Set keyring as the password storage method for this test + originalStorage := viper.GetString("password_storage") + viper.Set("password_storage", "keyring") + defer viper.Set("password_storage", originalStorage) + + // Create a test service without endpoint (invalid) + serviceID := "test-invalid-service" + projectID := "test-invalid-project" + service := api.Service{ + ServiceId: &serviceID, + ProjectId: &projectID, + Endpoint: nil, // Invalid - no endpoint + } + + // Call BuildConnectionString with withPassword=true - should fail + _, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + WithPassword: true, + }) + + if err == nil { + t.Fatal("Expected error for invalid service endpoint, but got none") + } + + // Verify we get an endpoint error + expectedError := "service endpoint not available" + if !strings.Contains(err.Error(), expectedError) { + t.Errorf("Expected error message to contain '%s', got: %v", expectedError, err) + } +} + +func TestBuildConnectionString_PoolerWarning(t *testing.T) { + // Service without connection pooler + service := api.Service{ + Endpoint: &api.Endpoint{ + Host: Ptr("test-host.tigerdata.com"), + Port: Ptr(5432), + }, + ConnectionPooler: nil, // No pooler available + } + + // Create a buffer to capture warnings + warnBuf := new(bytes.Buffer) + + // Request pooled connection when pooler is not available + connectionString, err := BuildConnectionString(service, ConnectionStringOptions{ + Pooled: true, + Role: "tsdbadmin", + WithPassword: false, + WarnWriter: warnBuf, + }) + + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + // Should return direct connection string + expectedString := "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require" + if connectionString != expectedString { + t.Errorf("Expected connection string %q, got %q", expectedString, connectionString) + } + + // Should have warning message + stderrOutput := warnBuf.String() + if !strings.Contains(stderrOutput, "Warning: Connection pooler not available") { + t.Errorf("Expected warning about pooler not available, but got: %q", stderrOutput) + } + + // Verify the warning mentions using direct connection + if !strings.Contains(stderrOutput, "using direct connection") { + t.Errorf("Expected warning to mention direct connection fallback, but got: %q", stderrOutput) + } +} From f7c37a4610532c88584b612e3810b05a45fcb501 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Tue, 30 Sep 2025 13:41:33 -0400 Subject: [PATCH 03/13] Get rid of buildConnectionConfig, roll into util.BuildConnectionString --- internal/tiger/cmd/db.go | 58 +++++----------- internal/tiger/cmd/db_test.go | 91 ++------------------------ internal/tiger/mcp/db.go | 4 +- internal/tiger/util/connection.go | 44 +++++++++++-- internal/tiger/util/connection_test.go | 26 ++++---- 5 files changed, 73 insertions(+), 150 deletions(-) diff --git a/internal/tiger/cmd/db.go b/internal/tiger/cmd/db.go index 90043d97..03624d04 100644 --- a/internal/tiger/cmd/db.go +++ b/internal/tiger/cmd/db.go @@ -60,10 +60,15 @@ Examples: return err } + passwordMode := util.PasswordExclude + if dbConnectionStringWithPassword { + passwordMode = util.PasswordRequired + } + connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: dbConnectionStringPooled, Role: dbConnectionStringRole, - WithPassword: dbConnectionStringWithPassword, + PasswordMode: passwordMode, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -137,11 +142,10 @@ Examples: return fmt.Errorf("psql client not found. Please install PostgreSQL client tools") } - // Get connection string using existing logic connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: dbConnectPooled, Role: dbConnectRole, - WithPassword: false, + PasswordMode: util.PasswordExclude, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -201,11 +205,11 @@ Examples: return exitWithCode(ExitInvalidParameters, err) } - // Build connection string for testing + // Build connection string for testing with password (if available) connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: dbTestConnectionPooled, Role: dbTestConnectionRole, - WithPassword: false, + PasswordMode: util.PasswordOptional, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -218,7 +222,7 @@ Examples: } // Test the connection - return testDatabaseConnection(connectionString, dbTestConnectionTimeout, service, cmd) + return testDatabaseConnection(cmd.Context(), connectionString, dbTestConnectionTimeout, cmd) }, } @@ -284,7 +288,7 @@ func getServiceDetails(cmd *cobra.Command, args []string) (api.Service, error) { } // Fetch service details - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second) defer cancel() resp, err := client.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, projectID, serviceID) @@ -371,50 +375,18 @@ func buildPsqlCommand(connectionString, psqlPath string, additionalFlags []strin return psqlCmd } -// buildConnectionConfig creates a pgx connection config with proper password handling -func buildConnectionConfig(connectionString string, service api.Service) (*pgx.ConnConfig, error) { - // Parse the connection string first to validate it - config, err := pgx.ParseConfig(connectionString) - if err != nil { - return nil, err - } - - // Set password from keyring storage if available - // pgpass storage works automatically since pgx checks ~/.pgpass file - storage := util.GetPasswordStorage() - if _, isKeyring := storage.(*util.KeyringStorage); isKeyring { - if password, err := storage.Get(service); err == nil && password != "" { - config.Password = password - } - // Note: If keyring password retrieval fails, we let pgx try without it - // This allows fallback to other authentication methods - } - - return config, nil -} - // testDatabaseConnection tests the database connection and returns appropriate exit codes -func testDatabaseConnection(connectionString string, timeout time.Duration, service api.Service, cmd *cobra.Command) error { +func testDatabaseConnection(ctx context.Context, connectionString string, timeout time.Duration, cmd *cobra.Command) error { // Create context with timeout if specified - var ctx context.Context var cancel context.CancelFunc - if timeout > 0 { - ctx, cancel = context.WithTimeout(context.Background(), timeout) + ctx, cancel = context.WithTimeout(ctx, timeout) defer cancel() - } else { - ctx = context.Background() - } - - // Build connection config with proper password handling - config, err := buildConnectionConfig(connectionString, service) - if err != nil { - fmt.Fprintf(cmd.ErrOrStderr(), "Failed to build connection config: %v\n", err) - return exitWithCode(ExitInvalidParameters, err) } // Attempt to connect to the database - conn, err := pgx.ConnectConfig(ctx, config) + // The connection string already includes the password (if available) thanks to PasswordOptional mode + conn, err := pgx.Connect(ctx, connectionString) if err != nil { // Determine the appropriate exit code based on error type if isContextDeadlineExceeded(err) { diff --git a/internal/tiger/cmd/db_test.go b/internal/tiger/cmd/db_test.go index b2737e13..0696babf 100644 --- a/internal/tiger/cmd/db_test.go +++ b/internal/tiger/cmd/db_test.go @@ -144,7 +144,7 @@ func TestDBConnectionString_PoolerWarning(t *testing.T) { connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: true, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: util.PasswordExclude, WarnWriter: errBuf, }) @@ -426,85 +426,6 @@ func TestBuildPsqlCommand_PgpassStorage_NoEnvVar(t *testing.T) { } } -func TestBuildConnectionConfig_KeyringPassword(t *testing.T) { - // This test verifies that buildConnectionConfig properly sets password from keyring - - // Set keyring as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "keyring") - defer viper.Set("password_storage", originalStorage) - - // Create a test service - serviceID := "test-connection-config-service" - projectID := "test-connection-config-project" - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - } - - // Store a test password in keyring - testPassword := "test-connection-config-password-789" - storage := util.GetPasswordStorage() - err := storage.Save(service, testPassword) - if err != nil { - t.Fatalf("Failed to save test password: %v", err) - } - defer storage.Remove(service) // Clean up after test - - connectionString := "postgresql://testuser@testhost:5432/testdb?sslmode=require" - - // Call the actual production function that builds the config - config, err := buildConnectionConfig(connectionString, service) - - if err != nil { - t.Fatalf("buildConnectionConfig failed: %v", err) - } - - if config == nil { - t.Fatal("buildConnectionConfig returned nil config") - } - - // Verify that the password was set in the config - if config.Password != testPassword { - t.Errorf("Expected password '%s' to be set in config, but got '%s'", testPassword, config.Password) - } -} - -func TestBuildConnectionConfig_PgpassStorage_NoPasswordSet(t *testing.T) { - // This test verifies that buildConnectionConfig doesn't set password for pgpass storage - - // Set pgpass as the password storage method for this test - originalStorage := viper.GetString("password_storage") - viper.Set("password_storage", "pgpass") - defer viper.Set("password_storage", originalStorage) - - // Create a test service - serviceID := "test-connection-config-pgpass" - projectID := "test-connection-config-project" - service := api.Service{ - ServiceId: &serviceID, - ProjectId: &projectID, - } - - connectionString := "postgresql://testuser@testhost:5432/testdb?sslmode=require" - - // Call the actual production function that builds the config - config, err := buildConnectionConfig(connectionString, service) - - if err != nil { - t.Fatalf("buildConnectionConfig failed: %v", err) - } - - if config == nil { - t.Fatal("buildConnectionConfig returned nil config") - } - - // Verify that no password was set in the config (pgx will check ~/.pgpass automatically) - if config.Password != "" { - t.Errorf("Expected no password to be set in config for pgpass storage, but got '%s'", config.Password) - } -} - func TestSeparateServiceAndPsqlArgs(t *testing.T) { testCases := []struct { name string @@ -670,8 +591,7 @@ func TestTestDatabaseConnection_InvalidConnectionString(t *testing.T) { // Test with malformed connection string (should return ExitInvalidParameters) invalidConnectionString := "this is not a valid connection string at all" - service := api.Service{} // Dummy service for test - err := testDatabaseConnection(invalidConnectionString, 1, service, cmd) + err := testDatabaseConnection(invalidConnectionString, 1, cmd) if err == nil { t.Error("Expected error for invalid connection string") @@ -699,9 +619,8 @@ func TestTestDatabaseConnection_Timeout(t *testing.T) { // Use a connection string to a non-routable IP to test timeout timeoutConnectionString := "postgresql://user:pass@192.0.2.1:5432/db?sslmode=disable&connect_timeout=1" - service := api.Service{} // Dummy service for test start := time.Now() - err := testDatabaseConnection(timeoutConnectionString, 1, service, cmd) // 1 second timeout + err := testDatabaseConnection(timeoutConnectionString, 1, cmd) // 1 second timeout duration := time.Since(start) if err == nil { @@ -927,7 +846,7 @@ func TestDBConnectionString_WithPassword(t *testing.T) { baseConnectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: util.PasswordExclude, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -948,7 +867,7 @@ func TestDBConnectionString_WithPassword(t *testing.T) { connectionStringWithPassword, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: util.PasswordRequired, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db.go index 82e9dc11..a32ffee0 100644 --- a/internal/tiger/mcp/db.go +++ b/internal/tiger/mcp/db.go @@ -132,8 +132,8 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ connString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, // MCP always includes password - WarnWriter: nil, // No warnings in MCP context + PasswordMode: util.PasswordRequired, // MCP always requires password + WarnWriter: nil, // No warnings in MCP context }) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to build connection string: %w", err) diff --git a/internal/tiger/util/connection.go b/internal/tiger/util/connection.go index 6f785ab6..caf97910 100644 --- a/internal/tiger/util/connection.go +++ b/internal/tiger/util/connection.go @@ -7,6 +7,23 @@ import ( "github.com/timescale/tiger-cli/internal/tiger/api" ) +// PasswordMode determines how passwords are handled in connection strings +type PasswordMode int + +const ( + // PasswordExclude means don't include password in connection string (default) + // Connection will rely on PGPASSWORD env var or ~/.pgpass file + PasswordExclude PasswordMode = iota + + // PasswordRequired means include password in connection string, return error if unavailable + // Used when user explicitly requests --with-password flag + PasswordRequired + + // PasswordOptional means include password if available, but don't error if unavailable + // Used for connection testing and psql launching where we want best-effort password inclusion + PasswordOptional +) + // ConnectionStringOptions configures how the connection string is built type ConnectionStringOptions struct { // Pooled determines whether to use the pooler endpoint (if available) @@ -15,9 +32,8 @@ type ConnectionStringOptions struct { // Role is the database role/username to use (e.g., "tsdbadmin") Role string - // WithPassword determines whether to include the password in the connection string - // If false, the connection string will not include a password (for use with PGPASSWORD env var or ~/.pgpass) - WithPassword bool + // PasswordMode determines how passwords are handled + PasswordMode PasswordMode // WarnWriter is an optional writer for warning messages (e.g., when pooler is requested but not available) // If nil, warnings are suppressed @@ -89,8 +105,10 @@ func BuildConnectionString(service api.Service, opts ConnectionStringOptions) (s // Build connection string in PostgreSQL URI format var connectionString string - if opts.WithPassword { - // Get password from storage + + switch opts.PasswordMode { + case PasswordRequired: + // Password is required - error if unavailable storage := GetPasswordStorage() password, err := storage.Get(service) if err != nil { @@ -113,7 +131,21 @@ func BuildConnectionString(service api.Service, opts ConnectionStringOptions) (s // Include password in connection string connectionString = fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", opts.Role, password, host, port, database) - } else { + + case PasswordOptional: + // Try to include password, but don't error if unavailable + storage := GetPasswordStorage() + password, err := storage.Get(service) + + // Only include password if we successfully retrieved it + if err == nil && password != "" { + connectionString = fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=require", opts.Role, password, host, port, database) + } else { + // Fall back to connection string without password + connectionString = fmt.Sprintf("postgresql://%s@%s:%d/%s?sslmode=require", opts.Role, host, port, database) + } + + default: // PasswordExclude // Build connection string without password (default behavior) // Password is handled separately via PGPASSWORD env var or ~/.pgpass file // This ensures credentials are never visible in process arguments diff --git a/internal/tiger/util/connection_test.go b/internal/tiger/util/connection_test.go index a9b7a790..046633b5 100644 --- a/internal/tiger/util/connection_test.go +++ b/internal/tiger/util/connection_test.go @@ -31,7 +31,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", expectError: false, @@ -47,7 +47,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: false, Role: "readonly", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectedString: "postgresql://readonly@test-host.tigerdata.com:5432/tsdb?sslmode=require", expectError: false, @@ -63,7 +63,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectedString: "postgresql://tsdbadmin@test-host.tigerdata.com:5432/tsdb?sslmode=require", expectError: false, @@ -85,7 +85,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: true, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectedString: "postgresql://tsdbadmin@pooler-host.tigerdata.com:6432/tsdb?sslmode=require", expectError: false, @@ -102,7 +102,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: true, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, WarnWriter: new(bytes.Buffer), // Enable warnings }, expectedString: "postgresql://tsdbadmin@direct-host.tigerdata.com:5432/tsdb?sslmode=require", @@ -117,7 +117,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectError: true, }, @@ -132,7 +132,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { opts: ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, }, expectError: true, }, @@ -217,7 +217,7 @@ func TestBuildConnectionString_WithPassword_KeyringStorage(t *testing.T) { result, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: PasswordRequired, }) if err != nil { @@ -269,7 +269,7 @@ func TestBuildConnectionString_WithPassword_PgpassStorage(t *testing.T) { result, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: PasswordRequired, }) if err != nil { @@ -312,7 +312,7 @@ func TestBuildConnectionString_WithPassword_NoStorage(t *testing.T) { _, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: PasswordRequired, }) if err == nil { @@ -350,7 +350,7 @@ func TestBuildConnectionString_WithPassword_NoPasswordAvailable(t *testing.T) { _, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: PasswordRequired, }) if err == nil { @@ -383,7 +383,7 @@ func TestBuildConnectionString_WithPassword_InvalidServiceEndpoint(t *testing.T) _, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - WithPassword: true, + PasswordMode: PasswordRequired, }) if err == nil { @@ -414,7 +414,7 @@ func TestBuildConnectionString_PoolerWarning(t *testing.T) { connectionString, err := BuildConnectionString(service, ConnectionStringOptions{ Pooled: true, Role: "tsdbadmin", - WithPassword: false, + PasswordMode: PasswordExclude, WarnWriter: warnBuf, }) From 900275b3bc451dcd4489ad9fead6eaa849586f4d Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sat, 4 Oct 2025 16:29:59 -0400 Subject: [PATCH 04/13] Fix circular reference issues post-merge --- internal/tiger/cmd/db.go | 15 ++++----- internal/tiger/cmd/db_test.go | 32 ++++++++++++------- internal/tiger/cmd/service.go | 11 ++++++- internal/tiger/mcp/db.go | 7 ++-- .../tiger/{util => password}/connection.go | 2 +- .../{util => password}/connection_test.go | 31 +++++++++--------- 6 files changed, 58 insertions(+), 40 deletions(-) rename internal/tiger/{util => password}/connection.go (99%) rename internal/tiger/{util => password}/connection_test.go (95%) diff --git a/internal/tiger/cmd/db.go b/internal/tiger/cmd/db.go index 79f3a07d..286234c9 100644 --- a/internal/tiger/cmd/db.go +++ b/internal/tiger/cmd/db.go @@ -15,7 +15,6 @@ import ( "github.com/timescale/tiger-cli/internal/tiger/api" "github.com/timescale/tiger-cli/internal/tiger/config" "github.com/timescale/tiger-cli/internal/tiger/password" - "github.com/timescale/tiger-cli/internal/tiger/util" ) var ( @@ -61,12 +60,12 @@ Examples: return err } - passwordMode := util.PasswordExclude + passwordMode := password.PasswordExclude if dbConnectionStringWithPassword { - passwordMode = util.PasswordRequired + passwordMode = password.PasswordRequired } - connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + connectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: dbConnectionStringPooled, Role: dbConnectionStringRole, PasswordMode: passwordMode, @@ -143,10 +142,10 @@ Examples: return fmt.Errorf("psql client not found. Please install PostgreSQL client tools") } - connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + connectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: dbConnectPooled, Role: dbConnectRole, - PasswordMode: util.PasswordExclude, + PasswordMode: password.PasswordExclude, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -207,10 +206,10 @@ Examples: } // Build connection string for testing with password (if available) - connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + connectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: dbTestConnectionPooled, Role: dbTestConnectionRole, - PasswordMode: util.PasswordOptional, + PasswordMode: password.PasswordOptional, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { diff --git a/internal/tiger/cmd/db_test.go b/internal/tiger/cmd/db_test.go index 3321ee72..51c830f3 100644 --- a/internal/tiger/cmd/db_test.go +++ b/internal/tiger/cmd/db_test.go @@ -2,6 +2,7 @@ package cmd import ( "bytes" + "context" "fmt" "os" "strings" @@ -121,7 +122,7 @@ func TestDBConnectionString_NoAuth(t *testing.T) { func TestDBConnectionString_PoolerWarning(t *testing.T) { // This test demonstrates that the warning functionality works - // by directly testing the util.BuildConnectionString function + // by directly testing the password.BuildConnectionString function // Service without connection pooler service := api.Service{ @@ -136,10 +137,10 @@ func TestDBConnectionString_PoolerWarning(t *testing.T) { errBuf := new(bytes.Buffer) // Request pooled connection when pooler is not available - connectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + connectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: true, Role: "tsdbadmin", - PasswordMode: util.PasswordExclude, + PasswordMode: password.PasswordExclude, WarnWriter: errBuf, }) @@ -576,7 +577,8 @@ func TestTestDatabaseConnection_InvalidConnectionString(t *testing.T) { // Test with malformed connection string (should return ExitInvalidParameters) invalidConnectionString := "this is not a valid connection string at all" - err := testDatabaseConnection(invalidConnectionString, 1, cmd) + ctx := context.Background() + err := testDatabaseConnection(ctx, invalidConnectionString, 1*time.Second, cmd) if err == nil { t.Error("Expected error for invalid connection string") @@ -604,8 +606,9 @@ func TestTestDatabaseConnection_Timeout(t *testing.T) { // Use a connection string to a non-routable IP to test timeout timeoutConnectionString := "postgresql://user:pass@192.0.2.1:5432/db?sslmode=disable&connect_timeout=1" + ctx := context.Background() start := time.Now() - err := testDatabaseConnection(timeoutConnectionString, 1, cmd) // 1 second timeout + err := testDatabaseConnection(ctx, timeoutConnectionString, 1*time.Second, cmd) // 1 second timeout duration := time.Since(start) if err == nil { @@ -802,7 +805,12 @@ func TestBuildConnectionString(t *testing.T) { errBuf := new(bytes.Buffer) cmd.SetErr(errBuf) - result, err := buildConnectionString(tc.service, tc.pooled, tc.role, false, cmd.ErrOrStderr()) + result, err := password.BuildConnectionString(tc.service, password.ConnectionStringOptions{ + Pooled: tc.pooled, + Role: tc.role, + PasswordMode: password.PasswordExclude, + WarnWriter: cmd.ErrOrStderr(), + }) if tc.expectError { if err == nil { @@ -970,12 +978,12 @@ func TestDBConnectionString_WithPassword(t *testing.T) { } defer storage.Remove(service) // Clean up after test - // Test util.BuildConnectionString without password (default behavior) + // Test password.BuildConnectionString without password (default behavior) cmd := &cobra.Command{} - baseConnectionString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + baseConnectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - PasswordMode: util.PasswordExclude, + PasswordMode: password.PasswordExclude, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { @@ -992,11 +1000,11 @@ func TestDBConnectionString_WithPassword(t *testing.T) { t.Errorf("Base connection string should not contain password, but it does: %s", baseConnectionString) } - // Test util.BuildConnectionString with password (simulating --with-password flag) - connectionStringWithPassword, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + // Test password.BuildConnectionString with password (simulating --with-password flag) + connectionStringWithPassword, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - PasswordMode: util.PasswordRequired, + PasswordMode: password.PasswordRequired, WarnWriter: cmd.ErrOrStderr(), }) if err != nil { diff --git a/internal/tiger/cmd/service.go b/internal/tiger/cmd/service.go index fbf98928..97f11d04 100644 --- a/internal/tiger/cmd/service.go +++ b/internal/tiger/cmd/service.go @@ -809,7 +809,16 @@ func prepareServiceForOutput(service api.Service, withPassword bool, output io.W } // Build connection string - connectionString, err := buildConnectionString(service, false, "tsdbadmin", withPassword, output) + passwordMode := password.PasswordExclude + if withPassword { + passwordMode = password.PasswordRequired + } + connectionString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ + Pooled: false, + Role: "tsdbadmin", + PasswordMode: passwordMode, + WarnWriter: output, + }) if err == nil { outputSvc.ConnectionString = &connectionString } diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db.go index a32ffee0..09bed735 100644 --- a/internal/tiger/mcp/db.go +++ b/internal/tiger/mcp/db.go @@ -12,6 +12,7 @@ import ( "go.uber.org/zap" "github.com/timescale/tiger-cli/internal/tiger/logging" + "github.com/timescale/tiger-cli/internal/tiger/password" "github.com/timescale/tiger-cli/internal/tiger/util" ) @@ -129,11 +130,11 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ service := *serviceResp.JSON200 // Build connection string with password (use direct connection, default role tsdbadmin) - connString, err := util.BuildConnectionString(service, util.ConnectionStringOptions{ + connString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ Pooled: false, Role: "tsdbadmin", - PasswordMode: util.PasswordRequired, // MCP always requires password - WarnWriter: nil, // No warnings in MCP context + PasswordMode: password.PasswordRequired, // MCP always requires password + WarnWriter: nil, // No warnings in MCP context }) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to build connection string: %w", err) diff --git a/internal/tiger/util/connection.go b/internal/tiger/password/connection.go similarity index 99% rename from internal/tiger/util/connection.go rename to internal/tiger/password/connection.go index caf97910..38aaf4e8 100644 --- a/internal/tiger/util/connection.go +++ b/internal/tiger/password/connection.go @@ -1,4 +1,4 @@ -package util +package password import ( "fmt" diff --git a/internal/tiger/util/connection_test.go b/internal/tiger/password/connection_test.go similarity index 95% rename from internal/tiger/util/connection_test.go rename to internal/tiger/password/connection_test.go index 046633b5..e34f3dd7 100644 --- a/internal/tiger/util/connection_test.go +++ b/internal/tiger/password/connection_test.go @@ -1,4 +1,4 @@ -package util +package password import ( "bytes" @@ -9,6 +9,7 @@ import ( "github.com/spf13/viper" "github.com/timescale/tiger-cli/internal/tiger/api" + "github.com/timescale/tiger-cli/internal/tiger/util" ) func TestBuildConnectionString_Basic(t *testing.T) { @@ -24,8 +25,8 @@ func TestBuildConnectionString_Basic(t *testing.T) { name: "Basic connection string without password", service: api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("test-host.tigerdata.com"), - Port: Ptr(5432), + Host: util.Ptr("test-host.tigerdata.com"), + Port: util.Ptr(5432), }, }, opts: ConnectionStringOptions{ @@ -40,8 +41,8 @@ func TestBuildConnectionString_Basic(t *testing.T) { name: "Connection string with custom role", service: api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("test-host.tigerdata.com"), - Port: Ptr(5432), + Host: util.Ptr("test-host.tigerdata.com"), + Port: util.Ptr(5432), }, }, opts: ConnectionStringOptions{ @@ -56,7 +57,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { name: "Connection string with default port", service: api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("test-host.tigerdata.com"), + Host: util.Ptr("test-host.tigerdata.com"), Port: nil, // Should use default 5432 }, }, @@ -72,13 +73,13 @@ func TestBuildConnectionString_Basic(t *testing.T) { name: "Pooled connection string", service: api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("direct-host.tigerdata.com"), - Port: Ptr(5432), + Host: util.Ptr("direct-host.tigerdata.com"), + Port: util.Ptr(5432), }, ConnectionPooler: &api.ConnectionPooler{ Endpoint: &api.Endpoint{ - Host: Ptr("pooler-host.tigerdata.com"), - Port: Ptr(6432), + Host: util.Ptr("pooler-host.tigerdata.com"), + Port: util.Ptr(6432), }, }, }, @@ -94,8 +95,8 @@ func TestBuildConnectionString_Basic(t *testing.T) { name: "Pooled connection fallback to direct when pooler unavailable", service: api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("direct-host.tigerdata.com"), - Port: Ptr(5432), + Host: util.Ptr("direct-host.tigerdata.com"), + Port: util.Ptr(5432), }, ConnectionPooler: nil, // No pooler available }, @@ -126,7 +127,7 @@ func TestBuildConnectionString_Basic(t *testing.T) { service: api.Service{ Endpoint: &api.Endpoint{ Host: nil, - Port: Ptr(5432), + Port: util.Ptr(5432), }, }, opts: ConnectionStringOptions{ @@ -401,8 +402,8 @@ func TestBuildConnectionString_PoolerWarning(t *testing.T) { // Service without connection pooler service := api.Service{ Endpoint: &api.Endpoint{ - Host: Ptr("test-host.tigerdata.com"), - Port: Ptr(5432), + Host: util.Ptr("test-host.tigerdata.com"), + Port: util.Ptr(5432), }, ConnectionPooler: nil, // No pooler available } From 31a53fa90f3a81f5a49fb004e9b48725eacda05a Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:03:21 -0400 Subject: [PATCH 05/13] Minor tool cleanup --- internal/tiger/mcp/db.go | 64 +++++++++++++++------------------------- 1 file changed, 23 insertions(+), 41 deletions(-) diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db.go index 09bed735..7ba6d641 100644 --- a/internal/tiger/mcp/db.go +++ b/internal/tiger/mcp/db.go @@ -18,22 +18,24 @@ import ( // DBExecuteQueryInput represents input for tiger_db_execute_query type DBExecuteQueryInput struct { - ServiceID string `json:"service_id,omitempty"` - Query string `json:"query"` - Timeout *int `json:"timeout,omitempty"` + ServiceID string `json:"service_id"` + Query string `json:"query"` + TimeoutSeconds int `json:"timeout_seconds,omitempty"` } func (DBExecuteQueryInput) Schema() *jsonschema.Schema { schema := util.Must(jsonschema.For[DBExecuteQueryInput](nil)) - schema.Properties["service_id"].Description = "Service ID to execute query on (uses default if not provided)" + schema.Properties["service_id"].Description = "The unique identifier of the service (10-character alphanumeric string). Use service_list to find service IDs." + schema.Properties["service_id"].Examples = []any{"e6ue9697jf", "u8me885b93"} + schema.Properties["service_id"].Pattern = "^[a-z0-9]{10}$" - schema.Properties["query"].Description = "SQL query to execute" + schema.Properties["query"].Description = "PostgreSQL query to execute" - schema.Properties["timeout"].Description = "Query timeout in seconds (default: 30)" - schema.Properties["timeout"].Minimum = util.Ptr(0.0) - schema.Properties["timeout"].Default = util.Must(json.Marshal(30)) - schema.Properties["timeout"].Examples = []any{10, 30, 60} + schema.Properties["timeout_seconds"].Description = "Query timeout in seconds" + schema.Properties["timeout_seconds"].Minimum = util.Ptr(0.0) + schema.Properties["timeout_seconds"].Default = util.Must(json.Marshal(30)) + schema.Properties["timeout_seconds"].Examples = []any{10, 30, 60} return schema } @@ -48,21 +50,14 @@ type DBExecuteQueryOutput struct { // registerDatabaseTools registers database operation tools with comprehensive schemas and descriptions func (s *Server) registerDatabaseTools() { - // tiger_db_execute_query mcp.AddTool(s.mcpServer, &mcp.Tool{ - Name: "tiger_db_execute_query", + Name: "db_execute_query", Title: "Execute SQL Query", Description: `Execute a SQL query against a service database. -This tool connects to a database service and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Perfect for data exploration, schema inspection, and database operations. +This tool connects to a PostgreSQL database service in TigerData Cloud and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Perfect for data exploration, schema inspection, and database operations. -IMPORTANT: Use with caution - this tool can execute any SQL statement including INSERT, UPDATE, DELETE, and DDL commands. Always review queries before execution. - -Perfect for: -- Querying data from tables and views -- Inspecting database schema -- Testing database connectivity with real queries -- Performing data analysis and exploration`, +WARNING: Use with caution - this tool can execute any SQL statement including INSERT, UPDATE, DELETE, and DDL commands. Always review queries before execution.`, InputSchema: DBExecuteQueryInput{}.Schema(), Annotations: &mcp.ToolAnnotations{ DestructiveHint: util.Ptr(true), // Can execute destructive SQL @@ -73,41 +68,29 @@ Perfect for: // handleDBExecuteQuery handles the tiger_db_execute_query MCP tool func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequest, input DBExecuteQueryInput) (*mcp.CallToolResult, DBExecuteQueryOutput, error) { - // Create fresh API client with current credentials - apiClient, err := s.createAPIClient() + // Load config and validate project ID + cfg, err := s.loadConfigWithProjectID() if err != nil { return nil, DBExecuteQueryOutput{}, err } - // Load fresh config and validate project ID is set - cfg, err := s.loadConfigWithProjectID() + // Create fresh API client with current credentials + apiClient, err := s.createAPIClient() if err != nil { return nil, DBExecuteQueryOutput{}, err } - // Get service ID (use default from config if not provided) - serviceID := input.ServiceID - if serviceID == "" { - if cfg.ServiceID == "" { - return nil, DBExecuteQueryOutput{}, fmt.Errorf("service ID is required. Please provide service_id or run 'tiger config set service_id '") - } - serviceID = cfg.ServiceID - } - - // Set default timeout if not provided - timeout := 30 * time.Second - if input.Timeout != nil { - timeout = time.Duration(*input.Timeout) * time.Second - } + // Convert timeout in seconds to time.Duration + timeout := time.Duration(input.TimeoutSeconds) * time.Second logging.Debug("MCP: Executing database query", zap.String("project_id", cfg.ProjectID), - zap.String("service_id", serviceID), + zap.String("service_id", input.ServiceID), zap.Duration("timeout", timeout), ) // Get service details to construct connection string - serviceResp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, cfg.ProjectID, serviceID) + serviceResp, err := apiClient.GetProjectsProjectIdServicesServiceIdWithResponse(ctx, cfg.ProjectID, input.ServiceID) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to get service details: %w", err) } @@ -122,7 +105,7 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ case 403: return nil, DBExecuteQueryOutput{}, fmt.Errorf("permission denied: insufficient access to service") case 404: - return nil, DBExecuteQueryOutput{}, fmt.Errorf("service '%s' not found in project '%s'", serviceID, cfg.ProjectID) + return nil, DBExecuteQueryOutput{}, fmt.Errorf("service '%s' not found in project '%s'", input.ServiceID, cfg.ProjectID) default: return nil, DBExecuteQueryOutput{}, fmt.Errorf("API request failed with status %d", serviceResp.StatusCode()) } @@ -134,7 +117,6 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ Pooled: false, Role: "tsdbadmin", PasswordMode: password.PasswordRequired, // MCP always requires password - WarnWriter: nil, // No warnings in MCP context }) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("failed to build connection string: %w", err) From 01210b8910fbc5d61433b1b2e0365bf5b549a44d Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:04:11 -0400 Subject: [PATCH 06/13] Rename db.go to db_tools.go for parity with service_tools.go --- internal/tiger/mcp/{db.go => db_tools.go} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename internal/tiger/mcp/{db.go => db_tools.go} (100%) diff --git a/internal/tiger/mcp/db.go b/internal/tiger/mcp/db_tools.go similarity index 100% rename from internal/tiger/mcp/db.go rename to internal/tiger/mcp/db_tools.go From 5e7005e385c5befa87ff95c1a6e71643c5b34493 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:11:22 -0400 Subject: [PATCH 07/13] Return rows affected instead of naive row count --- internal/tiger/mcp/db_tools.go | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index 7ba6d641..3b990557 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -44,10 +44,28 @@ func (DBExecuteQueryInput) Schema() *jsonschema.Schema { type DBExecuteQueryOutput struct { Columns []string `json:"columns"` Rows [][]any `json:"rows"` - RowCount int `json:"row_count"` + RowsAffected int64 `json:"rows_affected"` ExecutionTime string `json:"execution_time"` } +func (DBExecuteQueryOutput) Schema() *jsonschema.Schema { + schema := util.Must(jsonschema.For[DBExecuteQueryOutput](nil)) + + schema.Properties["columns"].Description = "Column names from the query result" + schema.Properties["columns"].Examples = []any{[]string{"id", "name", "created_at"}} + + schema.Properties["rows"].Description = "Result rows as arrays of values. Empty for commands that don't return rows (INSERT, UPDATE, DELETE, etc.)" + schema.Properties["rows"].Examples = []any{[][]any{{1, "alice", "2024-01-01"}, {2, "bob", "2024-01-02"}}} + + schema.Properties["rows_affected"].Description = "Number of rows affected by the query. For SELECT, this is the number of rows returned. For INSERT/UPDATE/DELETE, this is the number of rows modified. Returns 0 for statements that don't return or modify rows (e.g. CREATE TABLE)." + schema.Properties["rows_affected"].Examples = []any{5, 42, 1000} + + schema.Properties["execution_time"].Description = "Query execution time as a human-readable duration string" + schema.Properties["execution_time"].Examples = []any{"123ms", "1.5s", "45.2µs"} + + return schema +} + // registerDatabaseTools registers database operation tools with comprehensive schemas and descriptions func (s *Server) registerDatabaseTools() { mcp.AddTool(s.mcpServer, &mcp.Tool{ @@ -58,7 +76,8 @@ func (s *Server) registerDatabaseTools() { This tool connects to a PostgreSQL database service in TigerData Cloud and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Perfect for data exploration, schema inspection, and database operations. WARNING: Use with caution - this tool can execute any SQL statement including INSERT, UPDATE, DELETE, and DDL commands. Always review queries before execution.`, - InputSchema: DBExecuteQueryInput{}.Schema(), + InputSchema: DBExecuteQueryInput{}.Schema(), + OutputSchema: DBExecuteQueryOutput{}.Schema(), Annotations: &mcp.ToolAnnotations{ DestructiveHint: util.Ptr(true), // Can execute destructive SQL Title: "Execute SQL Query", @@ -167,7 +186,7 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ output := DBExecuteQueryOutput{ Columns: columns, Rows: resultRows, - RowCount: len(resultRows), + RowsAffected: rows.CommandTag().RowsAffected(), ExecutionTime: time.Since(startTime).String(), } From eab502c20fe1c0d5f5f307bf22ca29fd6fcfcd8b Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:15:02 -0400 Subject: [PATCH 08/13] Return column type information as well as name --- internal/tiger/mcp/db_tools.go | 37 +++++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index 3b990557..94f9b650 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -40,19 +40,29 @@ func (DBExecuteQueryInput) Schema() *jsonschema.Schema { return schema } +// DBExecuteQueryColumn represents a column in the query result +type DBExecuteQueryColumn struct { + Name string `json:"name"` + Type string `json:"type"` +} + // DBExecuteQueryOutput represents output for tiger_db_execute_query type DBExecuteQueryOutput struct { - Columns []string `json:"columns"` - Rows [][]any `json:"rows"` - RowsAffected int64 `json:"rows_affected"` - ExecutionTime string `json:"execution_time"` + Columns []DBExecuteQueryColumn `json:"columns"` + Rows [][]any `json:"rows"` + RowsAffected int64 `json:"rows_affected"` + ExecutionTime string `json:"execution_time"` } func (DBExecuteQueryOutput) Schema() *jsonschema.Schema { schema := util.Must(jsonschema.For[DBExecuteQueryOutput](nil)) - schema.Properties["columns"].Description = "Column names from the query result" - schema.Properties["columns"].Examples = []any{[]string{"id", "name", "created_at"}} + schema.Properties["columns"].Description = "Column metadata from the query result including name and PostgreSQL type" + schema.Properties["columns"].Examples = []any{[]DBExecuteQueryColumn{ + {Name: "id", Type: "int4"}, + {Name: "name", Type: "text"}, + {Name: "created_at", Type: "timestamptz"}, + }} schema.Properties["rows"].Description = "Result rows as arrays of values. Empty for commands that don't return rows (INSERT, UPDATE, DELETE, etc.)" schema.Properties["rows"].Examples = []any{[][]any{{1, "alice", "2024-01-01"}, {2, "bob", "2024-01-02"}}} @@ -160,11 +170,20 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ } defer rows.Close() - // Get column names from field descriptions + // Get column metadata from field descriptions fieldDescriptions := rows.FieldDescriptions() - columns := make([]string, len(fieldDescriptions)) + columns := make([]DBExecuteQueryColumn, len(fieldDescriptions)) for i, fd := range fieldDescriptions { - columns[i] = string(fd.Name) + // Get the type name from the connection's type map + dataType, ok := conn.TypeMap().TypeForOID(fd.DataTypeOID) + typeName := "unknown" + if ok && dataType != nil { + typeName = dataType.Name + } + columns[i] = DBExecuteQueryColumn{ + Name: fd.Name, + Type: typeName, + } } // Collect all rows From 60ce598a8b0f1510fb0f37ad834c70300cf3d15d Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:21:12 -0400 Subject: [PATCH 09/13] Add role and pooled options to make 'tiger db connect' command --- internal/tiger/mcp/db_tools.go | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index 94f9b650..109a0b37 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -21,6 +21,8 @@ type DBExecuteQueryInput struct { ServiceID string `json:"service_id"` Query string `json:"query"` TimeoutSeconds int `json:"timeout_seconds,omitempty"` + Role string `json:"role,omitempty"` + Pooled bool `json:"pooled,omitempty"` } func (DBExecuteQueryInput) Schema() *jsonschema.Schema { @@ -37,6 +39,14 @@ func (DBExecuteQueryInput) Schema() *jsonschema.Schema { schema.Properties["timeout_seconds"].Default = util.Must(json.Marshal(30)) schema.Properties["timeout_seconds"].Examples = []any{10, 30, 60} + schema.Properties["role"].Description = "Database role/username to connect as" + schema.Properties["role"].Default = util.Must(json.Marshal("tsdbadmin")) + schema.Properties["role"].Examples = []any{"tsdbadmin", "readonly", "postgres"} + + schema.Properties["pooled"].Description = "Use connection pooling (if available for the service)" + schema.Properties["pooled"].Default = util.Must(json.Marshal(false)) + schema.Properties["pooled"].Examples = []any{false, true} + return schema } @@ -116,6 +126,8 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ zap.String("project_id", cfg.ProjectID), zap.String("service_id", input.ServiceID), zap.Duration("timeout", timeout), + zap.String("role", input.Role), + zap.Bool("pooled", input.Pooled), ) // Get service details to construct connection string @@ -141,10 +153,10 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ service := *serviceResp.JSON200 - // Build connection string with password (use direct connection, default role tsdbadmin) + // Build connection string with password connString, err := password.BuildConnectionString(service, password.ConnectionStringOptions{ - Pooled: false, - Role: "tsdbadmin", + Pooled: input.Pooled, + Role: input.Role, PasswordMode: password.PasswordRequired, // MCP always requires password }) if err != nil { From 9efbca36b0d4587ad176af03fb6ca4945b2f48a1 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:23:10 -0400 Subject: [PATCH 10/13] Update spec_mcp.md to reflect changes --- specs/spec_mcp.md | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/specs/spec_mcp.md b/specs/spec_mcp.md index cd24328e..4abe04b8 100644 --- a/specs/spec_mcp.md +++ b/specs/spec_mcp.md @@ -292,25 +292,37 @@ Test database connectivity. Execute a SQL query on a service database. **Parameters:** -- `service_id` (string, optional): Service ID (uses default if not provided) +- `service_id` (string, required): Service ID - `query` (string, required): SQL query to execute -- `timeout` (number, optional): Query timeout in seconds (default: 30) +- `timeout_seconds` (number, optional): Query timeout in seconds (default: 30) +- `role` (string, optional): Database role/username to connect as (default: tsdbadmin) +- `pooled` (boolean, optional): Use connection pooling (default: false) -**Returns:** Query results with rows, columns, and execution metadata. +**Returns:** Query results with rows, columns (including types), rows affected count, and execution metadata. **Example Response:** ```json { - "columns": ["id", "name", "created_at"], + "columns": [ + {"name": "id", "type": "int4"}, + {"name": "name", "type": "text"}, + {"name": "created_at", "type": "timestamptz"} + ], "rows": [ [1, "example", "2024-01-01T00:00:00Z"], [2, "test", "2024-01-02T00:00:00Z"] ], - "row_count": 2, - "execution_time_ms": 15 + "rows_affected": 2, + "execution_time": "15.2ms" } ``` +**Note:** +- `rows_affected` returns the number of rows returned for SELECT queries, and the number of rows modified for INSERT/UPDATE/DELETE queries +- `columns` includes both the column name and PostgreSQL data type for each column +- Empty `rows` array for commands that don't return rows (INSERT, UPDATE, DELETE, DDL commands) +- For parity with `tiger db connect` command, supports custom roles and connection pooling + ### High-Availability Management #### `tiger_ha_show` From 6b901dc0e331dca7482ca70d951e04a3c2fe1469 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Sun, 5 Oct 2025 15:33:39 -0400 Subject: [PATCH 11/13] Fix nil rows/columns bug --- internal/tiger/mcp/db_tools.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index 109a0b37..3c8a3223 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -58,8 +58,8 @@ type DBExecuteQueryColumn struct { // DBExecuteQueryOutput represents output for tiger_db_execute_query type DBExecuteQueryOutput struct { - Columns []DBExecuteQueryColumn `json:"columns"` - Rows [][]any `json:"rows"` + Columns []DBExecuteQueryColumn `json:"columns,omitempty"` + Rows [][]any `json:"rows,omitempty"` RowsAffected int64 `json:"rows_affected"` ExecutionTime string `json:"execution_time"` } @@ -184,18 +184,18 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ // Get column metadata from field descriptions fieldDescriptions := rows.FieldDescriptions() - columns := make([]DBExecuteQueryColumn, len(fieldDescriptions)) - for i, fd := range fieldDescriptions { + var columns []DBExecuteQueryColumn + for _, fd := range fieldDescriptions { // Get the type name from the connection's type map dataType, ok := conn.TypeMap().TypeForOID(fd.DataTypeOID) typeName := "unknown" if ok && dataType != nil { typeName = dataType.Name } - columns[i] = DBExecuteQueryColumn{ + columns = append(columns, DBExecuteQueryColumn{ Name: fd.Name, Type: typeName, - } + }) } // Collect all rows From 178a33d2c4c181e528c2fc7eb40620d4066cd7f7 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Mon, 6 Oct 2025 11:58:06 -0400 Subject: [PATCH 12/13] Support parameterized queries --- internal/tiger/mcp/db_tools.go | 6 +++++- specs/spec_mcp.md | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index 3c8a3223..dc498119 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -20,6 +20,7 @@ import ( type DBExecuteQueryInput struct { ServiceID string `json:"service_id"` Query string `json:"query"` + Parameters []any `json:"parameters,omitempty"` TimeoutSeconds int `json:"timeout_seconds,omitempty"` Role string `json:"role,omitempty"` Pooled bool `json:"pooled,omitempty"` @@ -34,6 +35,9 @@ func (DBExecuteQueryInput) Schema() *jsonschema.Schema { schema.Properties["query"].Description = "PostgreSQL query to execute" + schema.Properties["parameters"].Description = "Query parameters for parameterized queries. Values are substituted for $1, $2, etc. placeholders in the query." + schema.Properties["parameters"].Examples = []any{[]any{1, "alice"}, []any{"2024-01-01", 100}} + schema.Properties["timeout_seconds"].Description = "Query timeout in seconds" schema.Properties["timeout_seconds"].Minimum = util.Ptr(0.0) schema.Properties["timeout_seconds"].Default = util.Must(json.Marshal(30)) @@ -176,7 +180,7 @@ func (s *Server) handleDBExecuteQuery(ctx context.Context, req *mcp.CallToolRequ // Execute query and measure time startTime := time.Now() - rows, err := conn.Query(queryCtx, input.Query) + rows, err := conn.Query(queryCtx, input.Query, input.Parameters...) if err != nil { return nil, DBExecuteQueryOutput{}, fmt.Errorf("query execution failed: %w", err) } diff --git a/specs/spec_mcp.md b/specs/spec_mcp.md index 4abe04b8..993bf2ff 100644 --- a/specs/spec_mcp.md +++ b/specs/spec_mcp.md @@ -294,6 +294,7 @@ Execute a SQL query on a service database. **Parameters:** - `service_id` (string, required): Service ID - `query` (string, required): SQL query to execute +- `parameters` (array, optional): Query parameters for parameterized queries. Values are substituted for $1, $2, etc. placeholders in the query. - `timeout_seconds` (number, optional): Query timeout in seconds (default: 30) - `role` (string, optional): Database role/username to connect as (default: tsdbadmin) - `pooled` (boolean, optional): Use connection pooling (default: false) From 20462b37df3392ac769a9b8fbec982a8cf491855 Mon Sep 17 00:00:00 2001 From: Nathan Cochran Date: Mon, 6 Oct 2025 12:13:04 -0400 Subject: [PATCH 13/13] Update tool description to make it clear that multi-statement queries are not supported --- internal/tiger/mcp/db_tools.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/tiger/mcp/db_tools.go b/internal/tiger/mcp/db_tools.go index dc498119..c71d92f8 100644 --- a/internal/tiger/mcp/db_tools.go +++ b/internal/tiger/mcp/db_tools.go @@ -95,9 +95,9 @@ func (s *Server) registerDatabaseTools() { mcp.AddTool(s.mcpServer, &mcp.Tool{ Name: "db_execute_query", Title: "Execute SQL Query", - Description: `Execute a SQL query against a service database. + Description: `Execute a single SQL query against a service database. -This tool connects to a PostgreSQL database service in TigerData Cloud and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Perfect for data exploration, schema inspection, and database operations. +This tool connects to a PostgreSQL database service in TigerData Cloud and executes the provided SQL query, returning the results with column names, row data, and execution metadata. Multi-statement queries are not supported. WARNING: Use with caution - this tool can execute any SQL statement including INSERT, UPDATE, DELETE, and DDL commands. Always review queries before execution.`, InputSchema: DBExecuteQueryInput{}.Schema(),