sql: add drop type to the prepared statement generator

Release justification: fixes for high-priority or high-severity bugs in existing functionality Previously, prepare statement doesnt have support for drop type in optimizer, so it panic during execution as no flags are generated. This patch fixes this by adding dropType to the optimizer. Resolves #61226 Release note: none Signed-off-by: Tharun <[email protected]>
cockroachdb · Mar 4, 2021 · 556ea20 · 556ea20
1 parent c0d8340
commit 556ea20
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 31 deletions.
diff --git a/pkg/sql/drop_type.go b/pkg/sql/drop_type.go
@@ -28,14 +28,9 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
-type typeToDrop struct {
-	desc   *typedesc.Mutable
-	fqName string
-}
-
 type dropTypeNode struct {
-	n  *tree.DropType
-	td map[descpb.ID]typeToDrop
+	n      *tree.DropType
+	toDrop map[descpb.ID]*typedesc.Mutable
 }
 
 // Use to satisfy the linter.
@@ -51,8 +46,8 @@ func (p *planner) DropType(ctx context.Context, n *tree.DropType) (planNode, err
 	}
 
 	node := &dropTypeNode{
-		n:  n,
-		td: make(map[descpb.ID]typeToDrop),
+		n:      n,
+		toDrop: make(map[descpb.ID]*typedesc.Mutable),
 	}
 	if n.DropBehavior == tree.DropCascade {
 		return nil, unimplemented.NewWithIssue(51480, "DROP TYPE CASCADE is not yet supported")
@@ -67,7 +62,7 @@ func (p *planner) DropType(ctx context.Context, n *tree.DropType) (planNode, err
 			continue
 		}
 		// If we've already seen this type, then skip it.
-		if _, ok := node.td[typeDesc.ID]; ok {
+		if _, ok := node.toDrop[typeDesc.ID]; ok {
 			continue
 		}
 		switch typeDesc.Kind {
@@ -105,23 +100,9 @@ func (p *planner) DropType(ctx context.Context, n *tree.DropType) (planNode, err
 		if err := p.canDropTypeDesc(ctx, mutArrayDesc, n.DropBehavior); err != nil {
 			return nil, err
 		}
-
 		// Record these descriptors for deletion.
-		node.td[typeDesc.ID] = typeToDrop{
-			desc:   typeDesc,
-			fqName: tree.AsStringWithFQNames(name, p.Ann()),
-		}
-		arrayFQName, err := getTypeNameFromTypeDescriptor(
-			oneAtATimeSchemaResolver{ctx, p},
-			mutArrayDesc,
-		)
-		if err != nil {
-			return nil, err
-		}
-		node.td[mutArrayDesc.ID] = typeToDrop{
-			desc:   mutArrayDesc,
-			fqName: arrayFQName.FQString(),
-		}
+		node.toDrop[typeDesc.ID] = typeDesc
+		node.toDrop[mutArrayDesc.ID] = mutArrayDesc
 	}
 	return node, nil
 }
@@ -148,13 +129,23 @@ func (p *planner) canDropTypeDesc(
 }
 
 func (n *dropTypeNode) startExec(params runParams) error {
-	for _, toDrop := range n.td {
-		typ, fqName := toDrop.desc, toDrop.fqName
-		if err := params.p.dropTypeImpl(params.ctx, typ, tree.AsStringWithFQNames(n.n, params.Ann()), true /* queueJob */); err != nil {
+	for _, typeDesc := range n.toDrop {
+		typeFQName, err := getTypeNameFromTypeDescriptor(
+			oneAtATimeSchemaResolver{params.ctx, params.p},
+			typeDesc,
+		)
+		if err != nil {
 			return err
 		}
+		err = params.p.dropTypeImpl(params.ctx, typeDesc, "dropping type "+typeFQName.FQString(), true /* queueJob */)
+		if err != nil {
+			return err
+		}
+		event := &eventpb.DropType{
+			TypeName: typeFQName.FQString(),
+		}
 		// Log a Drop Type event.
-		if err := params.p.logEvent(params.ctx, typ.ID, &eventpb.DropType{TypeName: fqName}); err != nil {
+		if err := params.p.logEvent(params.ctx, typeDesc.ID, event); err != nil {
 			return err
 		}
 	}

diff --git a/pkg/sql/explain_test.go b/pkg/sql/explain_test.go
@@ -37,6 +37,7 @@ func TestStatementReuses(t *testing.T) {
 		`CREATE SEQUENCE s`,
 		`CREATE INDEX woo ON a(b)`,
 		`CREATE USER woo`,
+		`CREATE TYPE test as ENUM('a')`,
 	}
 
 	for _, s := range initStmts {
@@ -54,6 +55,7 @@ func TestStatementReuses(t *testing.T) {
 		`DROP SEQUENCE s`,
 		`DROP VIEW v`,
 		`DROP USER woo`,
+		`DROP TYPE test`,
 
 		// Ditto ALTER first, so that erroneous side effects bork what's
 		// below.

diff --git a/pkg/sql/logictest/testdata/logic_test/prepare b/pkg/sql/logictest/testdata/logic_test/prepare
@@ -1355,6 +1355,10 @@ CREATE TABLE greeting_table (x greeting NOT NULL, y INT, INDEX (x, y))
 statement ok
 PREPARE enum_query AS SELECT x, y FROM greeting_table WHERE y = 2
 
+# Create prepared statement to drop type.
+statement ok
+PREPARE enum_drop AS DROP TYPE greeting
+
 # Now alter the enum to have a new value.
 statement ok
 ALTER TYPE greeting ADD VALUE 'howdy'
@@ -1368,3 +1372,11 @@ query TI
 EXECUTE enum_query
 ----
 howdy 2
+
+# drop table
+statement ok
+DROP TABLE greeting_table
+
+# drop the type using prepared statement.
+statement ok
+EXECUTE enum_drop
diff --git a/pkg/sql/plan_opt.go b/pkg/sql/plan_opt.go
@@ -61,7 +61,7 @@ func (p *planner) prepareUsingOptimizer(ctx context.Context) (planFlags, error)
 		*tree.CreateSequence,
 		*tree.CreateStats,
 		*tree.Deallocate, *tree.Discard, *tree.DropDatabase, *tree.DropIndex,
-		*tree.DropTable, *tree.DropView, *tree.DropSequence,
+		*tree.DropTable, *tree.DropView, *tree.DropSequence, *tree.DropType,
 		*tree.Execute,
 		*tree.Grant, *tree.GrantRole,
 		*tree.Prepare,