experimental postgres query: use pgconn directly

experimental/postgres/cmd/connect.go

@@ -11,7 +11,6 @@ import (
 
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/log"
-	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
 )
@@ -45,11 +44,11 @@ type retryConfig struct {
 	MaxDelay     time.Duration
 }
 
-// connectFunc is a seam for unit tests: production wires pgx.ConnectConfig,
+// connectFunc is a seam for unit tests: production wires pgconn.ConnectConfig,
 // tests inject failures (DNS, auth, ctx-cancel mid-connect). We deliberately
-// do not wrap *pgx.Conn behind an interface for query execution; that surface
+// do not wrap *pgconn.PgConn behind an interface for query execution; that surface
 // is exercised by integration tests against real Lakebase endpoints.
-type connectFunc func(ctx context.Context, cfg *pgx.ConnConfig) (*pgx.Conn, error)
+type connectFunc func(ctx context.Context, cfg *pgconn.Config) (*pgconn.PgConn, error)
 
 // buildPgxConfig parses a DSN that includes the real host so pgx derives the
 // right TLSConfig and Fallbacks for sslmode=require. An empty host in the DSN
@@ -73,14 +72,19 @@ type connectFunc func(ctx context.Context, cfg *pgx.ConnConfig) (*pgx.Conn, erro
 // terminate the query within 5s, fall back to deadlining the connection.
 // Zero DeadlineDelay would race the cancel-request and could leave the
 // connection unusable.
-func buildPgxConfig(c connectConfig) (*pgx.ConnConfig, error) {
+func buildPgxConfig(c connectConfig) (*pgconn.Config, error) {
 	dsn := fmt.Sprintf("postgresql://%s/%s?sslmode=require",
 		net.JoinHostPort(c.Host, strconv.Itoa(c.Port)),
 		url.PathEscape(c.Database))
-	cfg, err := pgx.ParseConfig(dsn)
+	cfg, err := pgconn.ParseConfig(dsn)
 	if err != nil {
 		return nil, fmt.Errorf("parse pgx config: %w", err)
 	}
+	// pgx.ParseConfig strips its own options, which a PGSERVICE file may
+	// set; pgconn would send them to the server as unknown parameters.
+	for _, k := range []string{"statement_cache_capacity", "description_cache_capacity", "default_query_exec_mode"} {
+		delete(cfg.RuntimeParams, k)
+	}
 	cfg.User = c.Username
 	cfg.Password = c.Password
 	cfg.ConnectTimeout = c.ConnectTimeout
@@ -102,7 +106,7 @@ func buildPgxConfig(c connectConfig) (*pgx.ConnConfig, error) {
 //
 // MaxAttempts must be >= 1 (caller validates). 1 means a single attempt
 // with no retries.
-func connectWithRetry(ctx context.Context, cfg *pgx.ConnConfig, rc retryConfig, dial connectFunc) (*pgx.Conn, error) {
+func connectWithRetry(ctx context.Context, cfg *pgconn.Config, rc retryConfig, dial connectFunc) (*pgconn.PgConn, error) {
 	delay := rc.InitialDelay
 	var lastErr error
 

experimental/postgres/cmd/connect_test.go

@@ -4,11 +4,12 @@ import (
 	"context"
 	"errors"
 	"net"
+	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -102,11 +103,11 @@ func TestConnectWithRetry_RespectsMaxAttempts(t *testing.T) {
 	ctx := testCtx(t)
 	calls := 0
 	dialErr := &pgconn.PgError{Code: "08006"}
-	dial := func(ctx context.Context, cfg *pgx.ConnConfig) (*pgx.Conn, error) {
+	dial := func(ctx context.Context, cfg *pgconn.Config) (*pgconn.PgConn, error) {
 		calls++
 		return nil, dialErr
 	}
-	cfg := &pgx.ConnConfig{}
+	cfg := &pgconn.Config{}
 	rc := retryConfig{MaxAttempts: 3, InitialDelay: 0, MaxDelay: 0}
 
 	_, err := connectWithRetry(ctx, cfg, rc, dial)
@@ -118,11 +119,11 @@ func TestConnectWithRetry_StopsOnNonRetryable(t *testing.T) {
 	ctx := testCtx(t)
 	calls := 0
 	authErr := &pgconn.PgError{Code: "28P01"}
-	dial := func(ctx context.Context, cfg *pgx.ConnConfig) (*pgx.Conn, error) {
+	dial := func(ctx context.Context, cfg *pgconn.Config) (*pgconn.PgConn, error) {
 		calls++
 		return nil, authErr
 	}
-	cfg := &pgx.ConnConfig{}
+	cfg := &pgconn.Config{}
 	rc := retryConfig{MaxAttempts: 3, InitialDelay: 0}
 
 	_, err := connectWithRetry(ctx, cfg, rc, dial)
@@ -153,3 +154,15 @@ func TestBuildPgxConfig(t *testing.T) {
 	require.NotNil(t, cfg.TLSConfig, "TLSConfig must be set for sslmode=require")
 	assert.Equal(t, "host.example.com", cfg.TLSConfig.ServerName)
 }
+
+func TestBuildPgxConfigDropsPgxOnlyServiceOptions(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "pg_service.conf")
+	require.NoError(t, os.WriteFile(path, []byte("[svc]\nstatement_cache_capacity=0\ndescription_cache_capacity=0\ndefault_query_exec_mode=exec\napplication_name=app\n"), 0o600))
+	t.Setenv("PGSERVICEFILE", path)
+	t.Setenv("PGSERVICE", "svc")
+
+	cfg, err := buildPgxConfig(connectConfig{Host: "host.example.com", Port: 5432, Database: "db"})
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{"application_name": "app"}, cfg.RuntimeParams)
+}

experimental/postgres/cmd/execute.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 )
 
@@ -17,7 +16,8 @@ type rowSink interface {
 	// empty slice so the sink can lock in its rows-vs-command shape.
 	Begin(fields []pgconn.FieldDescription) error
 	// Row delivers one decoded row. Values aligns with the fields passed to
-	// Begin and uses pgx's Go type mapping (int64, float64, time.Time, ...).
+	// Begin and uses decodeTextValue's Go type mapping (int64, float64,
+	// time.Time, ...).
 	Row(values []any) error
 	// End is called once on successful completion.
 	End(ctx context.Context, commandTag string) error
@@ -31,48 +31,44 @@ type rowSink interface {
 
 // executeOne runs a single SQL statement and streams the result through sink.
 //
-// We pass QueryExecModeExec explicitly (not the pgx default
-// QueryExecModeCacheStatement) for two reasons:
-//
-//  1. Statement caching has no benefit for a one-shot CLI: the connection is
-//     closed at the end of the command, so the cached prepared statement
-//     never gets reused.
-//  2. Exec mode uses Postgres' extended-protocol "exec" path with text-format
-//     result columns. We still call rows.Values() (not RawValues) so all
-//     three sinks see uniform Go-typed input; jsonValue/textValue then map
-//     those types back to canonical strings for text/CSV and to JSON-typed
-//     values for JSON. The wire format being text means pgx's decode is
-//     cheap (text -> Go) rather than binary -> Go.
-//
-// QueryExecModeExec still uses extended protocol with a single statement and
-// no implicit transaction wrap, so transaction-disallowed DDL like
-// CREATE DATABASE works.
-func executeOne(ctx context.Context, conn *pgx.Conn, sql string, sink rowSink) error {
-	rows, err := conn.Query(ctx, sql, pgx.QueryExecModeExec)
-	if err != nil {
-		return fmt.Errorf("query failed: %w", err)
-	}
-	defer rows.Close()
+// ExecParams uses Postgres' extended protocol with a single unnamed statement
+// and text-format result columns. There is no statement caching (no benefit
+// for a one-shot CLI) and no implicit transaction wrap, so
+// transaction-disallowed DDL like CREATE DATABASE works. Columns are decoded
+// from their text form by decodeTextValue so all three sinks see uniform
+// Go-typed input.
+func executeOne(ctx context.Context, conn *pgconn.PgConn, sql string, sink rowSink) error {
+	rr := conn.ExecParams(ctx, sql, nil, nil, nil, nil)
+	// Close is idempotent; the explicit call below reports the command tag
+	// and any server error.
+	defer func() { _, _ = rr.Close() }()
 
-	if err := sink.Begin(rows.FieldDescriptions()); err != nil {
+	fields := rr.FieldDescriptions()
+	if err := sink.Begin(fields); err != nil {
 		return err
 	}
 
-	for rows.Next() {
-		values, err := rows.Values()
-		if err != nil {
-			sink.OnError(err)
-			return fmt.Errorf("decode row: %w", err)
+	for rr.NextRow() {
+		raw := rr.Values()
+		values := make([]any, len(raw))
+		for i, b := range raw {
+			v, err := decodeTextValue(fields[i].DataTypeOID, b)
+			if err != nil {
+				sink.OnError(err)
+				return fmt.Errorf("decode row: %w", err)
+			}
+			values[i] = v
 		}
 		if err := sink.Row(values); err != nil {
 			sink.OnError(err)
 			return err
 		}
 	}
-	if err := rows.Err(); err != nil {
+	tag, err := rr.Close()
+	if err != nil {
 		sink.OnError(err)
 		return fmt.Errorf("query failed: %w", err)
 	}
 
-	return sink.End(ctx, rows.CommandTag().String())
+	return sink.End(ctx, tag.String())
 }

experimental/postgres/cmd/query.go

@@ -10,7 +10,7 @@ import (
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/experimental/libs/sqlcli"
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/spf13/cobra"
 )
 
@@ -192,7 +192,7 @@ func runQuery(ctx context.Context, cmd *cobra.Command, args []string, f queryFla
 	// to stderr and degrades to a no-op in non-interactive terminals.
 	sp := cmdio.NewSpinner(signalCtx)
 	sp.Update("Connecting to " + resolved.DisplayName)
-	conn, err := connectWithRetry(signalCtx, pgxCfg, rc, pgx.ConnectConfig)
+	conn, err := connectWithRetry(signalCtx, pgxCfg, rc, pgconn.ConnectConfig)
 	sp.Close()
 	if err != nil {
 		return err

experimental/postgres/cmd/result.go

@@ -6,7 +6,6 @@ import (
 	"time"
 
 	"github.com/databricks/cli/experimental/libs/sqlcli"
-	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 )
 
@@ -32,7 +31,7 @@ func (r *unitResult) IsRowsProducing() bool {
 // as a thin wrapper that hands a bufferSink to executeOne, so error wrapping
 // and the rowSink contract stay in one place rather than parallel-evolving
 // across two query loops.
-func runUnitBuffered(ctx context.Context, conn *pgx.Conn, unit sqlcli.Input) (*unitResult, error) {
+func runUnitBuffered(ctx context.Context, conn *pgconn.PgConn, unit sqlcli.Input) (*unitResult, error) {
 	start := time.Now()
 	r := &unitResult{Source: unit.Source, SQL: unit.SQL}
 	sink := &bufferSink{result: r}
@@ -54,7 +53,7 @@ type bufferSink struct {
 }
 
 func (s *bufferSink) Begin(fields []pgconn.FieldDescription) error {
-	// pgx reuses the FieldDescription backing array across queries on the same
+	// pgconn reuses the FieldDescription backing array across queries on the same
 	// connection (pgConn.fieldDescriptions is a fixed-size buffer that's
 	// re-sliced per statement). Clone here so a buffered unit holds onto its
 	// own column descriptions; otherwise the multi-input renderers see every