Make DROP/CREATE DATABASE pluggable

go/vt/vtgate/engine/create_database.go

@@ -0,0 +1,92 @@
+/*
+Copyright 2020 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package engine
+
+import (
+	"strings"
+
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/vtgate/vindexes"
+)
+
+var _ Primitive = (*DropCreateDatabase)(nil)
+
+// DropCreateDatabase is just a container around custom database provisioning plugins
+// The default behaviour is to just return an error
+type DropCreateDatabase struct {
+	name, verb string
+	plugin     func() error
+	noInputs
+	noTxNeeded
+}
+
+// CreateDropCreateDatabase creates the engine primitive
+func CreateDropCreateDatabase(dbName, verb string, plugin func() error) *DropCreateDatabase {
+	return &DropCreateDatabase{
+		name:   dbName,
+		verb:   verb,
+		plugin: plugin,
+	}
+}
+
+// RouteType implements the Primitive interface
+func (c *DropCreateDatabase) RouteType() string {
+	return c.verb + " database"
+}
+
+// GetKeyspaceName implements the Primitive interface
+func (c *DropCreateDatabase) GetKeyspaceName() string {
+	return c.name
+}
+
+// GetTableName implements the Primitive interface
+func (c *DropCreateDatabase) GetTableName() string {
+	return ""
+}
+
+// Execute implements the Primitive interface
+func (c *DropCreateDatabase) Execute(VCursor, map[string]*querypb.BindVariable, bool) (*sqltypes.Result, error) {
+	err := c.plugin()
+	if err != nil {
+		return nil, err
+	}
+
+	return &sqltypes.Result{}, nil
+}
+
+// StreamExecute implements the Primitive interface
+func (c *DropCreateDatabase) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	res, err := c.Execute(vcursor, bindVars, wantfields)
+	if err != nil {
+		return err
+	}
+	return callback(res)
+}
+
+// GetFields implements the Primitive interface
+func (c *DropCreateDatabase) GetFields(VCursor, map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	return &sqltypes.Result{}, nil
+}
+
+// description implements the Primitive interface
+func (c *DropCreateDatabase) description() PrimitiveDescription {
+	return PrimitiveDescription{
+		OperatorType: strings.ToUpper(c.verb + " DATABASE"),
+		Keyspace:     &vindexes.Keyspace{Name: c.name},
+	}
+}

go/vt/vtgate/planbuilder/builder.go

@@ -208,7 +208,7 @@ func buildDBDDLPlan(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Pr
 		if !ksExists {
 			return nil, mysql.NewSQLError(mysql.ERDbDropExists, mysql.SSUnknownSQLState, "Can't drop database '%s'; database doesn't exists", ksName)
 		}
-		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "drop database not allowed")
+		return engine.CreateDropCreateDatabase(dbDDL.DBName, "drop", func() error { return databaseCreator.DropDatabase(dbDDL) }), nil
 	case *sqlparser.AlterDatabase:
 		if !ksExists {
 			return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "Can't alter database '%s'; database doesn't exists", ksName)
@@ -221,7 +221,7 @@ func buildDBDDLPlan(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Pr
 		if !dbDDL.IfNotExists && ksExists {
 			return nil, mysql.NewSQLError(mysql.ERDbCreateExists, mysql.SSUnknownSQLState, "Can't create database '%s'; database exists", ksName)
 		}
-		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "create database not allowed")
+		return engine.CreateDropCreateDatabase(dbDDL.DBName, "create", func() error { return databaseCreator.CreateDatabase(dbDDL) }), nil
 	}
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unreachable code path: %s", sqlparser.String(dbDDLstmt))
 }

go/vt/vtgate/planbuilder/create_database_pluggable.go

@@ -0,0 +1,44 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package planbuilder
+
+import (
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vterrors"
+)
+
+// DropCreateDB is the interface that you need to implement to add a custom CREATE/DROP DATABASE handler
+type DropCreateDB interface {
+	CreateDatabase(ast *sqlparser.CreateDatabase) error
+	DropDatabase(ast *sqlparser.DropDatabase) error
+}
+
+type defaultHandler struct{}
+
+// CreateDatabase implements the DropCreateDB interface
+func (defaultHandler) CreateDatabase(*sqlparser.CreateDatabase) error {
+	return vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "create database not allowed")
+}
+
+// DropDatabase implements the DropCreateDB interface
+func (defaultHandler) DropDatabase(*sqlparser.DropDatabase) error {
+	return vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "drop database not allowed")
+}
+
+//goland:noinspection GoVarAndConstTypeMayBeOmitted
+var databaseCreator DropCreateDB = defaultHandler{}

go/vt/vtgate/planbuilder/create_database_pluggable_test.go

@@ -0,0 +1,70 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package planbuilder
+
+import (
+	"testing"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/vtgate/vindexes"
+)
+
+func TestCreateDB(t *testing.T) {
+	ks := &vindexes.Keyspace{Name: "main"}
+	vschema := &vschemaWrapper{
+		v: &vindexes.VSchema{
+			Keyspaces: map[string]*vindexes.KeyspaceSchema{"main": {Keyspace: ks}},
+		},
+		keyspace: ks,
+	}
+
+	// default behaviour
+	plan, err := TestBuilder("create database test", vschema)
+	require.NoError(t, err)
+
+	_, err = plan.Instructions.Execute(nil, nil, true)
+	require.EqualError(t, err, "create database not allowed")
+
+	// we make sure to restore the state so we don't destabilize other tests
+	before := databaseCreator
+	defer func() {
+		databaseCreator = before
+	}()
+
+	// setting custom behaviour for CREATE DATABASE
+	databaseCreator = &fakePlugin{}
+
+	plan, err = TestBuilder("create database test", vschema)
+	require.NoError(t, err)
+	_, err = plan.Instructions.Execute(nil, nil, true)
+	require.NoError(t, err)
+}
+
+type fakePlugin struct{}
+
+// CreateDatabase is a fake
+func (*fakePlugin) CreateDatabase(ast *sqlparser.CreateDatabase) error {
+	return nil
+}
+
+// DropDatabase is a fake
+func (*fakePlugin) DropDatabase(ast *sqlparser.DropDatabase) error {
+	return nil
+}

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

@@ -98,7 +98,17 @@
 
 # create db foo
 "create database foo"
-"create database not allowed"
+{
+  "QueryType": "DDL",
+  "Original": "create database foo",
+  "Instructions": {
+    "OperatorType": "CREATE DATABASE",
+    "Keyspace": {
+      "Name": "foo",
+      "Sharded": false
+    }
+  }
+}
 
 # create db main
 "create database main"
@@ -128,11 +138,31 @@
 
 # drop db main
 "drop database main"
-"drop database not allowed"
+{
+  "QueryType": "DDL",
+  "Original": "drop database main",
+  "Instructions": {
+    "OperatorType": "DROP DATABASE",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    }
+  }
+}
 
 # drop db if exists main
 "drop database if exists main"
-"drop database not allowed"
+{
+  "QueryType": "DDL",
+  "Original": "drop database if exists main",
+  "Instructions": {
+    "OperatorType": "DROP DATABASE",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    }
+  }
+}
 
 # drop db if exists foo
 "drop schema if exists foo"