From b7c45c1631fbbaf908ce8919b8e23057feb2e8ca Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Fri, 10 Jun 2022 17:20:06 -0400
Subject: [PATCH] sql: fix DECLARE with pgwire placeholders

Previously, it was not possible to prepare a DECLARE statement via
pgwire if it included placeholders. That limitation is now lifted.

Release note (bug fix): add missing support for preparing a DECLARE
cursor statement with placeholders.
---
 pkg/sql/pgwire/testdata/pgtest/portals | 38 ++++++++++++++++++++++++++
 pkg/sql/plan_opt.go                    | 10 +++++++
 pkg/sql/sql_cursor.go                  |  2 +-
 pkg/sql/sql_cursor_test.go             | 12 +++-----
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/pgwire/testdata/pgtest/portals b/pkg/sql/pgwire/testdata/pgtest/portals
index 4c3c63af1c5c..def589e3c10c 100644
--- a/pkg/sql/pgwire/testdata/pgtest/portals
+++ b/pkg/sql/pgwire/testdata/pgtest/portals
@@ -1356,3 +1356,41 @@ ReadyForQuery
 ----
 {"Type":"RowDescription","Fields":[{"Name":"g","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
 {"Type":"ReadyForQuery","TxStatus":"T"}
+
+# Test preparing a SQL-level DECLARE with placeholders
+
+send
+Parse {"Query": "DECLARE bar CURSOR FOR SELECT g::INT8 FROM generate_series(1, $1) g(g)", "ParameterOIDs": [21]}
+Bind {"PreparedStatement": "", "ParameterFormatCodes": [0], "ResultFormatCodes": [0], "Parameters": [{"text":"2"}]}
+Describe {"ObjectType": "P", "Name": ""}
+Execute
+Sync
+----
+
+until
+ReadyForQuery
+----
+{"Type":"ParseComplete"}
+{"Type":"BindComplete"}
+{"Type":"NoData"}
+{"Type":"CommandComplete","CommandTag":"DECLARE CURSOR"}
+{"Type":"ReadyForQuery","TxStatus":"T"}
+
+send
+Parse {"Query": "FETCH 2 bar"}
+Bind
+Describe {"ObjectType": "P", "Name": ""}
+Execute
+Sync
+----
+
+until
+ReadyForQuery
+----
+{"Type":"ParseComplete"}
+{"Type":"BindComplete"}
+{"Type":"RowDescription","Fields":[{"Name":"g","TableOID":0,"TableAttributeNumber":0,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
+{"Type":"DataRow","Values":[{"text":"1"}]}
+{"Type":"DataRow","Values":[{"text":"2"}]}
+{"Type":"CommandComplete","CommandTag":"FETCH 2"}
+{"Type":"ReadyForQuery","TxStatus":"T"}
diff --git a/pkg/sql/plan_opt.go b/pkg/sql/plan_opt.go
index 57c968b1d8d9..9cf9d7c4e202 100644
--- a/pkg/sql/plan_opt.go
+++ b/pkg/sql/plan_opt.go
@@ -108,6 +108,16 @@ func (p *planner) prepareUsingOptimizer(ctx context.Context) (planFlags, error)
 		}
 		stmt.Prepared.Columns = colinfo.ExplainPlanColumns
 		return opc.flags, nil
+	case *tree.DeclareCursor:
+		// Build memo for the purposes of typing placeholders.
+		// TODO(jordan): converting DeclareCursor to not be an opaque statement
+		// would be a better way to accomplish this goal. See CREATE TABLE for an
+		// example.
+		f := opc.optimizer.Factory()
+		bld := optbuilder.New(ctx, &p.semaCtx, p.EvalContext(), &opc.catalog, f, t.Select)
+		if err := bld.Build(); err != nil {
+			return opc.flags, err
+		}
 	}
 
 	if opc.useCache {
diff --git a/pkg/sql/sql_cursor.go b/pkg/sql/sql_cursor.go
index cb0c14b5a581..40faa9b8f197 100644
--- a/pkg/sql/sql_cursor.go
+++ b/pkg/sql/sql_cursor.go
@@ -86,7 +86,7 @@ func (p *planner) DeclareCursor(ctx context.Context, s *tree.DeclareCursor) (pla
 					"DECLARE CURSOR must not contain data-modifying statements in WITH")
 			}
 
-			statement := s.Select.String()
+			statement := formatWithPlaceholders(s.Select, p.EvalContext())
 			itCtx := context.Background()
 			rows, err := ie.QueryIterator(itCtx, "sql-cursor", p.txn, statement)
 			if err != nil {
diff --git a/pkg/sql/sql_cursor_test.go b/pkg/sql/sql_cursor_test.go
index 3b96c5528933..94f7f87b07c1 100644
--- a/pkg/sql/sql_cursor_test.go
+++ b/pkg/sql/sql_cursor_test.go
@@ -16,7 +16,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
@@ -110,15 +109,12 @@ func TestPrepareCursors(t *testing.T) {
 
 	// Make sure that we can use the automatic prepare support (when sending
 	// placeholders) to do the same thing.
-	// TODO (jordan): This currently doesn't work, because we don't fully walk
-	// the tree typechecking expressions when sending a DECLARE through the
-	// optimizer. See issue #77067.
-	// When this limitation is lifted, the rest of this test should be uncommented.
 	t.Run("prepare_declare_placeholder", func(t *testing.T) {
-		_, err = conn.ExecContext(ctx, "DECLARE foo CURSOR FOR SELECT 1 WHERE $1", true)
-		require.Contains(t, err.Error(), "could not determine data type of placeholder")
+		_, err = conn.ExecContext(ctx, "BEGIN TRANSACTION")
+		require.NoError(t, err)
 
-		skip.WithIssue(t, 77067, "placeholders in DECLARE not supported")
+		_, err = conn.ExecContext(ctx, "DECLARE foo CURSOR FOR SELECT 1 WHERE $1", true)
+		require.NoError(t, err)
 
 		stmt, err := conn.PrepareContext(ctx, "FETCH 1 foo")
 		require.NoError(t, err)