Merge branch 'fix/batch-cursor-field-descriptors-2626'

Author: jackc

batch_test.go

@@ -178,6 +178,57 @@ func TestConnSendBatchEmptyQuery(t *testing.T) {
 	})
 }
 
+// https://github.com/jackc/pgx/issues/2626
+func TestConnSendBatchCursorFetch(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		pgxtest.SkipCockroachDB(t, conn, "Server does not support cursors in implicit transactions")
+
+		tx, err := conn.Begin(ctx)
+		require.NoError(t, err)
+		defer tx.Rollback(ctx)
+
+		// The FETCH is prepared before the DECLARE has executed, so the cursor does not exist and the server
+		// describes the FETCH result as NoData. The actual fields are only known at execution time.
+		batch := &pgx.Batch{}
+		batch.Queue(`declare test_cursor cursor for select n, n::text as str from generate_series(1, 3) n`)
+		batch.Queue(`fetch all in test_cursor`)
+
+		br := tx.SendBatch(ctx, batch)
+
+		_, err = br.Exec()
+		require.NoError(t, err)
+
+		rows, err := br.Query()
+		require.NoError(t, err)
+
+		fds := rows.FieldDescriptions()
+		require.Len(t, fds, 2)
+		require.Equal(t, "n", fds[0].Name)
+		require.Equal(t, "str", fds[1].Name)
+
+		var ns []int32
+		var strs []string
+		for rows.Next() {
+			var n int32
+			var str string
+			require.NoError(t, rows.Scan(&n, &str))
+			ns = append(ns, n)
+			strs = append(strs, str)
+		}
+		require.NoError(t, rows.Err())
+		require.Equal(t, []int32{1, 2, 3}, ns)
+		require.Equal(t, []string{"1", "2", "3"}, strs)
+
+		err = br.Close()
+		require.NoError(t, err)
+	})
+}
+
 func TestConnSendBatchQueuedQuery(t *testing.T) {
 	t.Parallel()
 

pgconn/pgconn.go

@@ -1303,7 +1303,9 @@ func (pgConn *PgConn) ExecPrepared(ctx context.Context, stmtName string, paramVa
 //
 // This differs from [PgConn.ExecPrepared] in that it takes a [*StatementDescription] instead of the prepared statement name.
 // Because it has the [*StatementDescription] it can avoid the Describe Portal message that [PgConn.ExecPrepared] must send to get
-// the result column descriptions.
+// the result column descriptions. However, if the statement description has no fields then a Describe is still sent, as
+// an empty Fields may mean the results were not knowable at prepare time, e.g. a FETCH from a cursor that did not exist
+// yet.
 //
 // paramValues are the parameter values. It must be encoded in the format given by paramFormats.
 //
@@ -1364,7 +1366,10 @@ func (pgConn *PgConn) execExtendedPrefix(ctx context.Context, paramValues [][]by
 }
 
 func (pgConn *PgConn) execExtendedSuffix(result *ResultReader, statementDescription *StatementDescription, resultFormats []int16) {
-	if statementDescription == nil {
+	if statementDescription == nil || len(statementDescription.Fields) == 0 {
+		// The cached field descriptions are missing or empty. Empty field descriptions can occur when the statement's
+		// result set was not known at prepare time, e.g. a FETCH from a cursor that did not exist yet. Send a Describe
+		// so the server supplies the actual row description when the statement is executed.
 		pgConn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'P'})
 	}
 	pgConn.frontend.SendExecute(&pgproto3.Execute{})
@@ -2523,6 +2528,12 @@ func (p *Pipeline) SendQueryStatement(statementDescription *StatementDescription
 	}
 
 	p.conn.frontend.SendBind(&pgproto3.Bind{PreparedStatement: statementDescription.Name, ParameterFormatCodes: paramFormats, Parameters: paramValues, ResultFormatCodes: resultFormats})
+	if len(statementDescription.Fields) == 0 {
+		// The cached field descriptions are empty. This can occur when the statement's result set is
+		// not known at prepare time, e.g. a FETCH from a cursor that did not exist yet. Send a
+		// Describe so the server supplies the actual row description when the statement is executed.
+		p.conn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'P'})
+	}
 	p.conn.frontend.SendExecute(&pgproto3.Execute{})
 	p.state.PushBackRequestType(pipelineQueryStatement)
 	p.state.PushBackStatementData(statementDescription, resultFormats)
@@ -2729,12 +2740,36 @@ func (p *Pipeline) getResultsQueryStatement() (*ResultReader, error) {
 		return nil, err
 	}
 
+	sdFields := sd.Fields
+	if len(sdFields) == 0 {
+		// A Describe was sent for this statement (see SendQueryStatement). Read the server-provided
+		// row description which may include fields that were not known at prepare time.
+		msg, err := p.receiveMessage()
+		if err != nil {
+			return nil, err
+		}
+
+		switch msg := msg.(type) {
+		case *pgproto3.RowDescription:
+			sdFields = make([]FieldDescription, len(msg.Fields))
+			convertRowDescription(sdFields, msg)
+		case *pgproto3.NoData:
+			// Statement returns no rows.
+		case *pgproto3.ErrorResponse:
+			pgErr := ErrorResponseToPgError(msg)
+			p.state.HandleError(pgErr)
+			p.conn.resultReader.closed = true
+			return nil, pgErr
+		default:
+			return nil, p.handleUnexpectedMessage("QueryStatement RowDescription or NoData", msg)
+		}
+	}
+
 	msg, err := p.receiveMessage()
 	if err != nil {
 		return nil, err
 	}
 
-	sdFields := sd.Fields
 	fieldDescriptions := p.conn.getFieldDescriptionSlice(len(sdFields))
 	err = combineFieldDescriptionsAndResultFormats(fieldDescriptions, sdFields, resultFormats)
 	if err != nil {

pgconn/pgconn_test.go

@@ -1711,6 +1711,49 @@ func TestConnExecStatement(t *testing.T) {
 	ensureConnValid(t, pgConn)
 }
 
+// https://github.com/jackc/pgx/issues/2626
+func TestConnExecStatementCursorFetch(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgConn, err := pgconn.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer closeConn(t, pgConn)
+
+	if pgConn.ParameterStatus("crdb_version") != "" {
+		t.Skip("Server does not support cursors in implicit transactions")
+	}
+
+	// Prepare the FETCH before the cursor exists. The server describes the result as NoData so the statement
+	// description has no fields. The actual fields are only known at execution time.
+	sd, err := pgConn.Prepare(ctx, "ps_fetch", `fetch all in "exec_statement_cursor"`, nil)
+	require.NoError(t, err)
+	require.Empty(t, sd.Fields)
+
+	// DECLARE CURSOR requires an explicit transaction block.
+	_, err = pgConn.Exec(ctx, "begin").ReadAll()
+	require.NoError(t, err)
+
+	_, err = pgConn.Exec(ctx, `declare "exec_statement_cursor" cursor for select n, n::text from generate_series(1, 3) n`).ReadAll()
+	require.NoError(t, err)
+
+	result := pgConn.ExecStatement(ctx, sd, nil, nil, nil).Read()
+	require.NoError(t, result.Err)
+	require.Len(t, result.FieldDescriptions, 2)
+	require.Equal(t, uint32(pgtype.Int4OID), result.FieldDescriptions[0].DataTypeOID)
+	require.Equal(t, uint32(pgtype.TextOID), result.FieldDescriptions[1].DataTypeOID)
+	require.Len(t, result.Rows, 3)
+	require.Equal(t, "1", string(result.Rows[0][0]))
+	require.Equal(t, "3", string(result.Rows[2][1]))
+
+	_, err = pgConn.Exec(ctx, "rollback").ReadAll()
+	require.NoError(t, err)
+
+	ensureConnValid(t, pgConn)
+}
+
 type byteCounterConn struct {
 	conn         net.Conn
 	bytesRead    int
@@ -3682,6 +3725,71 @@ func TestPipelineQueryStatementBindError(t *testing.T) {
 	ensureConnValid(t, pgConn)
 }
 
+// https://github.com/jackc/pgx/issues/2626
+func TestPipelineQueryStatementCursorFetch(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgConn, err := pgconn.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer closeConn(t, pgConn)
+
+	if pgConn.ParameterStatus("crdb_version") != "" {
+		t.Skip("Server does not support cursors in implicit transactions")
+	}
+
+	// Prepare the FETCH before the cursor exists. The server describes the result as NoData so the statement
+	// description has no fields. The actual fields are only known at execution time.
+	sd, err := pgConn.Prepare(ctx, "ps_fetch", `fetch all in "pipeline_cursor"`, nil)
+	require.NoError(t, err)
+	require.Empty(t, sd.Fields)
+
+	// DECLARE CURSOR requires an explicit transaction block.
+	_, err = pgConn.Exec(ctx, "begin").ReadAll()
+	require.NoError(t, err)
+
+	pipeline := pgConn.StartPipeline(ctx)
+	pipeline.SendQueryParams(`declare "pipeline_cursor" cursor for select n, n::text from generate_series(1, 3) n`, nil, nil, nil, nil)
+	pipeline.SendQueryStatement(sd, nil, nil, nil)
+	err = pipeline.Sync()
+	require.NoError(t, err)
+
+	results, err := pipeline.GetResults()
+	require.NoError(t, err)
+	rr, ok := results.(*pgconn.ResultReader)
+	require.Truef(t, ok, "expected ResultReader, got: %#v", results)
+	readResult := rr.Read()
+	require.NoError(t, readResult.Err)
+
+	results, err = pipeline.GetResults()
+	require.NoError(t, err)
+	rr, ok = results.(*pgconn.ResultReader)
+	require.Truef(t, ok, "expected ResultReader, got: %#v", results)
+	readResult = rr.Read()
+	require.NoError(t, readResult.Err)
+	require.Len(t, readResult.FieldDescriptions, 2)
+	require.Equal(t, uint32(pgtype.Int4OID), readResult.FieldDescriptions[0].DataTypeOID)
+	require.Equal(t, uint32(pgtype.TextOID), readResult.FieldDescriptions[1].DataTypeOID)
+	require.Len(t, readResult.Rows, 3)
+	require.Equal(t, "1", string(readResult.Rows[0][0]))
+	require.Equal(t, "3", string(readResult.Rows[2][1]))
+
+	results, err = pipeline.GetResults()
+	require.NoError(t, err)
+	_, ok = results.(*pgconn.PipelineSync)
+	require.Truef(t, ok, "expected PipelineSync, got: %#v", results)
+
+	err = pipeline.Close()
+	require.NoError(t, err)
+
+	_, err = pgConn.Exec(ctx, "rollback").ReadAll()
+	require.NoError(t, err)
+
+	ensureConnValid(t, pgConn)
+}
+
 func TestPipelineGetResultsNilResultsOnError(t *testing.T) {
 	t.Parallel()
 

query_test.go

@@ -72,6 +72,53 @@ func TestConnQueryRowsFieldDescriptionsBeforeNext(t *testing.T) {
 	assert.Equal(t, "msg", rows.FieldDescriptions()[0].Name)
 }
 
+// https://github.com/jackc/pgx/issues/2626
+func TestConnQueryPreparedCursorFetch(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	pgxtest.SkipCockroachDB(t, conn, "Server does not support cursors in implicit transactions")
+
+	// Prepare the FETCH before the cursor exists. The server describes the result as NoData so the statement
+	// description has no fields. The actual fields are only known at execution time.
+	sd, err := conn.Prepare(ctx, "fetch_cursor", `fetch all in "query_cursor"`)
+	require.NoError(t, err)
+	require.Empty(t, sd.Fields)
+
+	tx, err := conn.Begin(ctx)
+	require.NoError(t, err)
+	defer tx.Rollback(ctx)
+
+	_, err = tx.Exec(ctx, `declare "query_cursor" cursor for select n, n::text as str from generate_series(1, 3) n`)
+	require.NoError(t, err)
+
+	rows, err := tx.Query(ctx, "fetch_cursor")
+	require.NoError(t, err)
+
+	fds := rows.FieldDescriptions()
+	require.Len(t, fds, 2)
+	require.Equal(t, "n", fds[0].Name)
+	require.Equal(t, "str", fds[1].Name)
+
+	var ns []int32
+	var strs []string
+	for rows.Next() {
+		var n int32
+		var str string
+		require.NoError(t, rows.Scan(&n, &str))
+		ns = append(ns, n)
+		strs = append(strs, str)
+	}
+	require.NoError(t, rows.Err())
+	require.Equal(t, []int32{1, 2, 3}, ns)
+	require.Equal(t, []string{"1", "2", "3"}, strs)
+}
+
 func TestConnQueryWithoutResultSetCommandTag(t *testing.T) {
 	t.Parallel()