From 04bf2f33c97f30658d2da2535be6c332501e026d Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Fri, 19 Aug 2022 15:49:16 -0400
Subject: [PATCH] sql: refactor DISCARD so it runs during execution

Previously, it would run during planning, which is not correct, or at
least, makes it harder to reason about what's going on.

Release note: None

Release justification: low risk refactor
---
 pkg/sql/discard.go                    | 29 ++++++++++++++++++---------
 pkg/sql/pgwire/testdata/pgtest/pgjdbc |  8 +-------
 pkg/sql/sem/tree/stmt.go              |  8 +++++++-
 pkg/sql/walk.go                       |  1 +
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/pkg/sql/discard.go b/pkg/sql/discard.go
index 3ea8e35382f0..d6227dea2a04 100644
--- a/pkg/sql/discard.go
+++ b/pkg/sql/discard.go
@@ -22,28 +22,39 @@ import (
 // Discard implements the DISCARD statement.
 // See https://www.postgresql.org/docs/9.6/static/sql-discard.html for details.
 func (p *planner) Discard(ctx context.Context, s *tree.Discard) (planNode, error) {
-	switch s.Mode {
+	return &discardNode{mode: s.Mode}, nil
+}
+
+type discardNode struct {
+	mode tree.DiscardMode
+}
+
+func (n *discardNode) Next(_ runParams) (bool, error) { return false, nil }
+func (n *discardNode) Values() tree.Datums            { return nil }
+func (n *discardNode) Close(_ context.Context)        {}
+func (n *discardNode) startExec(params runParams) error {
+	switch n.mode {
 	case tree.DiscardModeAll:
-		if !p.autoCommit {
-			return nil, pgerror.New(pgcode.ActiveSQLTransaction,
+		if !params.p.autoCommit {
+			return pgerror.New(pgcode.ActiveSQLTransaction,
 				"DISCARD ALL cannot run inside a transaction block")
 		}
 
 		// RESET ALL
-		if err := p.sessionDataMutatorIterator.applyOnEachMutatorError(
+		if err := params.p.sessionDataMutatorIterator.applyOnEachMutatorError(
 			func(m sessionDataMutator) error {
-				return resetSessionVars(ctx, m)
+				return resetSessionVars(params.ctx, m)
 			},
 		); err != nil {
-			return nil, err
+			return err
 		}
 
 		// DEALLOCATE ALL
-		p.preparedStatements.DeleteAll(ctx)
+		params.p.preparedStatements.DeleteAll(params.ctx)
 	default:
-		return nil, errors.AssertionFailedf("unknown mode for DISCARD: %d", s.Mode)
+		return errors.AssertionFailedf("unknown mode for DISCARD: %d", n.mode)
 	}
-	return newZeroNode(nil /* columns */), nil
+	return nil
 }
 
 func resetSessionVars(ctx context.Context, m sessionDataMutator) error {
diff --git a/pkg/sql/pgwire/testdata/pgtest/pgjdbc b/pkg/sql/pgwire/testdata/pgtest/pgjdbc
index d7d2c4329197..2494eb3530b3 100644
--- a/pkg/sql/pgwire/testdata/pgtest/pgjdbc
+++ b/pkg/sql/pgwire/testdata/pgtest/pgjdbc
@@ -42,13 +42,7 @@ send
 Query {"String": "DISCARD ALL"}
 ----
 
-until crdb_only ignore=ParameterStatus
-ReadyForQuery
-----
-{"Type":"CommandComplete","CommandTag":"DISCARD"}
-{"Type":"ReadyForQuery","TxStatus":"I"}
-
-until noncrdb_only ignore=ParameterStatus
+until ignore=ParameterStatus ignore=NoticeResponse
 ReadyForQuery
 ----
 {"Type":"CommandComplete","CommandTag":"DISCARD ALL"}
diff --git a/pkg/sql/sem/tree/stmt.go b/pkg/sql/sem/tree/stmt.go
index bb0fa47c2d76..d09f7e72f7fb 100644
--- a/pkg/sql/sem/tree/stmt.go
+++ b/pkg/sql/sem/tree/stmt.go
@@ -853,7 +853,13 @@ func (*Discard) StatementReturnType() StatementReturnType { return Ack }
 func (*Discard) StatementType() StatementType { return TypeTCL }
 
 // StatementTag returns a short string identifying the type of statement.
-func (*Discard) StatementTag() string { return "DISCARD" }
+func (d *Discard) StatementTag() string {
+	switch d.Mode {
+	case DiscardModeAll:
+		return "DISCARD ALL"
+	}
+	return "DISCARD"
+}
 
 // StatementReturnType implements the Statement interface.
 func (n *DeclareCursor) StatementReturnType() StatementReturnType { return Ack }
diff --git a/pkg/sql/walk.go b/pkg/sql/walk.go
index de6bd221537a..2a57fe543325 100644
--- a/pkg/sql/walk.go
+++ b/pkg/sql/walk.go
@@ -393,6 +393,7 @@ var planNodeNames = map[reflect.Type]string{
 	reflect.TypeOf(&delayedNode{}):                             "virtual table",
 	reflect.TypeOf(&deleteNode{}):                              "delete",
 	reflect.TypeOf(&deleteRangeNode{}):                         "delete range",
+	reflect.TypeOf(&discardNode{}):                             "discard",
 	reflect.TypeOf(&distinctNode{}):                            "distinct",
 	reflect.TypeOf(&dropDatabaseNode{}):                        "drop database",
 	reflect.TypeOf(&dropExternalConnectionNode{}):              "drop external connection",