From 5a1875faacbf88b28dff4763081ae76debf939c2 Mon Sep 17 00:00:00 2001 From: Taras <9948629+taraspos@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:30:35 +0100 Subject: [PATCH] mcp: use `QueryExecModeExec` for redshift access (#68250) --- lib/client/db/postgres/mcp/mcp.go | 25 ++++++++--- lib/client/db/postgres/mcp/mcp_test.go | 57 ++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/lib/client/db/postgres/mcp/mcp.go b/lib/client/db/postgres/mcp/mcp.go index 43e86128406..698d50f5603 100644 --- a/lib/client/db/postgres/mcp/mcp.go +++ b/lib/client/db/postgres/mcp/mcp.go @@ -28,6 +28,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/mark3labs/mcp-go/mcp" + "github.com/gravitational/teleport/api/types" dbmcp "github.com/gravitational/teleport/lib/client/db/mcp" clientmcp "github.com/gravitational/teleport/lib/client/mcp" "github.com/gravitational/teleport/lib/defaults" @@ -256,14 +257,28 @@ func buildConnConfig(db *dbmcp.Database) (*pgxpool.Config, error) { applicationNameParamName: applicationNameParamValue, } config.ConnConfig.TLSConfig = nil - // Use simple protocol to have a closer behavior to DB REPL and psql. - // - // This also avoids each query being prepared, binded and executed, reducing - // the amount of audit events per query executed. - config.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol + config.ConnConfig.DefaultQueryExecMode = queryExecModeForDatabase(db.DB) return config, nil } +func queryExecModeForDatabase(db types.Database) pgx.QueryExecMode { + switch db.GetType() { + case types.DatabaseTypeRedshift, types.DatabaseTypeRedshiftServerless: + // Redshift does not report standard_conforming_strings in the same way + // PostgreSQL does, which makes pgx reject simple protocol queries before + // sending them. QueryExecModeExec uses the extended query protocol without + // pgx statement caching, so it avoids the simple protocol sanitizer while + // staying suitable for ad-hoc SQL generated through MCP. + return pgx.QueryExecModeExec + default: + // Use simple protocol to have a closer behavior to DB REPL and psql. + // + // This also avoids each query being prepared, bound, and executed, reducing + // the amount of audit events per query executed. + return pgx.QueryExecModeSimpleProtocol + } +} + const ( // queryToolDatabaseParam is the name of the database URI param name from // query tool. diff --git a/lib/client/db/postgres/mcp/mcp_test.go b/lib/client/db/postgres/mcp/mcp_test.go index 698baeb4992..6e6ed161ff9 100644 --- a/lib/client/db/postgres/mcp/mcp_test.go +++ b/lib/client/db/postgres/mcp/mcp_test.go @@ -62,6 +62,48 @@ func TestFormatResult(t *testing.T) { } } +func TestQueryExecModeForDatabase(t *testing.T) { + for name, tc := range map[string]struct { + spec databaseSpec + want pgx.QueryExecMode + }{ + "postgres uses simple protocol": { + spec: databaseSpec{ + protocol: defaults.ProtocolPostgres, + }, + want: pgx.QueryExecModeSimpleProtocol, + }, + "redshift uses exec mode": { + spec: databaseSpec{ + protocol: defaults.ProtocolPostgres, + aws: types.AWS{ + Region: "us-east-1", + Redshift: types.Redshift{ClusterID: "redshift-cluster"}, + }, + }, + want: pgx.QueryExecModeExec, + }, + "redshift serverless uses exec mode": { + spec: databaseSpec{ + protocol: defaults.ProtocolPostgres, + aws: types.AWS{ + Region: "us-east-1", + RedshiftServerless: types.RedshiftServerless{ + WorkgroupName: "redshift-workgroup", + }, + }, + }, + want: pgx.QueryExecModeExec, + }, + } { + t.Run(name, func(t *testing.T) { + db, err := newTestDatabase(tc.spec) + require.NoError(t, err) + require.Equal(t, tc.want, queryExecModeForDatabase(db)) + }) + } +} + func TestFormatErrors(t *testing.T) { // Dummy listener that always drop connections. listener := listener.NewInMemoryListener() @@ -177,6 +219,21 @@ func TestFormatErrors(t *testing.T) { } } +type databaseSpec struct { + protocol string + aws types.AWS +} + +func newTestDatabase(spec databaseSpec) (types.Database, error) { + return types.NewDatabaseV3(types.Metadata{ + Name: "test-db", + }, types.DatabaseSpecV3{ + Protocol: spec.protocol, + URI: "localhost:5432", + AWS: spec.aws, + }) +} + func newMockRows(commandTag string, fields []string, rows [][]any) pgx.Rows { var fds []pgconn.FieldDescription for _, fieldName := range fields {