Merge pull request #7034 from planetscale/improve-dbddl

Improve database ddl plan

go/vt/vtgate/planbuilder/builder.go

@@ -19,6 +19,9 @@ package planbuilder
 import (
 	"errors"
 
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+
 	"vitess.io/vitess/go/vt/vterrors"
 
 	"vitess.io/vitess/go/vt/key"
@@ -133,6 +136,7 @@ type ContextVSchema interface {
 	AnyKeyspace() (*vindexes.Keyspace, error)
 	FirstSortedKeyspace() (*vindexes.Keyspace, error)
 	SysVarSetEnabled() bool
+	KeyspaceExists(keyspace string) bool
 }
 
 //-------------------------------------------------------------------------
@@ -340,7 +344,7 @@ func createInstructionFor(query string, stmt sqlparser.Statement, vschema Contex
 	case *sqlparser.Load:
 		return buildLoadPlan(query, vschema)
 	case *sqlparser.DBDDL:
-		return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: Database DDL %v", sqlparser.String(stmt))
+		return buildRoutePlan(stmt, vschema, buildDBDDLPlan)
 	case *sqlparser.SetTransaction:
 		return nil, ErrPlanNotSupported
 	case *sqlparser.Begin, *sqlparser.Commit, *sqlparser.Rollback, *sqlparser.Savepoint, *sqlparser.SRollback, *sqlparser.Release:
@@ -353,6 +357,35 @@ func createInstructionFor(query string, stmt sqlparser.Statement, vschema Contex
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "BUG: unexpected statement type: %T", stmt)
 }
 
+func buildDBDDLPlan(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Primitive, error) {
+	dbDDL := stmt.(*sqlparser.DBDDL)
+	ksExists := vschema.KeyspaceExists(dbDDL.DBName)
+	switch dbDDL.Action {
+	case sqlparser.CreateDBDDLAction:
+		if dbDDL.IfNotExists && ksExists {
+			return engine.NewRowsPrimitive(make([][]sqltypes.Value, 0), make([]*querypb.Field, 0)), nil
+		}
+		if !dbDDL.IfNotExists && ksExists {
+			return nil, vterrors.Errorf(vtrpcpb.Code_ALREADY_EXISTS, "cannot create database '%s'; database exists", dbDDL.DBName)
+		}
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "create database not allowed")
+	case sqlparser.AlterDBDDLAction:
+		if !ksExists {
+			return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "cannot alter database '%s'; database does not exists", dbDDL.DBName)
+		}
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "alter database not allowed")
+	case sqlparser.DropDBDDLAction:
+		if dbDDL.IfExists && !ksExists {
+			return engine.NewRowsPrimitive(make([][]sqltypes.Value, 0), make([]*querypb.Field, 0)), nil
+		}
+		if !ksExists {
+			return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "cannot drop database '%s'; database does not exists", dbDDL.DBName)
+		}
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "drop database not allowed")
+	}
+	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unreachable code path: %s", sqlparser.String(dbDDL))
+}
+
 func buildLoadPlan(query string, vschema ContextVSchema) (engine.Primitive, error) {
 	keyspace, err := vschema.DefaultKeyspace()
 	if err != nil {

go/vt/vtgate/planbuilder/plan_test.go

@@ -282,6 +282,13 @@ type vschemaWrapper struct {
 	sysVarEnabled bool
 }
 
+func (vw *vschemaWrapper) KeyspaceExists(keyspace string) bool {
+	if vw.keyspace != nil {
+		return vw.keyspace.Name == keyspace
+	}
+	return false
+}
+
 func (vw *vschemaWrapper) SysVarSetEnabled() bool {
 	return vw.sysVarEnabled
 }

go/vt/vtgate/planbuilder/testdata/ddl_cases.txt

@@ -96,6 +96,44 @@
 "alter table a.b ADD id int"
 "no keyspace with name [a] found"
 
+# create db foo
+"create database foo"
+"create database not allowed"
+
+# create db main
+"create database main"
+"cannot create database 'main'; database exists"
+
+# create db if not exists main
+"create database if not exists main"
+{
+  "QueryType": "DDL",
+  "Original": "create database if not exists main",
+  "Instructions": {
+    "OperatorType": "Rows"
+  }
+}
+
+# alter db foo
+"alter database foo collate utf8"
+"cannot alter database 'foo'; database does not exists"
+
+# alter db main
+"alter database main collate utf8"
+"alter database not allowed"
+
+# drop db foo
+"drop database foo"
+"cannot drop database 'foo'; database does not exists"
+
+# drop db main
+"drop database main"
+"drop database not allowed"
+
+# drop db if exists main
+"drop database if exists main"
+"drop database not allowed"
+
 # DDL
 "create index a on user(id)"
 {

go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt

@@ -338,10 +338,6 @@
 "delete music,user from music inner join user where music.id = user.id"
 "unsupported: multi-shard or vindex write statement"
 
-# Database DDL
-"create database foo"
-"unsupported: Database DDL create database foo"
-
 # order by inside and outside parenthesis select
 "(select 1 from user order by 1 desc) order by 1 asc limit 2"
 "can't do ORDER BY on top of ORDER BY"

go/vt/vtgate/vcursor_impl.go

@@ -318,6 +318,16 @@ func (vc *vcursorImpl) FirstSortedKeyspace() (*vindexes.Keyspace, error) {
 	return kss[keys[0]].Keyspace, nil
 }
 
+// SysVarSetEnabled implements the ContextVSchema interface
+func (vc *vcursorImpl) SysVarSetEnabled() bool {
+	return *sysVarSetEnabled
+}
+
+// KeyspaceExists provides whether the keyspace exists or not.
+func (vc *vcursorImpl) KeyspaceExists(ks string) bool {
+	return vc.vschema.Keyspaces[ks] != nil
+}
+
 // TargetString returns the current TargetString of the session.
 func (vc *vcursorImpl) TargetString() string {
 	return vc.safeSession.TargetString
@@ -574,11 +584,6 @@ func (vc *vcursorImpl) SetWorkload(workload querypb.ExecuteOptions_Workload) {
 	vc.safeSession.GetOrCreateOptions().Workload = workload
 }
 
-// SysVarSetEnabled implements the SessionActions interface
-func (vc *vcursorImpl) SysVarSetEnabled() bool {
-	return *sysVarSetEnabled
-}
-
 // SetFoundRows implements the SessionActions interface
 func (vc *vcursorImpl) SetFoundRows(foundRows uint64) {
 	vc.safeSession.FoundRows = foundRows