Skip to content

Commit c0aa573

Browse files
rimadeodharrimadeodhar
committed
pgwire: Add support for cursors enclosed in quotes
In CockroachDB and Postgres, it is possible to declare cursors with special identifiers enclosed within double quotes, for e.g. "1-2-3". Currently, we store the name as an unescaped string which causes problems during the pgwire DESCRIBE step for looking up the cursor. We should be storing using the tree.Name datatype for the cursor name while storing and looking up cursors. This PR updates the code to start using tree.Name instead of raw strings for handling cursor names. This fixes the issue where the pgwire DESCRIBE step fails while attempting to look up cursors with names containing special characters. Resolves #84261 Release note (bug fix): The pgwire DESCRIBE step no longer fails with an error while attempting to look up cursors declared with names containing special characters.
1 parent 9b3b3b6 commit c0aa573

9 files changed

Lines changed: 181 additions & 37 deletions

File tree

‎pkg/sql/conn_executor_prepare.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -337,7 +337,7 @@ func (ex *connExecutor) execBind(
337337
return retErr(pgerror.Newf(
338338
pgcode.DuplicateCursor, "portal %q already exists", portalName))
339339
}
340-
if cursor := ex.getCursorAccessor().getCursor(portalName); cursor != nil {
340+
if cursor := ex.getCursorAccessor().getCursor(tree.Name(portalName)); cursor != nil {
341341
return retErr(pgerror.Newf(
342342
pgcode.DuplicateCursor, "portal %q already exists as cursor", portalName))
343343
}
@@ -493,7 +493,7 @@ func (ex *connExecutor) addPortal(
493493
if _, ok := ex.extraTxnState.prepStmtsNamespace.portals[portalName]; ok {
494494
panic(errors.AssertionFailedf("portal already exists: %q", portalName))
495495
}
496-
if cursor := ex.getCursorAccessor().getCursor(portalName); cursor != nil {
496+
if cursor := ex.getCursorAccessor().getCursor(tree.Name(portalName)); cursor != nil {
497497
panic(errors.AssertionFailedf("portal already exists as cursor: %q", portalName))
498498
}
499499

@@ -572,7 +572,7 @@ func (ex *connExecutor) execDescribe(
572572

573573
switch descCmd.Type {
574574
case pgwirebase.PrepareStatement:
575-
ps, ok := ex.extraTxnState.prepStmtsNamespace.prepStmts[descCmd.Name]
575+
ps, ok := ex.extraTxnState.prepStmtsNamespace.prepStmts[string(descCmd.Name)]
576576
if !ok {
577577
return retErr(pgerror.Newf(
578578
pgcode.InvalidSQLStatementName,
@@ -602,7 +602,9 @@ func (ex *connExecutor) execDescribe(
602602
res.SetPrepStmtOutput(ctx, ps.Columns)
603603
}
604604
case pgwirebase.PreparePortal:
605-
portal, ok := ex.extraTxnState.prepStmtsNamespace.portals[descCmd.Name]
605+
// TODO(rimadeodhar): prepStmtsNamespace should also be updated to use tree.Name instead of string
606+
// for indexing internal maps.
607+
portal, ok := ex.extraTxnState.prepStmtsNamespace.portals[string(descCmd.Name)]
606608
if !ok {
607609
// Check SQL-level cursors.
608610
cursor := ex.getCursorAccessor().getCursor(descCmd.Name)

‎pkg/sql/conn_io.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ var _ Command = PrepareStmt{}
218218
// DescribeStmt is the Command for producing info about a prepared statement or
219219
// portal.
220220
type DescribeStmt struct {
221-
Name string
221+
Name tree.Name
222222
Type pgwirebase.PrepareType
223223
}
224224

‎pkg/sql/logictest/testdata/logic_test/cursor‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,3 +557,45 @@ CLOSE foo
557557

558558
statement ok
559559
ALTER TABLE a ADD COLUMN c INT
560+
561+
statement ok
562+
COMMIT;
563+
564+
statement ok
565+
BEGIN;
566+
567+
statement ok
568+
DECLARE "a"" b'c" CURSOR FOR SELECT 1;
569+
570+
query I
571+
FETCH 1 "a"" b'c";
572+
----
573+
1
574+
575+
statement ok
576+
CLOSE "a"" b'c";
577+
DECLARE "a b" CURSOR FOR SELECT 2;
578+
579+
query I
580+
FETCH 1 "a b";
581+
----
582+
2
583+
584+
statement ok
585+
CLOSE "a b";
586+
DECLARE "a\b" CURSOR FOR SELECT 3;
587+
588+
query I
589+
FETCH 1 "a\b";
590+
----
591+
3
592+
593+
statement ok
594+
CLOSE "a\b";
595+
596+
query error pq: at or near "b": syntax error
597+
FETCH 1 a b;
598+
599+
statement ok
600+
COMMIT;
601+

‎pkg/sql/logictest/testdata/logic_test/pg_catalog‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6034,3 +6034,31 @@ SELECT * FROM pg_sequences WHERE sequencename = 'serial'
60346034
----
60356035
schemaname sequencename sequenceowner data_type start_value min_value max_value increment_by cycle cache_size last_value
60366036
public serial root 20 101 1 9223372036854775807 5 false 1 NULL
6037+
6038+
statement ok
6039+
CREATE TABLE t (a INT PRIMARY KEY, b INT);
6040+
INSERT INTO t VALUES (1, 2), (2, 3);
6041+
6042+
statement ok
6043+
BEGIN;
6044+
6045+
statement ok
6046+
DECLARE "a"" b'c" CURSOR FOR TABLE t;
6047+
DECLARE "a b" CURSOR FOR TABLE t;
6048+
DECLARE "a\b" CURSOR FOR TABLE t;
6049+
6050+
## pg_catalog.pg_cursors
6051+
6052+
query TTBBB colnames
6053+
SELECT name, statement, is_holdable, is_binary, is_scrollable FROM pg_catalog.pg_cursors ORDER BY name;
6054+
----
6055+
name statement is_holdable is_binary is_scrollable
6056+
a b TABLE t false false false
6057+
a" b'c TABLE t false false false
6058+
a\b TABLE t false false false
6059+
6060+
statement ok
6061+
COMMIT;
6062+
6063+
statement ok
6064+
DROP TABLE t;

‎pkg/sql/pg_catalog.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3776,12 +3776,12 @@ https://www.postgresql.org/docs/14/view-pg-cursors.html`,
37763776
return err
37773777
}
37783778
if err := addRow(
3779-
tree.NewDString(name), /* name */
3780-
tree.NewDString(c.statement), /* statement */
3781-
tree.DBoolFalse, /* is_holdable */
3782-
tree.DBoolFalse, /* is_binary */
3783-
tree.DBoolFalse, /* is_scrollable */
3784-
tz, /* creation_date */
3779+
tree.NewDString(string(name)), /* name */
3780+
tree.NewDString(c.statement), /* statement */
3781+
tree.DBoolFalse, /* is_holdable */
3782+
tree.DBoolFalse, /* is_binary */
3783+
tree.DBoolFalse, /* is_scrollable */
3784+
tz, /* creation_date */
37853785
); err != nil {
37863786
return err
37873787
}

‎pkg/sql/pgwire/conn.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1011,7 +1011,7 @@ func (c *conn) handleDescribe(ctx context.Context, buf *pgwirebase.ReadBuffer) e
10111011
return c.stmtBuf.Push(
10121012
ctx,
10131013
sql.DescribeStmt{
1014-
Name: name,
1014+
Name: tree.Name(name),
10151015
Type: typ,
10161016
})
10171017
}

‎pkg/sql/pgwire/conn_test.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,9 @@ func TestConn(t *testing.T) {
159159
if err := finishQuery(generateError, conn); err != nil {
160160
t.Fatal(err)
161161
}
162+
expectExecStmt(ctx, t, "DECLARE \"a b\" CURSOR FOR SELECT 10", &rd, conn, queryStringComplete)
163+
expectExecStmt(ctx, t, "FETCH 1 \"a b\"", &rd, conn, queryStringComplete)
164+
expectExecStmt(ctx, t, "CLOSE \"a b\"", &rd, conn, queryStringComplete)
162165
// We got to the COMMIT at the end of the batch.
163166
expectExecStmt(ctx, t, "COMMIT TRANSACTION", &rd, conn, queryStringComplete)
164167
expectSync(ctx, t, &rd)
@@ -505,6 +508,9 @@ func client(ctx context.Context, serverAddr net.Addr, wg *sync.WaitGroup) error
505508
batch.Queue("BEGIN")
506509
batch.Queue("select 7")
507510
batch.Queue("select 8")
511+
batch.Queue("declare \"a b\" cursor for select 10")
512+
batch.Queue("fetch 1 \"a b\"")
513+
batch.Queue("close \"a b\"")
508514
batch.Queue("COMMIT")
509515

510516
batchResults := conn.SendBatch(ctx, batch)
@@ -676,7 +682,7 @@ func expectPrepareStmt(
676682
func expectDescribeStmt(
677683
ctx context.Context,
678684
t *testing.T,
679-
expName string,
685+
expName tree.Name,
680686
expType pgwirebase.PrepareType,
681687
rd *sql.StmtBufReader,
682688
c *conn,

‎pkg/sql/pgwire/testdata/pgtest/portals‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1486,3 +1486,71 @@ ReadyForQuery
14861486
{"Type":"ParseComplete"}
14871487
{"Type":"ErrorResponse","Code":"08P01","Message":"invalid DESCRIBE message subtype 0"}
14881488
{"Type":"ReadyForQuery","TxStatus":"I"}
1489+
1490+
# Check declaring cursor with a name enclosed in double quotes
1491+
send
1492+
Query {"String": "BEGIN"}
1493+
----
1494+
1495+
until
1496+
ReadyForQuery
1497+
----
1498+
{"Type":"CommandComplete","CommandTag":"BEGIN"}
1499+
{"Type":"ReadyForQuery","TxStatus":"T"}
1500+
1501+
send
1502+
Parse {"Query": "DECLARE \"a b\" CURSOR FOR SELECT generate_series(1, 10) AS bar"}
1503+
Bind
1504+
Describe {"ObjectType": "P", "Name": ""}
1505+
Execute
1506+
Sync
1507+
----
1508+
1509+
until
1510+
ReadyForQuery
1511+
----
1512+
{"Type":"ParseComplete"}
1513+
{"Type":"BindComplete"}
1514+
{"Type":"NoData"}
1515+
{"Type":"CommandComplete","CommandTag":"DECLARE CURSOR"}
1516+
{"Type":"ReadyForQuery","TxStatus":"T"}
1517+
1518+
send
1519+
Describe {"ObjectType": "P", "Name": "a b"}
1520+
Sync
1521+
----
1522+
1523+
until ignore_type_oids ignore_table_oids ignore_data_type_sizes
1524+
ReadyForQuery
1525+
----
1526+
{"Type":"RowDescription","Fields":[{"Name":"bar","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":0,"DataTypeSize":0,"TypeModifier":-1,"Format":0}]}
1527+
{"Type":"ReadyForQuery","TxStatus":"T"}
1528+
1529+
send
1530+
Parse {"Query": "FETCH 2 \"a b\""}
1531+
Bind
1532+
Describe {"ObjectType": "P", "Name": ""}
1533+
Execute
1534+
Sync
1535+
----
1536+
1537+
until ignore_type_oids ignore_table_oids ignore_data_type_sizes
1538+
ReadyForQuery
1539+
----
1540+
{"Type":"ParseComplete"}
1541+
{"Type":"BindComplete"}
1542+
{"Type":"RowDescription","Fields":[{"Name":"bar","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":0,"DataTypeSize":0,"TypeModifier":-1,"Format":0}]}
1543+
{"Type":"DataRow","Values":[{"text":"1"}]}
1544+
{"Type":"DataRow","Values":[{"text":"2"}]}
1545+
{"Type":"CommandComplete","CommandTag":"FETCH 2"}
1546+
{"Type":"ReadyForQuery","TxStatus":"T"}
1547+
1548+
send
1549+
Query {"String": "ROLLBACK"}
1550+
----
1551+
1552+
until
1553+
ReadyForQuery
1554+
----
1555+
{"Type":"CommandComplete","CommandTag":"ROLLBACK"}
1556+
{"Type":"ReadyForQuery","TxStatus":"I"}

‎pkg/sql/sql_cursor.go‎

Lines changed: 22 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -49,13 +49,12 @@ func (p *planner) DeclareCursor(ctx context.Context, s *tree.DeclareCursor) (pla
4949
}
5050

5151
ie := p.ExecCfg().InternalExecutorFactory.NewInternalExecutor(p.SessionData())
52-
cursorName := s.Name.String()
53-
if cursor := p.sqlCursors.getCursor(cursorName); cursor != nil {
54-
return nil, pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists", cursorName)
52+
if cursor := p.sqlCursors.getCursor(s.Name); cursor != nil {
53+
return nil, pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists", s.Name)
5554
}
5655

57-
if p.extendedEvalCtx.PreparedStatementState.HasPortal(cursorName) {
58-
return nil, pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists as portal", cursorName)
56+
if p.extendedEvalCtx.PreparedStatementState.HasPortal(string(s.Name)) {
57+
return nil, pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists as portal", s.Name)
5958
}
6059

6160
// Try to plan the cursor query to make sure that it's valid.
@@ -100,7 +99,7 @@ func (p *planner) DeclareCursor(ctx context.Context, s *tree.DeclareCursor) (pla
10099
statement: statement,
101100
created: timeutil.Now(),
102101
}
103-
if err := p.sqlCursors.addCursor(cursorName, cursor); err != nil {
102+
if err := p.sqlCursors.addCursor(s.Name, cursor); err != nil {
104103
// This case shouldn't happen because cursor names are scoped to a session,
105104
// and sessions can't have more than one statement running at once. But
106105
// let's be diligent and clean up if it somehow does happen anyway.
@@ -119,11 +118,10 @@ var errBackwardScan = pgerror.Newf(pgcode.ObjectNotInPrerequisiteState, "cursor
119118
func (p *planner) FetchCursor(
120119
_ context.Context, s *tree.CursorStmt, isMove bool,
121120
) (planNode, error) {
122-
cursorName := s.Name.String()
123-
cursor := p.sqlCursors.getCursor(cursorName)
121+
cursor := p.sqlCursors.getCursor(s.Name)
124122
if cursor == nil {
125123
return nil, pgerror.Newf(
126-
pgcode.InvalidCursorName, "cursor %q does not exist", cursorName,
124+
pgcode.InvalidCursorName, "cursor %q does not exist", s.Name,
127125
)
128126
}
129127
if s.Count < 0 || s.FetchType == tree.FetchBackwardAll {
@@ -243,7 +241,7 @@ func (p *planner) CloseCursor(ctx context.Context, n *tree.CloseCursor) (planNod
243241
return &delayedNode{
244242
name: n.String(),
245243
constructor: func(ctx context.Context, p *planner) (planNode, error) {
246-
return newZeroNode(nil /* columns */), p.sqlCursors.closeCursor(n.Name.String())
244+
return newZeroNode(nil /* columns */), p.sqlCursors.closeCursor(n.Name)
247245
},
248246
}, nil
249247
}
@@ -276,20 +274,20 @@ type sqlCursors interface {
276274
closeAll()
277275
// closeCursor closes the named cursor, returning an error if that cursor
278276
// didn't exist in the set.
279-
closeCursor(string) error
277+
closeCursor(tree.Name) error
280278
// getCursor returns the named cursor, returning nil if that cursor
281279
// didn't exist in the set.
282-
getCursor(string) *sqlCursor
280+
getCursor(tree.Name) *sqlCursor
283281
// addCursor adds a new cursor with the given name to the set, returning an
284282
// error if the cursor already existed in the set.
285-
addCursor(string, *sqlCursor) error
283+
addCursor(tree.Name, *sqlCursor) error
286284
// list returns all open cursors in the set.
287-
list() map[string]*sqlCursor
285+
list() map[tree.Name]*sqlCursor
288286
}
289287

290288
// cursorMap is a sqlCursors that's backed by an actual map.
291289
type cursorMap struct {
292-
cursors map[string]*sqlCursor
290+
cursors map[tree.Name]*sqlCursor
293291
}
294292

295293
func (c *cursorMap) closeAll() {
@@ -299,7 +297,7 @@ func (c *cursorMap) closeAll() {
299297
c.cursors = nil
300298
}
301299

302-
func (c *cursorMap) closeCursor(s string) error {
300+
func (c *cursorMap) closeCursor(s tree.Name) error {
303301
cursor, ok := c.cursors[s]
304302
if !ok {
305303
return pgerror.Newf(pgcode.InvalidCursorName, "cursor %q does not exist", s)
@@ -309,13 +307,13 @@ func (c *cursorMap) closeCursor(s string) error {
309307
return err
310308
}
311309

312-
func (c *cursorMap) getCursor(s string) *sqlCursor {
310+
func (c *cursorMap) getCursor(s tree.Name) *sqlCursor {
313311
return c.cursors[s]
314312
}
315313

316-
func (c *cursorMap) addCursor(s string, cursor *sqlCursor) error {
314+
func (c *cursorMap) addCursor(s tree.Name, cursor *sqlCursor) error {
317315
if c.cursors == nil {
318-
c.cursors = make(map[string]*sqlCursor)
316+
c.cursors = make(map[tree.Name]*sqlCursor)
319317
}
320318
if _, ok := c.cursors[s]; ok {
321319
return pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists", s)
@@ -324,7 +322,7 @@ func (c *cursorMap) addCursor(s string, cursor *sqlCursor) error {
324322
return nil
325323
}
326324

327-
func (c *cursorMap) list() map[string]*sqlCursor {
325+
func (c *cursorMap) list() map[tree.Name]*sqlCursor {
328326
return c.cursors
329327
}
330328

@@ -338,19 +336,19 @@ func (c connExCursorAccessor) closeAll() {
338336
c.ex.extraTxnState.sqlCursors.closeAll()
339337
}
340338

341-
func (c connExCursorAccessor) closeCursor(s string) error {
339+
func (c connExCursorAccessor) closeCursor(s tree.Name) error {
342340
return c.ex.extraTxnState.sqlCursors.closeCursor(s)
343341
}
344342

345-
func (c connExCursorAccessor) getCursor(s string) *sqlCursor {
343+
func (c connExCursorAccessor) getCursor(s tree.Name) *sqlCursor {
346344
return c.ex.extraTxnState.sqlCursors.getCursor(s)
347345
}
348346

349-
func (c connExCursorAccessor) addCursor(s string, cursor *sqlCursor) error {
347+
func (c connExCursorAccessor) addCursor(s tree.Name, cursor *sqlCursor) error {
350348
return c.ex.extraTxnState.sqlCursors.addCursor(s, cursor)
351349
}
352350

353-
func (c connExCursorAccessor) list() map[string]*sqlCursor {
351+
func (c connExCursorAccessor) list() map[tree.Name]*sqlCursor {
354352
return c.ex.extraTxnState.sqlCursors.list()
355353
}
356354

0 commit comments

Comments
 (0)