From 2b2faa3ba5d48906bddd88a8d9683cd574ac861c Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 16 Jul 2026 10:33:04 -0400 Subject: [PATCH 1/2] fix: compact Bind portal ownership --- server/conn.go | 189 ++++++-- server/conn_bind_test.go | 385 ++++++++++++++- server/conn_cursor.go | 21 +- server/conn_describe_test.go | 42 ++ server/conn_extended_query.go | 474 +++++++++++++++---- server/conn_pg_stat_activity.go | 12 +- server/conn_query_exec.go | 10 +- server/conn_querylog_lifecycle_test.go | 190 ++++++++ server/conn_results.go | 62 ++- server/conn_skip_until_sync_test.go | 63 +++ server/flightclient/flight_executor.go | 100 ++++ server/flightclient/interpolate_args_test.go | 92 +++- server/portal_lifecycle.go | 83 ++++ server/sqlcore/interfaces.go | 24 + server/sqlcore/sql_literal.go | 69 +++ 15 files changed, 1669 insertions(+), 147 deletions(-) create mode 100644 server/portal_lifecycle.go create mode 100644 server/sqlcore/sql_literal.go diff --git a/server/conn.go b/server/conn.go index a332b530..5cc750dd 100644 --- a/server/conn.go +++ b/server/conn.go @@ -94,57 +94,160 @@ type preparedStmt struct { cursorIsMove bool // FETCH is a MOVE: advance position without returning rows } +// portalState describes whether a portal can still be executed. A normal +// Execute (maxRows == 0) reaches a terminal state; existing maxRows behavior +// remains reusable until a future PortalSuspended implementation gives it +// lifecycle semantics. Keeping terminal shells matches PostgreSQL's error +// behavior for a repeated Execute while allowing the large Bind payload to be +// released now. +type portalState uint8 + +const ( + portalStateReady portalState = iota + portalStateDone + portalStateFailed +) + +// bindParam is a compact view into portal.bindBody. The PostgreSQL wire frame +// is capped below 2 GiB, so int32 offsets and lengths are sufficient. A +// negative length represents NULL; zero remains a distinct empty value. +type bindParam struct { + offset int32 + length int32 +} + +func (p bindParam) isNull() bool { return p.length < 0 } + type portal struct { - stmt *preparedStmt - paramValues [][]byte - paramFormats []int16 // 0=text, 1=binary for each parameter - resultFormats []int16 - described bool // true if Describe was called on this portal + stmt *preparedStmt + stmtName string // retained lightweight metadata for Close(S) after re-Parse + + // bindBody is the immutable Bind frame body allocated by wire.ReadMessage. + // Parameter data is never copied out of it: params stores offsets into this + // one backing allocation. Ownership transfers to this portal only after the + // entire Bind frame has been validated; releasePortalPayload drops it at a + // terminal lifecycle boundary. + bindBody []byte + params []bindParam + + paramFormats []int16 // 0=text, 1=binary; either one code or one per parameter + resultFormats []int16 // retained only until the portal reaches a terminal state + + // rowDescription is an already-encoded RowDescription body with this + // portal's result formats applied. It is intentionally the only + // result-metadata retained after terminal cleanup: it avoids a re-probe + // using released Bind values when a client later sends Describe(P). + rowDescription []byte + + described bool // true if Describe was called on this portal + state portalState + payloadReleased bool } -// decodeParams converts raw parameter bytes to Go values based on format codes. -// Returns (args, nil) on success, or (nil, error) for malformed binary data. -// On error, caller should send ErrorResponse with SQLSTATE 08P01. +func (p *portal) paramValue(index int) []byte { + if p == nil || index < 0 || index >= len(p.params) { + return nil + } + param := p.params[index] + if param.isNull() || param.offset < 0 || param.length < 0 { + return nil + } + start := int(param.offset) + end := start + int(param.length) + if start < 0 || end < start || end > len(p.bindBody) { + return nil + } + return p.bindBody[start:end] +} + +func (p *portal) paramTypeOID(index int) int32 { + if p != nil && p.stmt != nil && index >= 0 && index < len(p.stmt.paramTypes) { + return p.stmt.paramTypes[index] + } + return 0 +} + +func (p *portal) paramFormat(index int) int16 { + if p == nil { + return 0 + } + if len(p.paramFormats) == 1 { + return p.paramFormats[0] + } + if index >= 0 && index < len(p.paramFormats) { + return p.paramFormats[index] + } + return 0 +} + +// decodeParams converts compact raw parameter views to database/sql values. +// LocalExecutor and PinnedExecutor require this []any representation. Flight +// instead uses the SQLLiteralAppender methods below so its text path does not +// allocate a string/interface per value. func (p *portal) decodeParams() ([]interface{}, error) { - args := make([]interface{}, len(p.paramValues)) - for i, v := range p.paramValues { - if v == nil { + args := make([]interface{}, len(p.params)) + for i, param := range p.params { + if param.isNull() { args[i] = nil continue } + value := p.paramValue(i) + typeOID := p.paramTypeOID(i) + format := p.paramFormat(i) - // Get type OID for this parameter - typeOID := int32(0) // Unknown - if i < len(p.stmt.paramTypes) { - typeOID = p.stmt.paramTypes[i] + // Per PostgreSQL spec, UNKNOWN parameters are text even when a client + // supplied a binary format code. + if typeOID == 0 || format == 0 { + args[i] = string(value) + continue } - - // Get format code for this parameter - format := int16(0) // Default to text - if len(p.paramFormats) == 1 { - // Single format code applies to all parameters - format = p.paramFormats[0] - } else if i < len(p.paramFormats) { - // Per-parameter format codes - format = p.paramFormats[i] + decoded, err := decodeBinary(value, typeOID) + if err != nil { + return nil, fmt.Errorf("parameter %d: %w", i+1, err) } + args[i] = decoded + } + return args, nil +} - // CRITICAL: Per PostgreSQL spec, when type is unknown (OID 0), - // IGNORE binary format code and always treat as text. - // "Anything you have down as UNKNOWN, send as text." - if typeOID == 0 || format == 0 { - // Unknown type OR text format: treat as string - args[i] = string(v) - } else { - // Known type AND binary format: decode per type - val, err := decodeBinary(v, typeOID) - if err != nil { - return nil, fmt.Errorf("parameter %d: %w", i+1, err) - } - args[i] = val +// BindParameterCount and AppendBindParameterLiteral implement sqlcore's +// compact Flight interpolation capability. Raw text bytes are escaped directly +// into the final SQL builder, avoiding string(value) for every Bind parameter. +func (p *portal) BindParameterCount() int { return len(p.params) } + +func (p *portal) AppendBindParameterLiteral(dst *strings.Builder, index int) error { + if p == nil || index < 0 || index >= len(p.params) { + return fmt.Errorf("parameter %d is out of range", index+1) + } + if p.params[index].isNull() { + dst.WriteString("NULL") + return nil + } + value := p.paramValue(index) + if typeOID, format := p.paramTypeOID(index), p.paramFormat(index); typeOID == 0 || format == 0 { + sqlcore.AppendTextLiteralBytes(dst, value) + return nil + } + decoded, err := decodeBinary(value, p.paramTypeOID(index)) + if err != nil { + return fmt.Errorf("parameter %d: %w", index+1, err) + } + sqlcore.AppendSQLLiteral(dst, decoded) + return nil +} + +// validateBinaryParameters preserves the protocol-level 08P01 classification +// before a Flight interpolation error reaches generic executor handling. +func (p *portal) validateBinaryParameters() error { + for i, param := range p.params { + if param.isNull() || p.paramTypeOID(i) == 0 || p.paramFormat(i) == 0 { + continue + } + if _, err := decodeBinary(p.paramValue(i), p.paramTypeOID(i)); err != nil { + return fmt.Errorf("parameter %d: %w", i+1, err) } } - return args, nil + return nil } // Transaction status constants for PostgreSQL wire protocol @@ -1146,6 +1249,9 @@ func (c *clientConn) armIdleReadDeadline() { } func (c *clientConn) messageLoop() error { + // Every caller (standalone, control plane, and isolated worker) reaches + // this loop, so it is the single connection-close ownership boundary. + defer c.dropAllPortals("connection_close") atReadyBoundary := true for { if atReadyBoundary && c.drainRequested.Load() { @@ -1219,8 +1325,8 @@ func (c *clientConn) messageLoop() error { c.runExtendedQueryMessage(c.handleExecute, body) case wire.MsgSync: - // Extended query protocol - Sync: ends any skip-until-Sync error - // recovery, then reports readiness. + // Extended query protocol - Sync ends skip-until-Sync error recovery + // and reports readiness. It does not own portal cleanup. c.ignoreTillSync = false if err := c.writeReadyForQuery(c.txStatus); err != nil { return err @@ -1252,6 +1358,9 @@ func (c *clientConn) messageLoop() error { } func (c *clientConn) handleQuery(body []byte) (retErr error) { + // A simple Query destroys the unnamed extended-protocol portal before it + // executes, including an empty or failing query. + c.dropPortal("", "simple_query") query := string(bytes.TrimRight(body, "\x00")) query = strings.TrimSpace(query) diff --git a/server/conn_bind_test.go b/server/conn_bind_test.go index 7a38dec7..5d14958f 100644 --- a/server/conn_bind_test.go +++ b/server/conn_bind_test.go @@ -4,9 +4,20 @@ import ( "bufio" "bytes" "encoding/binary" + "io" + "strings" "testing" + + "github.com/posthog/duckgres/server/flightclient" ) +type bindTestValue struct { + null bool + data []byte +} + +var bindTestInterpolatedSQLSink string + // buildBindBody constructs a Bind message body (without the 'B' type byte and // length header): portal name, statement name, no format codes, one parameter // with the given declared length and actual data, no result format codes. @@ -38,6 +49,84 @@ func bindMessageBody(portal, stmt string, fields ...any) []byte { return buf.Bytes() } +func bindTestBody(portal, stmt string, paramFormats []int16, values []bindTestValue, resultFormats []int16) []byte { + var body bytes.Buffer + body.WriteString(portal) + body.WriteByte(0) + body.WriteString(stmt) + body.WriteByte(0) + _ = binary.Write(&body, binary.BigEndian, int16(len(paramFormats))) + for _, format := range paramFormats { + _ = binary.Write(&body, binary.BigEndian, format) + } + _ = binary.Write(&body, binary.BigEndian, int16(len(values))) + for _, value := range values { + if value.null { + _ = binary.Write(&body, binary.BigEndian, int32(-1)) + continue + } + _ = binary.Write(&body, binary.BigEndian, int32(len(value.data))) + body.Write(value.data) + } + _ = binary.Write(&body, binary.BigEndian, int16(len(resultFormats))) + for _, format := range resultFormats { + _ = binary.Write(&body, binary.BigEndian, format) + } + return body.Bytes() +} + +func executeTestBody(portal string) []byte { + return executeTestBodyWithMaxRows(portal, 0) +} + +func executeTestBodyWithMaxRows(portal string, maxRows int32) []byte { + var body bytes.Buffer + body.WriteString(portal) + body.WriteByte(0) + _ = binary.Write(&body, binary.BigEndian, maxRows) + return body.Bytes() +} + +func parseTestBody(stmt, query string, paramTypes []int32) []byte { + var body bytes.Buffer + body.WriteString(stmt) + body.WriteByte(0) + body.WriteString(query) + body.WriteByte(0) + _ = binary.Write(&body, binary.BigEndian, int16(len(paramTypes))) + for _, oid := range paramTypes { + _ = binary.Write(&body, binary.BigEndian, oid) + } + return body.Bytes() +} + +func bindPortalForTest(t *testing.T, c *clientConn, portal, stmt string, paramFormats []int16, values []bindTestValue, resultFormats []int16) *portal { + t.Helper() + c.handleBind(bindTestBody(portal, stmt, paramFormats, values, resultFormats)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush Bind: %v", err) + } + p := c.portals[portal] + if p == nil { + t.Fatalf("Bind did not install portal %q", portal) + } + return p +} + +func requireBindPayload(t *testing.T, p *portal) { + t.Helper() + if p == nil || p.payloadReleased || p.bindBody == nil { + t.Fatalf("portal did not retain its Bind frame: portal=%p released=%v body=%v", p, p != nil && p.payloadReleased, p != nil && p.bindBody != nil) + } +} + +func requireReleasedBindPayload(t *testing.T, p *portal) { + t.Helper() + if p == nil || !p.payloadReleased || p.bindBody != nil || p.params != nil || p.paramFormats != nil || p.resultFormats != nil { + t.Fatalf("portal Bind payload was not released: portal=%p released=%v body=%v params=%v paramFormats=%v resultFormats=%v", p, p != nil && p.payloadReleased, p != nil && p.bindBody != nil, p != nil && p.params != nil, p != nil && p.paramFormats != nil, p != nil && p.resultFormats != nil) + } +} + func newBindTestConn(out *bytes.Buffer) *clientConn { return &clientConn{ writer: bufio.NewWriter(out), @@ -135,8 +224,8 @@ func TestHandleBindAcceptsValidParam(t *testing.T) { if !ok { t.Fatal("expected portal to be created") } - if len(p.paramValues) != 1 || string(p.paramValues[0]) != "abc" { - t.Fatalf("unexpected param values: %v", p.paramValues) + if p.payloadReleased || p.bindBody == nil || len(p.params) != 1 || string(p.paramValue(0)) != "abc" { + t.Fatalf("unexpected compact parameter view: released=%v body=%v params=%v value=%q", p.payloadReleased, p.bindBody, p.params, p.paramValue(0)) } } @@ -155,7 +244,295 @@ func TestHandleBindValidMessageStillBinds(t *testing.T) { if !ok { t.Fatal("expected portal to be created") } - if len(p.paramValues) != 1 || p.paramValues[0] != nil { - t.Fatalf("expected one NULL param value, got %v", p.paramValues) + if len(p.params) != 1 || !p.params[0].isNull() || p.paramValue(0) != nil { + t.Fatalf("expected one NULL compact parameter, got %#v", p.params) + } +} + +// FETCH exposes the columns of an already-declared cursor, so its output +// cardinality is only known once the cursor schema is available during +// Describe or Execute. A client may nevertheless provide one result-format +// code per output column in Bind. +func TestHandleBindDefersCursorFetchResultFormatValidation(t *testing.T) { + var out bytes.Buffer + c := newBindTestConn(&out) + c.stmts["fetch"] = &preparedStmt{ + query: "FETCH 4 FROM cursor_name", + cursorOp: cursorOpFetch, + cursorName: "cursor_name", + } + + // No parameters, then two result format codes for the cursor's two + // columns. Their exact cardinality is validated once the schema is known. + c.handleBind(bindMessageBody("fetch-portal", "fetch", int16(0), int16(0), int16(2), int16(0), int16(1))) + _ = c.writer.Flush() + + if bytes.Contains(out.Bytes(), []byte("08P01")) { + t.Fatalf("cursor FETCH Bind must defer result format validation, got: %q", out.Bytes()) + } + if _, ok := c.portals["fetch-portal"]; !ok { + t.Fatal("expected cursor FETCH portal to be created") + } +} + +func TestHandleBindResultFormatCompatibilityForMutations(t *testing.T) { + value := []bindTestValue{{data: []byte("value")}} + tests := []struct { + name string + stmt *preparedStmt + formats []int16 + wantError bool + }{ + { + name: "explicit returning exact count", + stmt: &preparedStmt{ + query: "INSERT INTO t VALUES ($1) RETURNING first, second", + convertedQuery: "INSERT INTO t VALUES (?) RETURNING first, second", + numParams: 1, + }, + formats: []int16{0, 1}, + }, + { + name: "explicit returning mismatched count", + stmt: &preparedStmt{ + query: "INSERT INTO t VALUES ($1) RETURNING first, second", + convertedQuery: "INSERT INTO t VALUES (?) RETURNING first, second", + numParams: 1, + }, + formats: []int16{0, 1, 0}, + wantError: true, + }, + { + name: "wildcard returning deferred", + stmt: &preparedStmt{ + query: "INSERT INTO t VALUES ($1) RETURNING *", + convertedQuery: "INSERT INTO t VALUES (?) RETURNING *", + numParams: 1, + }, + formats: []int16{0, 1}, + }, + { + name: "writable cte deferred", + stmt: &preparedStmt{ + query: "WITH changed AS (UPDATE t SET value = $1 RETURNING id) SELECT id FROM changed", + convertedQuery: "SELECT id FROM changed", + numParams: 1, + statements: []string{"UPDATE t SET value = ?", "SELECT id FROM changed"}, + }, + formats: []int16{0, 1}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var out bytes.Buffer + executor := &lifecycleExecutor{} + c := newBindTestConn(&out) + c.executor = executor + c.stmts["mutation"] = tt.stmt + c.handleBind(bindTestBody("p", "mutation", nil, value, tt.formats)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush Bind: %v", err) + } + gotError := bytes.Contains(out.Bytes(), []byte("08P01")) + if gotError != tt.wantError { + t.Fatalf("Bind error = %v, want %v; output=%q", gotError, tt.wantError, out.Bytes()) + } + p := c.portals["p"] + if tt.wantError { + if p != nil { + t.Fatal("mismatched explicit RETURNING formats installed a portal") + } + } else { + if p == nil || p.state != portalStateReady { + t.Fatalf("compatible Bind portal = %#v, want ready portal", p) + } + requireBindPayload(t, p) + } + if got := executor.queryCalls.Load(); got != 0 { + t.Fatalf("Bind invoked Query %d times", got) + } + if got := executor.execCalls.Load(); got != 0 { + t.Fatalf("Bind invoked Exec %d times", got) + } + }) + } +} + +func TestHandleBindOwnershipCleanup(t *testing.T) { + var out bytes.Buffer + c := newBindTestConn(&out) + c.stmts["s1"] = &preparedStmt{query: "INSERT INTO t VALUES ($1)", convertedQuery: "INSERT INTO t VALUES (?)", numParams: 1} + + first := bindPortalForTest(t, c, "", "s1", nil, []bindTestValue{{data: []byte("first")}}, nil) + second := bindPortalForTest(t, c, "", "s1", nil, []bindTestValue{{data: []byte("second")}}, nil) + if c.portals[""] != second { + t.Fatal("unnamed Bind did not replace the previous portal") + } + requireReleasedBindPayload(t, first) + requireBindPayload(t, second) + + c.handleClose([]byte{'P', 0}) + if _, ok := c.portals[""]; ok { + t.Fatal("Close(P) did not remove unnamed portal") + } + requireReleasedBindPayload(t, second) + + named := bindPortalForTest(t, c, "named", "s1", nil, []bindTestValue{{data: []byte("value")}}, nil) + c.stmts["s1"] = &preparedStmt{query: "INSERT INTO t VALUES ($1)", convertedQuery: "INSERT INTO t VALUES (?)", numParams: 1} + c.handleClose([]byte{'S', 's', '1', 0}) + if _, ok := c.portals["named"]; ok { + t.Fatal("Close(S) did not remove portal after statement re-Parse") + } + requireReleasedBindPayload(t, named) +} + +func TestCompactBindPreservesNullEmptyTextAndBinary(t *testing.T) { + var out bytes.Buffer + c := newBindTestConn(&out) + c.stmts["values"] = &preparedStmt{ + query: "SELECT $1, $2, $3, $4", + convertedQuery: "SELECT ?, ?, ?, ?", + paramTypes: []int32{OidText, OidText, OidInt4, 0}, + numParams: 4, + } + p := bindPortalForTest(t, c, "values", "values", []int16{0, 0, 1, 1}, []bindTestValue{ + {null: true}, + {data: []byte{}}, + {data: []byte{0, 0, 0, 42}}, + {data: []byte("unknown-as-text")}, + }, nil) + args, err := p.decodeParams() + if err != nil { + t.Fatalf("decode params: %v", err) + } + if args[0] != nil { + t.Fatalf("NULL parameter decoded as %#v, want nil", args[0]) + } + if got, ok := args[1].(string); !ok || got != "" { + t.Fatalf("empty text parameter decoded as %#v, want empty string", args[1]) + } + if got, ok := args[2].(int32); !ok || got != 42 { + t.Fatalf("binary int4 decoded as %#v, want int32(42)", args[2]) + } + if got, ok := args[3].(string); !ok || got != "unknown-as-text" { + t.Fatalf("unknown OID with binary format decoded as %#v, want text", args[3]) + } + interpolated, err := flightclient.InterpolateBoundArgs("SELECT ?, ?, ?, ?", p) + if err != nil { + t.Fatalf("compact Flight interpolation: %v", err) + } + if want := "SELECT NULL, '', 42, 'unknown-as-text'"; interpolated != want { + t.Fatalf("compact Flight interpolation = %q, want %q", interpolated, want) + } +} + +func compactBindBenchmarkInput(count int, withPlaceholders bool) ([]byte, *preparedStmt) { + values := make([]bindTestValue, count) + for i := range values { + values[i] = bindTestValue{data: []byte("value")} + } + query := "INSERT INTO t VALUES ($1)" + if withPlaceholders { + query = "INSERT INTO t VALUES (" + strings.TrimSuffix(strings.Repeat("?,", count), ",") + ")" + } + return bindTestBody("", "bulk", nil, values, nil), &preparedStmt{ + query: query, + convertedQuery: query, + numParams: count, + } +} + +func TestCompactBind27000ParamAllocationsAreConstant(t *testing.T) { + measure := func(body []byte, stmt *preparedStmt) float64 { + return testing.AllocsPerRun(10, func() { + c := &clientConn{ + writer: bufio.NewWriter(io.Discard), + stmts: map[string]*preparedStmt{"bulk": stmt}, + portals: make(map[string]*portal), + } + c.handleBind(body) + if c.portals[""] == nil { + t.Fatal("Bind did not install a portal") + } + }) + } + smallBody, smallStmt := compactBindBenchmarkInput(1, false) + largeBody, largeStmt := compactBindBenchmarkInput(27_000, false) + small := measure(smallBody, smallStmt) + large := measure(largeBody, largeStmt) + if large > small+2 { + t.Fatalf("27,000 Bind parameters allocated %.0f objects versus %.0f for one parameter; want O(1) growth", large, small) + } +} + +func TestBindToFlightInterpolation27000TextAllocationsAreConstant(t *testing.T) { + measure := func(body []byte, stmt *preparedStmt) float64 { + return testing.AllocsPerRun(10, func() { + c := &clientConn{ + writer: bufio.NewWriter(io.Discard), + stmts: map[string]*preparedStmt{"bulk": stmt}, + portals: make(map[string]*portal), + } + c.handleBind(body) + p := c.portals[""] + if p == nil { + t.Fatal("Bind did not install a portal") + } + var err error + bindTestInterpolatedSQLSink, err = flightclient.InterpolateBoundArgs(stmt.convertedQuery, p) + if err != nil { + t.Fatal(err) + } + }) + } + smallBody, smallStmt := compactBindBenchmarkInput(1, true) + largeBody, largeStmt := compactBindBenchmarkInput(27_000, true) + small := measure(smallBody, smallStmt) + large := measure(largeBody, largeStmt) + if large > small+2 { + t.Fatalf("27,000 text Bind parameters allocated %.0f objects versus %.0f for one parameter; want O(1) growth", large, small) + } +} + +func BenchmarkHandleBind27000Params(b *testing.B) { + body, stmt := compactBindBenchmarkInput(27_000, false) + b.ReportAllocs() + b.SetBytes(int64(len(body))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + c := &clientConn{ + writer: bufio.NewWriter(io.Discard), + stmts: map[string]*preparedStmt{"bulk": stmt}, + portals: make(map[string]*portal), + } + c.handleBind(body) + if c.portals[""] == nil { + b.Fatal("Bind did not install a portal") + } + } +} + +func BenchmarkBindToFlightInterpolation27000Params(b *testing.B) { + body, stmt := compactBindBenchmarkInput(27_000, true) + b.ReportAllocs() + b.SetBytes(int64(len(body))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + c := &clientConn{ + writer: bufio.NewWriter(io.Discard), + stmts: map[string]*preparedStmt{"bulk": stmt}, + portals: make(map[string]*portal), + } + c.handleBind(body) + p := c.portals[""] + if p == nil { + b.Fatal("Bind did not install a portal") + } + interpolated, err := flightclient.InterpolateBoundArgs(stmt.convertedQuery, p) + if err != nil { + b.Fatal(err) + } + bindTestInterpolatedSQLSink = interpolated } } diff --git a/server/conn_cursor.go b/server/conn_cursor.go index 82e9c869..71fc2e92 100644 --- a/server/conn_cursor.go +++ b/server/conn_cursor.go @@ -197,14 +197,18 @@ func (c *clientConn) handlePgCursorsQuery(cursorName string) error { func (c *clientConn) handlePgCursorsQueryExtended(p *portal) { // Resolve cursor name: either from literal in query or from bind parameter cursorName := p.stmt.cursorName - if cursorName == "" && len(p.paramValues) > 0 && p.paramValues[0] != nil { - cursorName = string(p.paramValues[0]) + if cursorName == "" && p.BindParameterCount() > 0 && !p.params[0].isNull() { + cursorName = string(p.paramValue(0)) } _, exists := c.cursors[cursorName] + if !c.validPortalResultFormats(p, 1) { + return + } + p.rowDescription = pgCursorsRowDescriptionBody(p.resultFormats) if !p.stmt.described { - _ = c.sendPgCursorsRowDescriptionWithFormats(p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) } rowCount := 0 @@ -218,6 +222,10 @@ func (c *clientConn) handlePgCursorsQueryExtended(p *portal) { // sendPgCursorsRowDescription sends a RowDescription for a pg_cursors query result (single int4 column). func (c *clientConn) sendPgCursorsRowDescriptionWithFormats(formatCodes []int16) error { + return wire.WriteMessage(c.writer, wire.MsgRowDescription, pgCursorsRowDescriptionBody(formatCodes)) +} + +func pgCursorsRowDescriptionBody(formatCodes []int16) []byte { var buf bytes.Buffer _ = binary.Write(&buf, binary.BigEndian, int16(1)) // 1 column buf.WriteString("?column?") @@ -234,7 +242,7 @@ func (c *clientConn) sendPgCursorsRowDescriptionWithFormats(formatCodes []int16) format = formatCodes[0] } _ = binary.Write(&buf, binary.BigEndian, format) - return wire.WriteMessage(c.writer, wire.MsgRowDescription, buf.Bytes()) + return buf.Bytes() } // handleDeclareCursor handles DECLARE cursor in the Simple Query protocol. @@ -481,8 +489,11 @@ func (c *clientConn) handleFetchCursorExtended(p *portal) { } // Send RowDescription if Describe wasn't already called + if !c.cachePortalRowDescription(p, cursor.cols, cursor.colTypes) { + return + } if !p.described && len(cursor.cols) > 0 { - if err := c.sendRowDescriptionWithFormats(cursor.cols, cursor.colTypes, p.resultFormats); err != nil { + if err := c.writeCachedPortalRowDescription(p); err != nil { return } } diff --git a/server/conn_describe_test.go b/server/conn_describe_test.go index f7a35214..28099d0d 100644 --- a/server/conn_describe_test.go +++ b/server/conn_describe_test.go @@ -6,6 +6,8 @@ import ( "context" "errors" "testing" + + "github.com/posthog/duckgres/server/wire" ) type describeRecordingExecutor struct { @@ -164,3 +166,43 @@ func TestHandleDescribePortalPreservesExistingLimit(t *testing.T) { t.Fatalf("unexpected describe probe query: %q", got) } } + +func TestPortalLifecycleTerminalDescribeUsesCachedRowDescription(t *testing.T) { + exec := &describeRecordingExecutor{} + var out bytes.Buffer + c := &clientConn{ + executor: exec, + writer: bufio.NewWriter(&out), + portals: make(map[string]*portal), + cursors: make(map[string]*cursorState), + } + p := &portal{ + stmt: &preparedStmt{query: "SELECT $1", convertedQuery: "SELECT ?"}, + stmtName: "select", + bindBody: []byte("retained-bind-frame"), + params: []bindParam{{offset: 0, length: 1}}, + resultFormats: []int16{1}, + state: portalStateReady, + } + if !c.cachePortalRowDescription(p, []string{"value"}, []ColumnTyper{describeColumnType("VARCHAR")}) { + t.Fatal("cachePortalRowDescription failed") + } + c.finishPortal(p, portalStateDone, "terminal_success") + c.portals["p"] = p + requireReleasedBindPayload(t, p) + if p.stmt != nil { + t.Fatalf("terminal portal retained prepared statement %p", p.stmt) + } + + c.handleDescribe([]byte{'P', 'p', 0}) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush terminal Describe: %v", err) + } + if got := len(exec.queries); got != 0 { + t.Fatalf("terminal Describe queried %d times, want cached metadata only", got) + } + frames := scanWireFrames(t, out.Bytes()) + if len(frames) != 1 || frames[0].msgType != wire.MsgRowDescription { + t.Fatalf("terminal Describe frames = %q, want RowDescription", frameTypes(frames)) + } +} diff --git a/server/conn_extended_query.go b/server/conn_extended_query.go index 4d8aa2b5..12f796d4 100644 --- a/server/conn_extended_query.go +++ b/server/conn_extended_query.go @@ -5,12 +5,12 @@ import ( "context" "encoding/binary" "fmt" - "io" "strings" "time" pg_query "github.com/pganalyze/pg_query_go/v6" "github.com/posthog/duckgres/server/observe" + "github.com/posthog/duckgres/server/sqlcore" "github.com/posthog/duckgres/server/usersecrets" "github.com/posthog/duckgres/server/wire" "go.opentelemetry.io/otel/attribute" @@ -193,8 +193,8 @@ func (c *clientConn) handleParse(body []byte) { isIgnoredSet: result.IsIgnoredSet, isNoOp: result.IsNoOp, noOpTag: result.NoOpTag, - querySourceSet: result.QuerySourceSet, // SET duckgres.query_source (custom GUC) - querySourceShow: result.QuerySourceShow, // SHOW duckgres.query_source + querySourceSet: result.QuerySourceSet, // SET duckgres.query_source (custom GUC) + querySourceShow: result.QuerySourceShow, // SHOW duckgres.query_source statements: result.Statements, // Multi-statement rewrite (writable CTE) cleanupStatements: result.CleanupStatements, // Cleanup statements } @@ -380,6 +380,13 @@ func (c *clientConn) handleDescribe(body []byte) { c.sendError("ERROR", "34000", fmt.Sprintf("portal %q does not exist", name)) return } + // Terminal portals have released their Bind frame. Describe(P) must use + // the compact wire metadata captured while the portal was executable, + // never re-probe with an empty argument list. + if p.state != portalStateReady { + _ = c.writeCachedPortalRowDescription(p) + return + } // Handle cursor operations in portal Describe switch p.stmt.cursorOp { @@ -397,15 +404,29 @@ func (c *clientConn) handleDescribe(body []byte) { _ = wire.WriteNoData(c.writer) return } + if !c.cachePortalRowDescription(p, cols, colTypes) { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } p.described = true - _ = c.sendRowDescriptionWithFormats(cols, colTypes, p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) return case cursorOpPgCursorsQuery: - _ = c.sendPgCursorsRowDescriptionWithFormats(p.resultFormats) + if !c.validPortalResultFormats(p, 1) { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } + p.rowDescription = pgCursorsRowDescriptionBody(p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) p.described = true return case cursorOpPgStatActivity: - _ = c.sendPgStatActivityRowDescriptionWithFormats(p.resultFormats) + if !c.validPortalResultFormats(p, len(pgStatActivityColumns)) { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } + p.rowDescription = pgStatActivityRowDescriptionBody(p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) p.described = true return } @@ -434,19 +455,36 @@ func (c *clientConn) handleDescribe(body []byte) { // EXPLAIN [ANALYZE]: synthesize the single plan column without executing // (see the statement-Describe branch above). if isExplainStmt(p.stmt.query) { - _ = c.sendRowDescriptionWithFormats([]string{explainPlanColumn(p.stmt.query)}, []ColumnTyper{staticColumnType("VARCHAR")}, p.resultFormats) + if !c.cachePortalRowDescription(p, []string{explainPlanColumn(p.stmt.query)}, []ColumnTyper{staticColumnType("VARCHAR")}) { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } + _ = c.writeCachedPortalRowDescription(p) p.described = true p.stmt.described = true return } - // For SELECT, we need to describe the result columns - // We'll do a trial query with LIMIT 0 to get column info - args, err := p.decodeParams() - if err != nil { - // PostgreSQL returns 08P01 (protocol violation) for malformed binary data - c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) - return + // For SELECT, we need to describe the result columns. Local executors + // decode to database/sql arguments, while Flight receives the same + // compact literal appender used by Execute. + boundExecutor, usesBoundExecutor := c.executor.(sqlcore.BoundQueryExecutor) + var args []interface{} + if usesBoundExecutor { + if err := p.validateBinaryParameters(); err != nil { + c.finishPortal(p, portalStateFailed, "terminal_failure") + c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) + return + } + } else { + var err error + args, err = p.decodeParams() + if err != nil { + // PostgreSQL returns 08P01 (protocol violation) for malformed binary data + c.finishPortal(p, portalStateFailed, "terminal_failure") + c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) + return + } } // Try to get column info without fully executing expensive queries. @@ -456,7 +494,15 @@ func (c *clientConn) handleDescribe(body []byte) { describeQuery = describeQuery + " LIMIT 0" } - rows, err := c.executor.Query(describeQuery, args...) + var ( + rows RowSet + err error + ) + if usesBoundExecutor { + rows, err = boundExecutor.QueryWithBoundParams(describeQuery, p) + } else { + rows, err = c.executor.Query(describeQuery, args...) + } if err != nil { // Can't describe - send NoData _ = wire.WriteNoData(c.writer) @@ -479,9 +525,13 @@ func (c *clientConn) handleDescribe(body []byte) { // RowDescription. Without this, JDBC drivers that reuse named statements // (Bind/Execute without re-Describing) get an unexpected RowDescription // and desync their message queue. + if !c.cachePortalRowDescription(p, cols, colTypes) { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } p.described = true p.stmt.described = true - _ = c.sendRowDescriptionWithFormats(cols, colTypes, p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) default: c.sendError("ERROR", "08P01", "invalid Describe type") @@ -512,6 +562,25 @@ func (c *clientConn) handleExecute(body []byte) { c.sendError("ERROR", "34000", fmt.Sprintf("portal %q does not exist", portalName)) return } + if p.state != portalStateReady { + c.sendError("ERROR", "55000", fmt.Sprintf("portal %q is not executable", portalName)) + return + } + // A normal Execute is terminal. Preserve the existing maxRows behavior + // until PortalSuspended gets a dedicated implementation: limited executions + // keep their portal executable and retain its Bind payload. + terminalExecute := maxRows == 0 + errorsBeforeExecute := c.errorResponsesSent + defer func() { + if !terminalExecute { + return + } + if c.errorResponsesSent != errorsBeforeExecute { + c.finishPortal(p, portalStateFailed, "terminal_failure") + return + } + c.finishPortal(p, portalStateDone, "terminal_success") + }() // Redacted form for everything observable (pg_stat_activity, spans, // logs): CREATE SECRET option lists carry credential material. @@ -564,12 +633,23 @@ func (c *clientConn) handleExecute(body []byte) { ) defer span.End() - // Convert parameter values to interface{}, handling binary format - args, err := p.decodeParams() - if err != nil { - // PostgreSQL returns 08P01 (protocol violation) for malformed binary data - c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) - return + // Local database/sql executors need []any. Flight can instead append the + // compact portal views directly into its final SQL request buffer, avoiding + // string(value) and interface boxing for every text parameter. + boundExecutor, usesBoundExecutor := c.executor.(sqlcore.BoundQueryExecutor) + var args []interface{} + if usesBoundExecutor { + if err := p.validateBinaryParameters(); err != nil { + c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) + return + } + } else { + args, err = p.decodeParams() + if err != nil { + // PostgreSQL returns 08P01 (protocol violation) for malformed binary data + c.sendError("ERROR", "08P01", fmt.Sprintf("insufficient data left in message: %v", err)) + return + } } upperQuery := strings.ToUpper(strings.TrimSpace(p.stmt.query)) @@ -584,7 +664,7 @@ func (c *clientConn) handleExecute(body []byte) { return } - c.logger().Debug("Execute portal.", "portal", portalName, "params", len(args), "query", loggableQuery) + c.logger().Debug("Execute portal.", "portal", portalName, "params", p.BindParameterCount(), "query", loggableQuery) // duckgres.query_source custom GUC (SET / SHOW): intercepted session-side, // never forwarded to DuckDB. Determined by the transpiler during Parse. @@ -595,8 +675,11 @@ func (c *clientConn) handleExecute(body []byte) { return } if p.stmt.querySourceShow { + if !c.cachePortalRowDescription(p, []string{querySourceGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) { + return + } if !p.described { - _ = c.sendRowDescription([]string{querySourceGUCName}, []ColumnTyper{staticColumnType("VARCHAR")}) + _ = c.writeCachedPortalRowDescription(p) } _ = c.sendDataRowWithFormats([]interface{}{c.QuerySource()}, p.resultFormats, nil) _ = c.writeCommandComplete("SHOW") @@ -622,7 +705,7 @@ func (c *clientConn) handleExecute(body []byte) { // Handle multi-statement results (e.g., writable CTE rewrites) if len(p.stmt.statements) > 0 { c.logger().Debug("Execute multi-statement.", "statements", len(p.stmt.statements), "cleanup", len(p.stmt.cleanupStatements)) - c.executeMultiStatementExtended(p.stmt.statements, p.stmt.cleanupStatements, args, p.resultFormats, p.described) + c.executeMultiStatementExtended(p, p.stmt.statements, p.stmt.cleanupStatements, args, p.resultFormats, p.described) return } @@ -661,9 +744,18 @@ func (c *clientConn) handleExecute(body []byte) { // Non-result-returning query: use Exec with converted query runExec := func() (ExecResult, error) { - result, err := c.executor.Exec(convertedQuery, args...) + var result ExecResult + var err error + if usesBoundExecutor { + result, err = boundExecutor.ExecWithBoundParams(convertedQuery, p) + } else { + result, err = c.executor.Exec(convertedQuery, args...) + } if err != nil { if fallbackResult, handled, fallbackErr := c.execCompatibilityFallback(convertedQuery, err, func(fallbackQuery string) (ExecResult, error) { + if usesBoundExecutor { + return boundExecutor.ExecWithBoundParams(fallbackQuery, p) + } return c.executor.Exec(fallbackQuery, args...) }); handled { return fallbackResult, fallbackErr @@ -722,6 +814,9 @@ func (c *clientConn) handleExecute(body []byte) { // Result-returning query: use Query with converted query runQuery := func() (RowSet, error) { + if usesBoundExecutor { + return boundExecutor.QueryWithBoundParams(convertedQuery, p) + } return c.executor.Query(convertedQuery, args...) } @@ -773,6 +868,11 @@ func (c *clientConn) handleExecute(body []byte) { // Get column types for binary encoding colTypes, _ := rows.ColumnTypes() + // Cache serialized portal metadata even when Describe(S) made + // p.described true and Execute therefore suppresses RowDescription. + if !c.cachePortalRowDescription(p, cols, colTypes) { + return + } typeOIDs := make([]int32, len(cols)) for i, ct := range colTypes { typeOIDs[i] = getTypeInfo(ct).OID @@ -784,7 +884,7 @@ func (c *clientConn) handleExecute(body []byte) { // Skip if there are no columns - queries that return 0 columns (like // DDL accidentally routed here) don't need RowDescription. if !p.described && len(cols) > 0 { - if err := c.sendRowDescriptionWithFormats(cols, colTypes, p.resultFormats); err != nil { + if err := c.writeCachedPortalRowDescription(p); err != nil { return } } @@ -857,9 +957,19 @@ func (c *clientConn) handleClose(body []byte) { switch closeType { case 'S': - delete(c.stmts, name) + stmt := c.stmts[name] + // Closing a statement also closes all portals derived from that + // statement, including portals whose statement was subsequently + // re-Parsed under the same name. + c.dropPortalsForStatement(stmt, name, "close_statement") + if stmt != nil { + delete(c.stmts, name) + } case 'P': - delete(c.portals, name) + c.dropPortal(name, "close_portal") + default: + c.sendError("ERROR", "08P01", "invalid Close type") + return } _ = wire.WriteCloseComplete(c.writer) @@ -929,32 +1039,26 @@ func (c *clientConn) handleBind(body []byte) { // - Number of result format codes (int16) // - Result format codes (int16 each) - reader := bytes.NewReader(body) - - // Read portal name - portalName, err := readCString(reader) + reader := bindFrameReader{body: body} + portalName, err := reader.readCString() if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } - - // Read statement name - stmtName, err := readCString(reader) + stmtName, err := reader.readCString() if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } - // Look up prepared statement ps, ok := c.stmts[stmtName] if !ok { c.sendError("ERROR", "26000", fmt.Sprintf("prepared statement %q does not exist", stmtName)) return } - // Read parameter format codes - var numParamFormats int16 - if err := binary.Read(reader, binary.BigEndian, &numParamFormats); err != nil { + numParamFormats, err := reader.readInt16() + if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } @@ -962,17 +1066,25 @@ func (c *clientConn) handleBind(body []byte) { c.sendError("ERROR", "08P01", "invalid parameter format count in Bind message") return } - paramFormats := make([]int16, numParamFormats) - for i := int16(0); i < numParamFormats; i++ { - if err := binary.Read(reader, binary.BigEndian, ¶mFormats[i]); err != nil { - c.sendError("ERROR", "08P01", "invalid Bind message") - return + var paramFormats []int16 + if numParamFormats > 0 { + paramFormats = make([]int16, int(numParamFormats)) + for i := range paramFormats { + format, err := reader.readInt16() + if err != nil { + c.sendError("ERROR", "08P01", "invalid Bind message") + return + } + if format != 0 && format != 1 { + c.sendError("ERROR", "08P01", "invalid parameter format code in Bind message") + return + } + paramFormats[i] = format } } - // Read parameter values - var numParams int16 - if err := binary.Read(reader, binary.BigEndian, &numParams); err != nil { + numParams, err := reader.readInt16() + if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } @@ -980,39 +1092,42 @@ func (c *clientConn) handleBind(body []byte) { c.sendError("ERROR", "08P01", "invalid parameter count in Bind message") return } - paramValues := make([][]byte, numParams) - for i := int16(0); i < numParams; i++ { - var length int32 - if err := binary.Read(reader, binary.BigEndian, &length); err != nil { + if int(numParams) != ps.numParams { + c.sendError("ERROR", "08P01", fmt.Sprintf("bind message supplies %d parameters, but prepared statement %q requires %d", numParams, stmtName, ps.numParams)) + return + } + if len(paramFormats) != 0 && len(paramFormats) != 1 && len(paramFormats) != int(numParams) { + c.sendError("ERROR", "08P01", "invalid parameter format count in Bind message") + return + } + + params := make([]bindParam, int(numParams)) + for i := range params { + length, err := reader.readInt32() + if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } if length == -1 { - paramValues[i] = nil // NULL - } else if length < 0 { - // Only -1 (NULL) is a valid negative length. + params[i].length = -1 + continue + } + if length < 0 { + // Only -1 (NULL) is a valid negative length. This preserves the + // #717/#720 malformed-length behavior without allocating per value. c.sendError("ERROR", "08P01", "invalid parameter length in Bind message") return - } else { - // The length field is client-controlled; bound the allocation by - // the remaining bytes of the already-framed Bind message body — a - // parameter value can never legitimately exceed it. Without this - // check a client could reserve multi-GiB per parameter (#717). - if int64(length) > int64(reader.Len()) { - c.sendError("ERROR", "08P01", fmt.Sprintf("invalid Bind message: parameter %d length %d exceeds remaining message size %d", i+1, length, reader.Len())) - return - } - paramValues[i] = make([]byte, length) - if _, err := io.ReadFull(reader, paramValues[i]); err != nil { - c.sendError("ERROR", "08P01", "invalid Bind message") - return - } } + if int64(length) > int64(reader.remaining()) { + c.sendError("ERROR", "08P01", fmt.Sprintf("invalid Bind message: parameter %d length %d exceeds remaining message size %d", i+1, length, reader.remaining())) + return + } + params[i] = bindParam{offset: int32(reader.pos), length: length} + reader.pos += int(length) } - // Read result format codes - var numResultFormats int16 - if err := binary.Read(reader, binary.BigEndian, &numResultFormats); err != nil { + numResultFormats, err := reader.readInt16() + if err != nil { c.sendError("ERROR", "08P01", "invalid Bind message") return } @@ -1020,24 +1135,217 @@ func (c *clientConn) handleBind(body []byte) { c.sendError("ERROR", "08P01", "invalid result format count in Bind message") return } - resultFormats := make([]int16, numResultFormats) - for i := int16(0); i < numResultFormats; i++ { - if err := binary.Read(reader, binary.BigEndian, &resultFormats[i]); err != nil { - c.sendError("ERROR", "08P01", "invalid Bind message") - return + var resultFormats []int16 + if numResultFormats > 0 { + resultFormats = make([]int16, int(numResultFormats)) + for i := range resultFormats { + format, err := reader.readInt16() + if err != nil { + c.sendError("ERROR", "08P01", "invalid Bind message") + return + } + if format != 0 && format != 1 { + c.sendError("ERROR", "08P01", "invalid result format code in Bind message") + return + } + resultFormats[i] = format } } + if reader.remaining() != 0 { + c.sendError("ERROR", "08P01", "invalid Bind message: trailing data") + return + } + // Validate result-format vectors when their cardinality is already known. + // Unknown mutation output is deferred until actual metadata is available. + if !bindResultFormatCountSafeBeforeExecute(ps, len(resultFormats)) { + c.sendError("ERROR", "08P01", "invalid result format count in Bind message") + return + } - // Close existing portal with same name - delete(c.portals, portalName) + // PostgreSQL replaces only the unnamed portal. A duplicate named portal is + // an error and must leave the original portal/accounting untouched. + existing := c.portals[portalName] + if portalName != "" && existing != nil { + c.sendError("ERROR", "42P03", fmt.Sprintf("portal %q already exists", portalName)) + return + } - c.portals[portalName] = &portal{ + if portalName == "" && existing != nil { + c.dropPortal(portalName, "unnamed_rebind") + } + c.installPortal(portalName, &portal{ stmt: ps, - paramValues: paramValues, + stmtName: stmtName, + bindBody: body, + params: params, paramFormats: paramFormats, resultFormats: resultFormats, - described: ps.described, // Inherit from statement if Describe(S) was called - } + described: ps.described, // Inherit from statement Describe state. + state: portalStateReady, + }) _ = wire.WriteBindComplete(c.writer) } + +// bindResultFormatCountSafeBeforeExecute reports whether a Bind result-format +// vector can be accepted without risking execution of a mutation merely to +// learn its result-column count. PostgreSQL allows zero formats, one format, +// or exactly one per result column. +func bindResultFormatCountSafeBeforeExecute(ps *preparedStmt, formats int) bool { + if formats <= 1 { + return true + } + if ps == nil { + return false + } + // A non-MOVE FETCH returns the columns of a previously declared cursor. + // Its schema is not stored on the FETCH prepared statement, but becomes + // available during Describe or Execute before any cursor row is advanced. + // Defer an exact multi-code count check to that point. MOVE has no result + // columns and therefore cannot accept a multi-code vector. + if ps.cursorOp == cursorOpFetch { + return !ps.cursorIsMove + } + if !queryReturnsResults(ps.query) { + return false + } + // EXPLAIN always returns one plan column. In particular, EXPLAIN ANALYZE + // may execute a wrapped write, so its cardinality must be rejected here. + if isExplainStmt(ps.query) { + return false + } + if len(ps.statements) > 0 { + if !queryReturnsResults(ps.statements[len(ps.statements)-1]) { + return false + } + // TODO: validate deferred mutation result formats before execution + // without running the mutation merely to discover its schema. + return true + } + if isDMLReturning(ps.query) { + columns, known := explicitDMLReturningColumnCount(ps.query) + return !known || formats == columns + } + return true +} + +// explicitDMLReturningColumnCount returns a cardinality only when the original +// DML RETURNING list is a set of explicit targets. A wildcard depends on the +// table schema, so callers must not execute a mutation merely to resolve it. +func explicitDMLReturningColumnCount(query string) (int, bool) { + tree, err := pg_query.Parse(query) + if err != nil || len(tree.Stmts) != 1 || tree.Stmts[0] == nil || tree.Stmts[0].Stmt == nil { + return 0, false + } + + var targets []*pg_query.Node + switch stmt := tree.Stmts[0].Stmt.Node.(type) { + case *pg_query.Node_InsertStmt: + if stmt.InsertStmt != nil { + targets = stmt.InsertStmt.ReturningList + } + case *pg_query.Node_UpdateStmt: + if stmt.UpdateStmt != nil { + targets = stmt.UpdateStmt.ReturningList + } + case *pg_query.Node_DeleteStmt: + if stmt.DeleteStmt != nil { + targets = stmt.DeleteStmt.ReturningList + } + default: + return 0, false + } + if len(targets) == 0 { + return 0, false + } + for _, targetNode := range targets { + target := targetNode.GetResTarget() + if target == nil || resultTargetHasWildcard(target) { + return 0, false + } + } + return len(targets), true +} + +func resultTargetHasWildcard(target *pg_query.ResTarget) bool { + if resultTargetNodeHasWildcard(target.Val) { + return true + } + for _, indirection := range target.Indirection { + if resultTargetNodeHasWildcard(indirection) { + return true + } + } + return false +} + +// resultTargetNodeHasWildcard follows the AST shapes that expand one target +// into several output columns: `*`, `table.*`, and `(composite_expr).*`. +// Function calls such as count(*) remain a single explicit output target. +func resultTargetNodeHasWildcard(node *pg_query.Node) bool { + if node == nil { + return false + } + if node.GetAStar() != nil { + return true + } + if column := node.GetColumnRef(); column != nil { + for _, field := range column.Fields { + if field.GetAStar() != nil { + return true + } + } + } + if indirection := node.GetAIndirection(); indirection != nil { + if resultTargetNodeHasWildcard(indirection.Arg) { + return true + } + for _, part := range indirection.Indirection { + if resultTargetNodeHasWildcard(part) { + return true + } + } + } + return false +} + +// bindFrameReader parses an already framed Bind body without allocating a +// separate []byte for every parameter value. Its offsets are retained by the +// portal as compact descriptors into the immutable body. +type bindFrameReader struct { + body []byte + pos int +} + +func (r *bindFrameReader) remaining() int { return len(r.body) - r.pos } + +func (r *bindFrameReader) readCString() (string, error) { + start := r.pos + for r.pos < len(r.body) { + if r.body[r.pos] == 0 { + value := string(r.body[start:r.pos]) + r.pos++ + return value, nil + } + r.pos++ + } + return "", fmt.Errorf("unterminated cstring") +} + +func (r *bindFrameReader) readInt16() (int16, error) { + if r.remaining() < 2 { + return 0, fmt.Errorf("short int16") + } + value := int16(binary.BigEndian.Uint16(r.body[r.pos:])) + r.pos += 2 + return value, nil +} + +func (r *bindFrameReader) readInt32() (int32, error) { + if r.remaining() < 4 { + return 0, fmt.Errorf("short int32") + } + value := int32(binary.BigEndian.Uint32(r.body[r.pos:])) + r.pos += 4 + return value, nil +} diff --git a/server/conn_pg_stat_activity.go b/server/conn_pg_stat_activity.go index 78a6c631..8664ef25 100644 --- a/server/conn_pg_stat_activity.go +++ b/server/conn_pg_stat_activity.go @@ -98,8 +98,12 @@ func (c *clientConn) handlePgStatActivity() error { // handlePgStatActivityExtended handles SELECT FROM pg_stat_activity in the Extended Query protocol. func (c *clientConn) handlePgStatActivityExtended(p *portal) { + if !c.validPortalResultFormats(p, len(pgStatActivityColumns)) { + return + } + p.rowDescription = pgStatActivityRowDescriptionBody(p.resultFormats) if !p.stmt.described && !p.described { - _ = c.sendPgStatActivityRowDescriptionWithFormats(p.resultFormats) + _ = c.writeCachedPortalRowDescription(p) } conns := c.server.listConns() @@ -114,6 +118,10 @@ func (c *clientConn) handlePgStatActivityExtended(p *portal) { // sendPgStatActivityRowDescription sends a RowDescription for pg_stat_activity. func (c *clientConn) sendPgStatActivityRowDescriptionWithFormats(formatCodes []int16) error { + return wire.WriteMessage(c.writer, wire.MsgRowDescription, pgStatActivityRowDescriptionBody(formatCodes)) +} + +func pgStatActivityRowDescriptionBody(formatCodes []int16) []byte { var buf bytes.Buffer _ = binary.Write(&buf, binary.BigEndian, int16(len(pgStatActivityColumns))) for i, col := range pgStatActivityColumns { @@ -132,7 +140,7 @@ func (c *clientConn) sendPgStatActivityRowDescriptionWithFormats(formatCodes []i } _ = binary.Write(&buf, binary.BigEndian, format) } - return wire.WriteMessage(c.writer, wire.MsgRowDescription, buf.Bytes()) + return buf.Bytes() } // sendPgStatActivityDataRow sends a DataRow for a single connection in pg_stat_activity. diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index b42de724..1d656366 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -904,7 +904,7 @@ func (c *clientConn) executeCleanup(cleanup []string) { // executeMultiStatementExtended handles execution of multi-statement query rewrites // for the extended query protocol (Parse/Bind/Execute). // Unlike executeMultiStatement, this does NOT send ReadyForQuery (that's done by Sync). -func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup []string, args []interface{}, resultFormats []int16, described bool) { +func (c *clientConn) executeMultiStatementExtended(p *portal, statements []string, cleanup []string, args []interface{}, resultFormats []int16, described bool) { if len(statements) == 0 { _ = wire.WriteEmptyQueryResponse(c.writer) return @@ -930,7 +930,7 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmt) setupStart := time.Now() c.logQueryStarted(stmt) - result, err := c.executor.Exec(stmt, args...) + result, err := c.execPortal(p, stmt, args) var setupRows int64 if result != nil { setupRows, _ = result.RowsAffected() @@ -962,7 +962,7 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup if queryReturnsResults(finalStmt) { // Result-returning query: obtain cursor FIRST, cleanup SECOND, stream THIRD - rows, err := c.executor.Query(finalStmt, args...) + rows, err := c.queryPortal(p, finalStmt, args) if err != nil { finalErr = err c.logger().Error("Multi-stmt-ext final query error.", "query", finalStmt, "error", err) @@ -980,11 +980,11 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // is tracked by streamRowsToClientExtended; the deferred Finished // log uses 0 as an approximation (logQuery still records the // precise count via the structured channel). - c.streamRowsToClientExtended(rows, cmdType, resultFormats, described, finalStmt) + c.streamRowsToClientExtended(p, rows, cmdType, resultFormats, described, finalStmt) } else { // Non-result query (DML without RETURNING, DDL, etc.): execute then cleanup - result, err := c.executor.Exec(finalStmt, args...) + result, err := c.execPortal(p, finalStmt, args) if err != nil { finalErr = err c.logger().Error("Multi-stmt-ext final exec error.", "query", finalStmt, "error", err) diff --git a/server/conn_querylog_lifecycle_test.go b/server/conn_querylog_lifecycle_test.go index a2c6d695..2209b5e7 100644 --- a/server/conn_querylog_lifecycle_test.go +++ b/server/conn_querylog_lifecycle_test.go @@ -11,6 +11,9 @@ import ( "sync/atomic" "testing" "time" + + "github.com/posthog/duckgres/server/sqlcore" + "github.com/posthog/duckgres/server/wire" ) // newLifecycleClientConn builds a clientConn with both reader and writer @@ -76,6 +79,21 @@ func (e *lifecycleExecutor) ConnContext(_ context.Context) (RawConn, error) { func (e *lifecycleExecutor) PingContext(_ context.Context) error { return nil } func (e *lifecycleExecutor) Close() error { return nil } +type boundLifecycleExecutor struct { + lifecycleExecutor + boundCalls atomic.Int32 +} + +func (e *boundLifecycleExecutor) QueryWithBoundParams(_ string, _ sqlcore.SQLLiteralAppender) (RowSet, error) { + e.boundCalls.Add(1) + return e.queryRows, e.queryErr +} + +func (e *boundLifecycleExecutor) ExecWithBoundParams(_ string, _ sqlcore.SQLLiteralAppender) (ExecResult, error) { + e.boundCalls.Add(1) + return e.execResult, e.execErr +} + // emptyExecResult is an ExecResult that reports 0 rows affected — sufficient // for these tests, which assert lifecycle log presence not row-count math. type emptyExecResult struct{} @@ -343,3 +361,175 @@ func TestLifecyclePairFiresOnHandleExecuteQuery(t *testing.T) { c.handleExecute(body) assertLifecyclePair(t, buf, "handleExecute-Query") } + +func TestPortalLifecycleTerminalPortalsReleaseLargeStatementsAfterReparse(t *testing.T) { + const ( + largeQuerySize = 256 << 10 + largeParamOIDs = 27_000 + ) + for _, tc := range []struct { + name string + binary bool + }{ + {name: "success"}, + {name: "binary failure", binary: true}, + } { + t.Run(tc.name, func(t *testing.T) { + c, cleanup := newLifecycleClientConn(t) + defer cleanup() + var out bytes.Buffer + c.writer = bufio.NewWriter(&out) + c.passthrough = true + + var ( + executor QueryExecutor + bound *boundLifecycleExecutor + ) + if tc.binary { + bound = &boundLifecycleExecutor{} + executor = bound + } else { + executor = &lifecycleExecutor{execResult: emptyExecResult{}} + } + c.executor = executor + + paramTypes := make([]int32, largeParamOIDs) + paramFormats := []int16(nil) + value := bindTestValue{data: []byte("value")} + wantState := portalStateDone + if tc.binary { + paramTypes[0] = OidInt4 + paramFormats = []int16{1} + value = bindTestValue{data: []byte{0, 0}} + wantState = portalStateFailed + } + largeQuery := "INSERT INTO t VALUES ($1) /* " + strings.Repeat("x", largeQuerySize) + " */" + c.stmts["large"] = &preparedStmt{ + query: largeQuery, + convertedQuery: "INSERT INTO t VALUES (?) /* " + strings.Repeat("x", largeQuerySize) + " */", + paramTypes: paramTypes, + numParams: 1, + } + + p := bindPortalForTest(t, c, "terminal", "large", paramFormats, []bindTestValue{value}, nil) + out.Reset() + c.handleExecute(executeTestBody("terminal")) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush Execute: %v", err) + } + if p.state != wantState { + t.Fatalf("terminal state = %v, want %v", p.state, wantState) + } + requireReleasedBindPayload(t, p) + if p.stmt != nil { + t.Fatalf("terminal portal retained prepared statement %p", p.stmt) + } + if tc.binary { + if !bytes.Contains(out.Bytes(), []byte("08P01")) { + t.Fatalf("malformed binary Execute response = %q, want 08P01", out.Bytes()) + } + if got := bound.boundCalls.Load(); got != 0 { + t.Fatalf("malformed binary Bind reached Flight executor %d times", got) + } + } + + out.Reset() + c.handleExecute(executeTestBody("terminal")) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush repeated Execute: %v", err) + } + if !bytes.Contains(out.Bytes(), []byte("55000")) { + t.Fatalf("repeated terminal Execute response = %q, want 55000", out.Bytes()) + } + + c.handleClose([]byte{'S', 'm', 'i', 's', 's', 'i', 'n', 'g', 0}) + if c.portals["terminal"] != p { + t.Fatal("Close(S) for a missing statement dropped terminal portal") + } + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush missing Close(S): %v", err) + } + out.Reset() + + replacementQuery := "SELECT 1 /* " + strings.Repeat("r", largeQuerySize) + " */" + for i := 0; i < 3; i++ { + previous := c.stmts["large"] + c.handleParse(parseTestBody("large", replacementQuery, nil)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush re-Parse %d: %v", i, err) + } + frames := scanWireFrames(t, out.Bytes()) + if len(frames) != 1 || frames[0].msgType != wire.MsgParseComplete { + t.Fatalf("re-Parse %d frames = %q, want ParseComplete", i, frameTypes(frames)) + } + replacement := c.stmts["large"] + if replacement == nil || replacement == previous || replacement.query != replacementQuery { + t.Fatalf("re-Parse %d did not replace large statement: got %p, previous %p", i, replacement, previous) + } + out.Reset() + } + + c.handleClose([]byte{'S', 'l', 'a', 'r', 'g', 'e', 0}) + if _, ok := c.portals["terminal"]; ok { + t.Fatal("Close(S) did not remove terminal portal after re-Parse") + } + }) + } +} + +func TestPortalLifecycleMaxRowsKeepsLegacyReadyPortal(t *testing.T) { + c, cleanup := newLifecycleClientConn(t) + defer cleanup() + var out bytes.Buffer + c.writer = bufio.NewWriter(&out) + executor := &lifecycleExecutor{ + queryRows: &streamingRowSet{ + cols: []string{"value"}, + colTypers: []ColumnTyper{stringColumnTyper{}}, + rows: [][]any{{"one"}, {"two"}, {"three"}, {"four"}}, + }, + } + c.executor = executor + c.stmts["select"] = &preparedStmt{ + query: "SELECT $1", + convertedQuery: "SELECT ?", + numParams: 1, + } + p := bindPortalForTest(t, c, "limited", "select", nil, []bindTestValue{{data: []byte("value")}}, nil) + + for i := 0; i < 2; i++ { + out.Reset() + c.handleExecute(executeTestBodyWithMaxRows("limited", 1)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush limited Execute %d: %v", i, err) + } + if p.state != portalStateReady { + t.Fatalf("limited Execute %d state = %v, want ready", i, p.state) + } + if p.stmt == nil { + t.Fatalf("limited Execute %d released the statement", i) + } + requireBindPayload(t, p) + } + if got := executor.queryCalls.Load(); got != 2 { + t.Fatalf("limited Execute query calls = %d, want 2", got) + } +} + +func TestPortalLifecycleSimpleQueryDropsUnnamedPortal(t *testing.T) { + c, cleanup := newLifecycleClientConn(t) + defer cleanup() + var out bytes.Buffer + c.writer = bufio.NewWriter(&out) + c.executor = &lifecycleExecutor{execResult: emptyExecResult{}} + c.stmts["insert"] = &preparedStmt{query: "INSERT INTO t VALUES ($1)", convertedQuery: "INSERT INTO t VALUES (?)", numParams: 1} + p := bindPortalForTest(t, c, "", "insert", nil, []bindTestValue{{data: []byte("value")}}, nil) + + if err := c.handleQuery([]byte("UPDATE t SET value = 1\x00")); err != nil { + t.Fatalf("handle simple query: %v", err) + } + if _, ok := c.portals[""]; ok { + t.Fatal("simple Query did not remove unnamed portal") + } + requireReleasedBindPayload(t, p) +} diff --git a/server/conn_results.go b/server/conn_results.go index 71f35f19..1776b4d6 100644 --- a/server/conn_results.go +++ b/server/conn_results.go @@ -16,7 +16,7 @@ import ( // streamRowsToClientExtended sends result rows for the extended query protocol. // Unlike streamRowsToClient, this does NOT send ReadyForQuery, and supports // binary result formats and the described flag. -func (c *clientConn) streamRowsToClientExtended(rows RowSet, cmdType string, resultFormats []int16, described bool, query string) { +func (c *clientConn) streamRowsToClientExtended(p *portal, rows RowSet, cmdType string, resultFormats []int16, described bool, query string) { // Get column info cols, err := rows.Columns() if err != nil { @@ -33,6 +33,9 @@ func (c *clientConn) streamRowsToClientExtended(rows RowSet, cmdType string, res c.setTxError() return } + if !c.cachePortalRowDescription(p, cols, colTypes) { + return + } // Get type OIDs for binary encoding typeOIDs := make([]int32, len(cols)) @@ -42,7 +45,7 @@ func (c *clientConn) streamRowsToClientExtended(rows RowSet, cmdType string, res // Send RowDescription if Describe wasn't called before Execute if !described && len(cols) > 0 { - if err := c.sendRowDescriptionWithFormats(cols, colTypes, resultFormats); err != nil { + if err := c.writeCachedPortalRowDescription(p); err != nil { return } } @@ -174,6 +177,55 @@ func (c *clientConn) sendRowDescription(cols []string, colTypes []ColumnTyper) e // - single element: applies to all columns // - one per column: per-column format func (c *clientConn) sendRowDescriptionWithFormats(cols []string, colTypes []ColumnTyper, formatCodes []int16) error { + return c.writeRowDescriptionBody(c.rowDescriptionBody(cols, colTypes, formatCodes)) +} + +// cachePortalRowDescription stores exactly the wire metadata a terminal +// portal needs for a later Describe(P). It intentionally serializes the +// description instead of retaining RowSet/driver objects or the full Bind +// result-format slice. PostgreSQL allows zero formats, one format for every +// output column, or one format per actual output column; the count is first +// knowable here for arbitrary SQL. +func (c *clientConn) cachePortalRowDescription(p *portal, cols []string, colTypes []ColumnTyper) bool { + if p == nil { + return false + } + if len(p.resultFormats) != 0 && len(p.resultFormats) != 1 && len(p.resultFormats) != len(cols) { + c.sendError("ERROR", "08P01", "invalid result format count in Bind message") + return false + } + if len(cols) == 0 { + p.rowDescription = nil + return true + } + p.rowDescription = c.rowDescriptionBody(cols, colTypes, p.resultFormats) + return true +} + +func (c *clientConn) validPortalResultFormats(p *portal, columns int) bool { + if p == nil || len(p.resultFormats) == 0 || len(p.resultFormats) == 1 || len(p.resultFormats) == columns { + return true + } + c.sendError("ERROR", "08P01", "invalid result format count in Bind message") + return false +} + +func (c *clientConn) writeCachedPortalRowDescription(p *portal) error { + if p == nil || len(p.rowDescription) == 0 { + return wire.WriteNoData(c.writer) + } + return c.writeRowDescriptionBody(p.rowDescription) +} + +func (c *clientConn) writeRowDescriptionBody(body []byte) error { + if err := wire.WriteMessage(c.writer, wire.MsgRowDescription, body); err != nil { + c.markActiveQueryMetricsError(err) + return err + } + return nil +} + +func (c *clientConn) rowDescriptionBody(cols []string, colTypes []ColumnTyper, formatCodes []int16) []byte { var buf bytes.Buffer // Number of fields @@ -218,11 +270,7 @@ func (c *clientConn) sendRowDescriptionWithFormats(cols []string, colTypes []Col _ = binary.Write(&buf, binary.BigEndian, format) } - if err := wire.WriteMessage(c.writer, wire.MsgRowDescription, buf.Bytes()); err != nil { - c.markActiveQueryMetricsError(err) - return err - } - return nil + return buf.Bytes() } func (c *clientConn) mapTypeOIDWithColumnName(colName string, colType ColumnTyper) int32 { diff --git a/server/conn_skip_until_sync_test.go b/server/conn_skip_until_sync_test.go index b251eea8..b78cf65b 100644 --- a/server/conn_skip_until_sync_test.go +++ b/server/conn_skip_until_sync_test.go @@ -9,6 +9,7 @@ import ( "net" "strings" "testing" + "time" "github.com/posthog/duckgres/server/wire" ) @@ -235,6 +236,68 @@ func TestExtendedQueryErrorDiscardsPipelineUntilSync(t *testing.T) { } } +func TestPortalLifecycleSyncKeepsPortalUntilConnectionClose(t *testing.T) { + serverSide, clientSide := net.Pipe() + defer func() { _ = serverSide.Close() }() + defer func() { _ = clientSide.Close() }() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stmt := &preparedStmt{query: "SELECT 1", convertedQuery: "SELECT 1"} + p := &portal{ + stmt: stmt, + stmtName: "s", + bindBody: []byte("x"), + params: []bindParam{{offset: 0, length: 1}}, + state: portalStateReady, + } + c := &clientConn{ + server: &Server{activeQueries: make(map[BackendKey]context.CancelFunc)}, + conn: serverSide, + reader: bufio.NewReader(serverSide), + writer: bufio.NewWriter(serverSide), + ctx: ctx, + cancel: cancel, + txStatus: txStatusIdle, + stmts: map[string]*preparedStmt{"s": stmt}, + portals: map[string]*portal{"p": p}, + cursors: make(map[string]*cursorState), + } + done := make(chan error, 1) + go func() { done <- c.messageLoop() }() + + if err := wire.WriteMessage(clientSide, wire.MsgSync, nil); err != nil { + t.Fatalf("write Sync: %v", err) + } + msgType, body, err := wire.ReadMessage(clientSide) + if err != nil { + t.Fatalf("read Sync response: %v", err) + } + if msgType != wire.MsgReadyForQuery || len(body) != 1 || body[0] != txStatusIdle { + t.Fatalf("Sync response = type %q body %q, want ReadyForQuery(I)", msgType, body) + } + if got := c.portals["p"]; got != p { + t.Fatalf("Sync dropped portal: got %p, want %p", got, p) + } + requireBindPayload(t, p) + + if err := clientSide.Close(); err != nil { + t.Fatalf("close client: %v", err) + } + select { + case err := <-done: + if err != nil { + t.Fatalf("messageLoop returned %v", err) + } + case <-time.After(time.Second): + t.Fatal("messageLoop did not stop after connection close") + } + if _, ok := c.portals["p"]; ok { + t.Fatal("connection close did not remove portal") + } + requireReleasedBindPayload(t, p) +} + // TestAbortedTransactionStillAllowsExtendedRollback pins the recovery // trigger: it is the error event during extended-query processing, NOT // txStatus == 'E'. A connection whose transaction is aborted must still diff --git a/server/flightclient/flight_executor.go b/server/flightclient/flight_executor.go index 61e95d87..9e7b8e42 100644 --- a/server/flightclient/flight_executor.go +++ b/server/flightclient/flight_executor.go @@ -352,6 +352,42 @@ func (e *FlightExecutor) Exec(query string, args ...any) (sqlcore.ExecResult, er return e.ExecContext(context.Background(), query, args...) } +// QueryWithBoundParams executes a query using compact Bind parameters. Text +// values are appended straight into the final Flight SQL request buffer by the +// portal-owned appender, avoiding an intermediate string/interface per value. +func (e *FlightExecutor) QueryWithBoundParams(query string, params sqlcore.SQLLiteralAppender) (sqlcore.RowSet, error) { + // Match QueryContext's cheap terminal checks before building potentially + // large SQL from a compact Bind frame. + if e.dead.Load() { + return nil, ErrWorkerDead + } + if sqlcore.IsEmptyQuery(query) { + return &emptyRowSet{}, nil + } + interpolated, err := interpolateBoundArgs(query, params) + if err != nil { + return nil, err + } + return e.Query(interpolated) +} + +// ExecWithBoundParams is the Exec counterpart of QueryWithBoundParams. +func (e *FlightExecutor) ExecWithBoundParams(query string, params sqlcore.SQLLiteralAppender) (sqlcore.ExecResult, error) { + // Match ExecContext's cheap terminal checks before building potentially + // large SQL from a compact Bind frame. + if e.dead.Load() { + return nil, ErrWorkerDead + } + if sqlcore.IsEmptyQuery(query) { + return &flightExecResult{rowsAffected: 0}, nil + } + interpolated, err := interpolateBoundArgs(query, params) + if err != nil { + return nil, err + } + return e.Exec(interpolated) +} + func (e *FlightExecutor) ConnContext(ctx context.Context) (sqlcore.RawConn, error) { return nil, fmt.Errorf("ConnContext not supported in Flight mode (use batched INSERT for COPY FROM)") } @@ -1132,6 +1168,70 @@ func interpolateArgs(query string, args []any) string { return b.String() } +// InterpolateBoundArgs is interpolateArgs for a compact portal parameter +// source. It preserves the same placeholder and quoted/comment scanning rules +// while asking the source to append each literal directly to the one final SQL +// builder. That keeps a large text Bind from creating one transient string per +// parameter before the Flight RPC is made. +func InterpolateBoundArgs(query string, params sqlcore.SQLLiteralAppender) (string, error) { + if params == nil || params.BindParameterCount() == 0 { + return query, nil + } + + var b strings.Builder + b.Grow(len(query) + 16*params.BindParameterCount()) + nextPositional := 0 + + for i := 0; i < len(query); { + ch := query[i] + switch { + case ch == '\'' || ch == '"': + end := scanQuoted(query, i, ch) + b.WriteString(query[i:end]) + i = end + case ch == '-' && i+1 < len(query) && query[i+1] == '-': + end := indexOrEnd(query, i+2, "\n") + b.WriteString(query[i:end]) + i = end + case ch == '/' && i+1 < len(query) && query[i+1] == '*': + end := blockCommentEnd(query, i+2) + b.WriteString(query[i:end]) + i = end + case ch == '?': + if nextPositional < params.BindParameterCount() { + if err := params.AppendBindParameterLiteral(&b, nextPositional); err != nil { + return "", err + } + nextPositional++ + } else { + b.WriteByte(ch) + } + i++ + case ch == '$' && i+1 < len(query) && query[i+1] >= '1' && query[i+1] <= '9': + j := i + 1 + for j < len(query) && query[j] >= '0' && query[j] <= '9' { + j++ + } + if n, err := strconv.Atoi(query[i+1 : j]); err == nil && n >= 1 && n <= params.BindParameterCount() { + if err := params.AppendBindParameterLiteral(&b, n-1); err != nil { + return "", err + } + } else { + b.WriteString(query[i:j]) + } + i = j + default: + b.WriteByte(ch) + i++ + } + } + return b.String(), nil +} + +func interpolateBoundArgs(query string, params sqlcore.SQLLiteralAppender) (string, error) { + return InterpolateBoundArgs(query, params) +} + // scanQuoted returns the index just past a quoted region starting at start // (query[start] == quote), treating a doubled quote (” or "") as an escape. func scanQuoted(query string, start int, quote byte) int { diff --git a/server/flightclient/interpolate_args_test.go b/server/flightclient/interpolate_args_test.go index 531b7d08..48597048 100644 --- a/server/flightclient/interpolate_args_test.go +++ b/server/flightclient/interpolate_args_test.go @@ -1,6 +1,42 @@ package flightclient -import "testing" +import ( + "errors" + "strings" + "testing" + + "github.com/posthog/duckgres/server/sqlcore" +) + +type testLiteralAppender struct { + literals []string + errAt int +} + +func (a testLiteralAppender) BindParameterCount() int { return len(a.literals) } + +func (a testLiteralAppender) AppendBindParameterLiteral(dst *strings.Builder, index int) error { + if index == a.errAt { + return errors.New("malformed binary value") + } + dst.WriteString(a.literals[index]) + return nil +} + +var _ sqlcore.SQLLiteralAppender = testLiteralAppender{} + +type countingLiteralAppender struct { + calls int +} + +func (a *countingLiteralAppender) BindParameterCount() int { return 1 } + +func (a *countingLiteralAppender) AppendBindParameterLiteral(*strings.Builder, int) error { + a.calls++ + return errors.New("literal appender should not be called") +} + +var _ sqlcore.SQLLiteralAppender = (*countingLiteralAppender)(nil) func TestInterpolateArgs(t *testing.T) { tests := []struct { @@ -97,3 +133,57 @@ func TestInterpolateArgs(t *testing.T) { }) } } + +func TestInterpolateBoundArgs(t *testing.T) { + source := testLiteralAppender{literals: []string{"'o''brien'", "NULL", "42"}, errAt: -1} + got, err := interpolateBoundArgs("SELECT ?, $3, $1, '?' /* $2 */", source) + if err != nil { + t.Fatalf("interpolateBoundArgs() error = %v", err) + } + const want = "SELECT 'o''brien', 42, 'o''brien', '?' /* $2 */" + if got != want { + t.Fatalf("interpolateBoundArgs() = %q, want %q", got, want) + } + + _, err = interpolateBoundArgs("SELECT ?", testLiteralAppender{literals: []string{"ignored"}, errAt: 0}) + if err == nil || !strings.Contains(err.Error(), "malformed binary value") { + t.Fatalf("interpolateBoundArgs() error = %v, want literal decode error", err) + } +} + +func TestBoundParamsPreflightSkipsInterpolation(t *testing.T) { + t.Run("dead worker", func(t *testing.T) { + exec := &FlightExecutor{} + exec.MarkDead() + params := &countingLiteralAppender{} + + if _, err := exec.QueryWithBoundParams("SELECT ?", params); !errors.Is(err, ErrWorkerDead) { + t.Fatalf("QueryWithBoundParams() error = %v, want ErrWorkerDead", err) + } + if _, err := exec.ExecWithBoundParams("SELECT ?", params); !errors.Is(err, ErrWorkerDead) { + t.Fatalf("ExecWithBoundParams() error = %v, want ErrWorkerDead", err) + } + if params.calls != 0 { + t.Fatalf("literal appender calls = %d, want 0", params.calls) + } + }) + + t.Run("empty query", func(t *testing.T) { + exec := &FlightExecutor{} + params := &countingLiteralAppender{} + + rows, err := exec.QueryWithBoundParams("-- ping", params) + if err != nil { + t.Fatalf("QueryWithBoundParams() error = %v", err) + } + if rows.Next() { + t.Fatal("empty query returned a row") + } + if _, err := exec.ExecWithBoundParams("-- ping", params); err != nil { + t.Fatalf("ExecWithBoundParams() error = %v", err) + } + if params.calls != 0 { + t.Fatalf("literal appender calls = %d, want 0", params.calls) + } + }) +} diff --git a/server/portal_lifecycle.go b/server/portal_lifecycle.go new file mode 100644 index 00000000..f7307346 --- /dev/null +++ b/server/portal_lifecycle.go @@ -0,0 +1,83 @@ +package server + +import "github.com/posthog/duckgres/server/sqlcore" + +// installPortal transfers ownership of an already-validated Bind frame and +// compact metadata to the portal. +func (c *clientConn) installPortal(name string, p *portal) { + if c.portals == nil { + c.portals = make(map[string]*portal) + } + c.portals[name] = p +} + +// releasePortalPayload is idempotent. It releases the Bind frame and every +// Bind-derived slice, while retaining cached RowDescription metadata needed by +// terminal Describe(P). +func (c *clientConn) releasePortalPayload(p *portal, _ string) { + if p == nil || p.payloadReleased { + return + } + p.bindBody = nil + p.params = nil + p.paramFormats = nil + p.resultFormats = nil + p.payloadReleased = true +} + +// finishPortal records a terminal Execute result and releases the heavy Bind +// payload. A future PortalSuspended implementation can keep state Ready until +// the portal actually reaches a terminal response. +func (c *clientConn) finishPortal(p *portal, state portalState, reason string) { + if p == nil || p.state != portalStateReady { + return + } + p.state = state + c.releasePortalPayload(p, reason) + // Terminal Describe(P) replays rowDescription and Close(S) matches + // stmtName, so the shell must not pin the prepared statement. + p.stmt = nil +} + +func (c *clientConn) dropPortal(name, reason string) { + if c == nil || c.portals == nil { + return + } + p, ok := c.portals[name] + if !ok { + return + } + delete(c.portals, name) + c.releasePortalPayload(p, reason) +} + +func (c *clientConn) dropAllPortals(reason string) { + for name := range c.portals { + c.dropPortal(name, reason) + } +} + +func (c *clientConn) dropPortalsForStatement(stmt *preparedStmt, stmtName, reason string) { + for name, p := range c.portals { + if (stmt != nil && p.stmt == stmt) || p.stmtName == stmtName { + c.dropPortal(name, reason) + } + } +} + +// execPortal/queryPortal keep multi-statement rewrites on the same compact +// Flight path as ordinary extended execution. Local executors retain their +// existing database/sql []any behavior. +func (c *clientConn) execPortal(p *portal, query string, args []interface{}) (ExecResult, error) { + if executor, ok := c.executor.(sqlcore.BoundQueryExecutor); ok { + return executor.ExecWithBoundParams(query, p) + } + return c.executor.Exec(query, args...) +} + +func (c *clientConn) queryPortal(p *portal, query string, args []interface{}) (RowSet, error) { + if executor, ok := c.executor.(sqlcore.BoundQueryExecutor); ok { + return executor.QueryWithBoundParams(query, p) + } + return c.executor.Query(query, args...) +} diff --git a/server/sqlcore/interfaces.go b/server/sqlcore/interfaces.go index 9916fd52..882b3d5a 100644 --- a/server/sqlcore/interfaces.go +++ b/server/sqlcore/interfaces.go @@ -12,6 +12,7 @@ package sqlcore import ( "context" "io" + "strings" ) // ColumnTyper provides type name information for a database column. @@ -59,6 +60,29 @@ type QueryExecutor interface { LastProfilingOutput() string } +// SQLLiteralAppender exposes compact, already-validated bound parameters to an +// executor that must inline them into SQL. AppendBindParameterLiteral writes a +// single SQL literal directly into dst. Implementations must keep the backing +// parameter bytes alive until the executor returns. +// +// This deliberately avoids []any/string conversion for each text Bind value: +// remote Flight SQL execution needs one final SQL request buffer, but does not +// need thousands of intermediate Go strings. Local executors continue to use +// QueryExecutor's database/sql-compatible argument methods. +type SQLLiteralAppender interface { + BindParameterCount() int + AppendBindParameterLiteral(dst *strings.Builder, index int) error +} + +// BoundQueryExecutor is an optional QueryExecutor capability for backends +// (currently Flight SQL) that need a final interpolated SQL string. It accepts +// a compact literal appender so text Bind values can be written straight into +// that final buffer rather than first becoming one string/interface per value. +type BoundQueryExecutor interface { + QueryWithBoundParams(query string, params SQLLiteralAppender) (RowSet, error) + ExecWithBoundParams(query string, params SQLLiteralAppender) (ExecResult, error) +} + // CopyFromStdinExecutor is an optional capability implemented by executors // that need the CSV bytes shipped to a different filesystem than the one // the COPY FROM STDIN handler is running on. The remote (Flight) executor diff --git a/server/sqlcore/sql_literal.go b/server/sqlcore/sql_literal.go new file mode 100644 index 00000000..1eb34658 --- /dev/null +++ b/server/sqlcore/sql_literal.go @@ -0,0 +1,69 @@ +package sqlcore + +import ( + "encoding/hex" + "fmt" + "math/big" + "strconv" + "strings" + "time" +) + +// AppendTextLiteralBytes appends data as a DuckDB standard-conforming SQL +// string literal. It writes bytes directly so callers that retain a PostgreSQL +// Bind frame do not need to allocate an intermediate string for text values. +func AppendTextLiteralBytes(dst *strings.Builder, data []byte) { + dst.WriteByte('\'') + for _, ch := range data { + if ch == '\'' { + dst.WriteByte('\'') + } + dst.WriteByte(ch) + } + dst.WriteByte('\'') +} + +// AppendSQLLiteral appends the same literal representation historically used +// by Flight interpolation. It is shared with compact Bind interpolation so +// local and remote paths retain matching NULL, binary, numeric, and time +// semantics. +func AppendSQLLiteral(dst *strings.Builder, value any) { + if value == nil { + dst.WriteString("NULL") + return + } + switch v := value.(type) { + case string: + AppendTextLiteralBytes(dst, []byte(v)) + case []byte: + dst.WriteString(`'\x`) + dst.WriteString(hex.EncodeToString(v)) + dst.WriteString(`'::BLOB`) + case bool: + if v { + dst.WriteString("TRUE") + } else { + dst.WriteString("FALSE") + } + case time.Time: + AppendTextLiteralBytes(dst, []byte(v.Format("2006-01-02 15:04:05.999999"))) + case *big.Int: + dst.WriteString(v.String()) + case int: + dst.WriteString(strconv.Itoa(v)) + case int8: + dst.WriteString(strconv.FormatInt(int64(v), 10)) + case int16: + dst.WriteString(strconv.FormatInt(int64(v), 10)) + case int32: + dst.WriteString(strconv.FormatInt(int64(v), 10)) + case int64: + dst.WriteString(strconv.FormatInt(v, 10)) + case float32: + dst.WriteString(strconv.FormatFloat(float64(v), 'g', -1, 32)) + case float64: + dst.WriteString(strconv.FormatFloat(v, 'g', -1, 64)) + default: + AppendTextLiteralBytes(dst, []byte(fmt.Sprintf("%v", v))) + } +} From a308f36edd14b16521a1d6e9b1739cc82093a2e0 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 16 Jul 2026 11:27:59 -0400 Subject: [PATCH 2/2] fix: finalize failed portals on Execute errors --- server/conn_extended_query.go | 14 ++-- server/conn_querylog_lifecycle_test.go | 90 ++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 8 deletions(-) diff --git a/server/conn_extended_query.go b/server/conn_extended_query.go index 12f796d4..037fdce3 100644 --- a/server/conn_extended_query.go +++ b/server/conn_extended_query.go @@ -566,20 +566,18 @@ func (c *clientConn) handleExecute(body []byte) { c.sendError("ERROR", "55000", fmt.Sprintf("portal %q is not executable", portalName)) return } - // A normal Execute is terminal. Preserve the existing maxRows behavior - // until PortalSuspended gets a dedicated implementation: limited executions - // keep their portal executable and retain its Bind payload. - terminalExecute := maxRows == 0 + // Failed Execute calls are terminal. Preserve the existing successful + // maxRows behavior until PortalSuspended gets a dedicated implementation: + // only successful limited executions keep their portal and Bind payload. errorsBeforeExecute := c.errorResponsesSent defer func() { - if !terminalExecute { - return - } if c.errorResponsesSent != errorsBeforeExecute { c.finishPortal(p, portalStateFailed, "terminal_failure") return } - c.finishPortal(p, portalStateDone, "terminal_success") + if maxRows <= 0 { + c.finishPortal(p, portalStateDone, "terminal_success") + } }() // Redacted form for everything observable (pg_stat_activity, spans, diff --git a/server/conn_querylog_lifecycle_test.go b/server/conn_querylog_lifecycle_test.go index 2209b5e7..ce8bbc9a 100644 --- a/server/conn_querylog_lifecycle_test.go +++ b/server/conn_querylog_lifecycle_test.go @@ -516,6 +516,96 @@ func TestPortalLifecycleMaxRowsKeepsLegacyReadyPortal(t *testing.T) { } } +func TestPortalLifecycleFailedLimitedExecuteReleasesPortal(t *testing.T) { + c, cleanup := newLifecycleClientConn(t) + defer cleanup() + var out bytes.Buffer + c.writer = bufio.NewWriter(&out) + executor := &lifecycleExecutor{execErr: errors.New("worker is unavailable")} + c.executor = executor + c.stmts["insert"] = &preparedStmt{ + query: "INSERT INTO t VALUES ($1)", + convertedQuery: "INSERT INTO t VALUES (?)", + numParams: 1, + } + p := bindPortalForTest(t, c, "failed-limited", "insert", []int16{0}, []bindTestValue{{data: []byte("value")}}, []int16{0}) + if p.params == nil || p.paramFormats == nil || p.resultFormats == nil { + t.Fatal("test setup did not retain every Bind-owned descriptor") + } + + out.Reset() + c.handleExecute(executeTestBodyWithMaxRows("failed-limited", 1)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush limited Execute: %v", err) + } + if got := frameTypes(scanWireFrames(t, out.Bytes())); !strings.Contains(got, string(wire.MsgErrorResponse)) { + t.Fatalf("limited Execute response frames = %q, want ErrorResponse", got) + } + if p.state != portalStateFailed { + t.Errorf("limited failed Execute state = %v, want failed", p.state) + } + if p.stmt != nil { + t.Errorf("limited failed Execute retained prepared statement %p", p.stmt) + } + + out.Reset() + c.handleExecute(executeTestBodyWithMaxRows("failed-limited", 1)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush repeated limited Execute: %v", err) + } + if !bytes.Contains(out.Bytes(), []byte("55000")) { + t.Errorf("repeated failed limited Execute response = %q, want 55000", out.Bytes()) + } + if got := executor.execCalls.Load(); got != 1 { + t.Errorf("failed limited Execute calls = %d, want 1", got) + } + requireReleasedBindPayload(t, p) +} + +func TestPortalLifecycleNegativeMaxRowsIsUnlimitedAndTerminal(t *testing.T) { + c, cleanup := newLifecycleClientConn(t) + defer cleanup() + var out bytes.Buffer + c.writer = bufio.NewWriter(&out) + executor := &lifecycleExecutor{ + queryRows: &streamingRowSet{ + cols: []string{"value"}, + colTypers: []ColumnTyper{stringColumnTyper{}}, + rows: [][]any{{"one"}, {"two"}}, + }, + } + c.executor = executor + c.stmts["select"] = &preparedStmt{ + query: "SELECT $1", + convertedQuery: "SELECT ?", + numParams: 1, + } + p := bindPortalForTest(t, c, "negative", "select", []int16{0}, []bindTestValue{{data: []byte("value")}}, []int16{0}) + if p.params == nil || p.paramFormats == nil || p.resultFormats == nil { + t.Fatal("test setup did not retain every Bind-owned descriptor") + } + + out.Reset() + c.handleExecute(executeTestBodyWithMaxRows("negative", -1)) + if err := c.writer.Flush(); err != nil { + t.Fatalf("flush negative maxRows Execute: %v", err) + } + frames := scanWireFrames(t, out.Bytes()) + if got := strings.Count(frameTypes(frames), string(wire.MsgDataRow)); got != 2 { + t.Errorf("negative maxRows DataRow count = %d, want 2", got) + } + if got := frameTypes(frames); !strings.Contains(got, string(wire.MsgCommandComplete)) { + t.Errorf("negative maxRows response frames = %q, want CommandComplete", got) + } + if p.state != portalStateDone { + t.Errorf("negative maxRows Execute state = %v, want done", p.state) + } + if p.stmt != nil { + t.Errorf("negative maxRows Execute retained prepared statement %p", p.stmt) + } + requireReleasedBindPayload(t, p) +} + func TestPortalLifecycleSimpleQueryDropsUnnamedPortal(t *testing.T) { c, cleanup := newLifecycleClientConn(t) defer cleanup()