diff --git a/experimental/postgres/cmd/connect.go b/experimental/postgres/cmd/connect.go index 79e7545b152..3c5bc2a3b96 100644 --- a/experimental/postgres/cmd/connect.go +++ b/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 diff --git a/experimental/postgres/cmd/connect_test.go b/experimental/postgres/cmd/connect_test.go index fd294ef2765..330ac52046d 100644 --- a/experimental/postgres/cmd/connect_test.go +++ b/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) +} diff --git a/experimental/postgres/cmd/execute.go b/experimental/postgres/cmd/execute.go index a4fe26ad820..8c8d36a2032 100644 --- a/experimental/postgres/cmd/execute.go +++ b/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()) } diff --git a/experimental/postgres/cmd/query.go b/experimental/postgres/cmd/query.go index 53684423f52..fa4c27bcba0 100644 --- a/experimental/postgres/cmd/query.go +++ b/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 diff --git a/experimental/postgres/cmd/result.go b/experimental/postgres/cmd/result.go index d5a0826d4c4..380f34a88da 100644 --- a/experimental/postgres/cmd/result.go +++ b/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 diff --git a/experimental/postgres/cmd/value.go b/experimental/postgres/cmd/value.go index 1578c7efecf..582917953f4 100644 --- a/experimental/postgres/cmd/value.go +++ b/experimental/postgres/cmd/value.go @@ -1,6 +1,7 @@ package postgrescmd import ( + "bytes" "encoding/base64" "encoding/hex" "encoding/json" @@ -9,6 +10,7 @@ import ( "math/big" "strconv" "time" + "unicode/utf8" "github.com/jackc/pgx/v5/pgtype" ) @@ -19,7 +21,63 @@ import ( // consumers. We render those as JSON strings to preserve the original digits. const safeIntegerBound = 1<<53 - 1 -// textValue renders a Go value (as decoded by pgx) to its canonical Postgres +// decodeTextValue converts a text-format result column into the Go value the +// sinks render. Integers, floats, booleans, bytea and date/time types map to +// native Go types; everything else (numeric, uuid, interval, arrays, json, +// ...) stays as the server's canonical text, which is already the +// representation we want to print. Date/time values use pgtype's own text +// parsers, so they decode exactly as they did with pgx.Conn; infinity stays +// as text. +func decodeTextValue(oid uint32, raw []byte) (any, error) { + if raw == nil { + return nil, nil + } + if oid == pgtype.ByteaOID { + hexPart, ok := bytes.CutPrefix(raw, []byte(`\x`)) + if !ok { + return nil, fmt.Errorf("unsupported bytea text format %q", raw) + } + b := make([]byte, hex.DecodedLen(len(hexPart))) + _, err := hex.Decode(b, hexPart) + return b, err + } + s := string(raw) + switch oid { + case pgtype.Int2OID: + v, err := strconv.ParseInt(s, 10, 16) + return int16(v), err + case pgtype.Int4OID: + v, err := strconv.ParseInt(s, 10, 32) + return int32(v), err + case pgtype.Int8OID: + return strconv.ParseInt(s, 10, 64) + case pgtype.Float4OID: + v, err := strconv.ParseFloat(s, 32) + return float32(v), err + case pgtype.Float8OID: + return strconv.ParseFloat(s, 64) + case pgtype.BoolOID: + return s == "t", nil + case pgtype.DateOID: + var d pgtype.Date + if err := d.Scan(s); err == nil && d.InfinityModifier == pgtype.Finite { + return d.Time, nil + } + case pgtype.TimestampOID: + var ts pgtype.Timestamp + if err := ts.Scan(s); err == nil && ts.InfinityModifier == pgtype.Finite { + return ts.Time, nil + } + case pgtype.TimestamptzOID: + var ts pgtype.Timestamptz + if err := ts.Scan(s); err == nil && ts.InfinityModifier == pgtype.Finite { + return ts.Time, nil + } + } + return s, nil +} + +// textValue renders a Go value (as decoded by decodeTextValue) to its canonical Postgres // text representation. Used by --output text and --output csv. // // NULL renders as the literal "NULL" so it lines up with the column rather @@ -74,7 +132,7 @@ func floatTextForm(f float64) string { return strconv.FormatFloat(f, 'g', -1, 64) } -// jsonValue renders a Go value (as decoded by pgx) to a JSON-encodable +// jsonValue renders a Go value (as decoded by decodeTextValue) to a JSON-encodable // representation. Returns a value the standard json.Marshal can handle // directly and the JSON shape we want; never returns Go values that would // silently lose information (e.g. NaN, oversized integers). @@ -101,7 +159,7 @@ func jsonValue(v any) any { case int16, int32: return x case int64: - // pgx decodes Postgres int8 to Go int64. Outside the IEEE-754 safe + // Postgres int8 decodes to Go int64. Outside the IEEE-754 safe // integer range we render as a string so JavaScript-style consumers // don't silently lose precision. if x > safeIntegerBound || x < -safeIntegerBound { @@ -113,11 +171,7 @@ func jsonValue(v any) any { case float64: return jsonFloat(x) case []byte: - // Postgres jsonb / json arrive as []byte holding raw JSON. Anything - // else we'd like to base64-encode. We can't tell them apart from the - // Go type alone; the sink calls jsonValueWithOID for oid-aware - // disambiguation. This bare path is the conservative fallback and - // treats unknown bytes as base64 (lossless and correct for bytea). + // bytea decodes to []byte; base64 keeps it lossless in JSON. return base64.StdEncoding.EncodeToString(x) case time.Time: return x.UTC().Format(time.RFC3339Nano) @@ -146,8 +200,9 @@ func jsonFloat(f float64) any { } // jsonValueWithOID applies oid-aware overrides on top of jsonValue. The two -// places this matters today are JSON/JSONB and bytea: both arrive from pgx as -// []byte but want different JSON shapes (raw JSON passthrough vs base64). +// places this matters today are JSON/JSONB and bytea: JSON arrives as its +// text and must pass through raw rather than as a JSON string; bytea arrives +// as []byte and is base64-encoded. func jsonValueWithOID(v any, oid uint32) any { if v == nil { return nil @@ -155,15 +210,11 @@ func jsonValueWithOID(v any, oid uint32) any { switch oid { case pgtype.JSONOID, pgtype.JSONBOID: - // pgx returns json/jsonb as already-decoded Go values when no codec - // is registered; with the default codec, they decode to map/slice/etc. - // In QueryExecModeExec text-mode, pgx returns the raw JSON bytes as - // string (since the wire is text). We accept both shapes. switch x := v.(type) { case []byte: - return json.RawMessage(x) + return rawJSON(x) case string: - return json.RawMessage(x) + return rawJSON([]byte(x)) } case pgtype.ByteaOID: if b, ok := v.([]byte); ok { @@ -173,3 +224,18 @@ func jsonValueWithOID(v any, oid uint32) any { return jsonValue(v) } + +// rawJSON passes JSON text through unchanged. Invalid UTF-8 (possible with a +// non-UTF8 client_encoding) would be copied into the output verbatim, so +// such values are decoded instead, which replaces it the way encoding/json +// does for strings. +func rawJSON(b []byte) any { + if utf8.Valid(b) { + return json.RawMessage(b) + } + var decoded any + if err := json.Unmarshal(b, &decoded); err == nil { + return decoded + } + return string(b) +} diff --git a/experimental/postgres/cmd/value_test.go b/experimental/postgres/cmd/value_test.go index d52edae90bc..f1a6bd27b80 100644 --- a/experimental/postgres/cmd/value_test.go +++ b/experimental/postgres/cmd/value_test.go @@ -8,6 +8,7 @@ import ( "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestJSONValue_PrimitiveTypes(t *testing.T) { @@ -93,3 +94,41 @@ func TestTextValue_FiniteFloat(t *testing.T) { assert.Equal(t, "3.14", textValue(float64(3.14))) assert.Equal(t, "0", textValue(float64(0))) } + +func TestDecodeTextValue_DateTime(t *testing.T) { + v, err := decodeTextValue(pgtype.DateOID, []byte("10000-01-02")) + require.NoError(t, err) + assert.Equal(t, time.Date(10000, 1, 2, 0, 0, 0, 0, time.UTC), v) + + v, err = decodeTextValue(pgtype.DateOID, []byte("0044-03-15 BC")) + require.NoError(t, err) + assert.Equal(t, time.Date(-43, 3, 15, 0, 0, 0, 0, time.UTC), v) + + v, err = decodeTextValue(pgtype.TimestampOID, []byte("2024-01-02 03:04:05.123456")) + require.NoError(t, err) + assert.Equal(t, time.Date(2024, 1, 2, 3, 4, 5, 123456000, time.UTC), v) + + v, err = decodeTextValue(pgtype.TimestamptzOID, []byte("2024-06-01 10:00:00.5+02")) + require.NoError(t, err) + assert.True(t, time.Date(2024, 6, 1, 8, 0, 0, 500000000, time.UTC).Equal(v.(time.Time))) + + v, err = decodeTextValue(pgtype.DateOID, []byte("infinity")) + require.NoError(t, err) + assert.Equal(t, "infinity", v) +} + +func TestDecodeTextValue_Bytea(t *testing.T) { + v, err := decodeTextValue(pgtype.ByteaOID, []byte(`\xdeadbeef`)) + require.NoError(t, err) + assert.Equal(t, []byte{0xde, 0xad, 0xbe, 0xef}, v) + + _, err = decodeTextValue(pgtype.ByteaOID, []byte(`\336`)) + assert.Error(t, err) +} + +func TestJSONValueWithOID_InvalidUTF8(t *testing.T) { + v := jsonValueWithOID("{\"x\":\"\xe9\"}", pgtype.JSONOID) + encoded, err := json.Marshal(v) + require.NoError(t, err) + assert.Equal(t, `{"x":"�"}`, string(encoded)) +}