From b7400d43f5b1affedb9cf9cdb4f7f82f5f9ee670 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Tue, 26 Jan 2021 07:56:49 +0100
Subject: [PATCH] Start of making CREATE DATABASE pluggable

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 go/vt/vtgate/planbuilder/builder.go           |  2 +-
 .../planbuilder/create_database_pluggable.go  | 36 +++++++++++
 .../create_database_pluggable_test.go         | 59 +++++++++++++++++++
 3 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 go/vt/vtgate/planbuilder/create_database_pluggable.go
 create mode 100644 go/vt/vtgate/planbuilder/create_database_pluggable_test.go

diff --git a/go/vt/vtgate/planbuilder/builder.go b/go/vt/vtgate/planbuilder/builder.go
index 3575425dd23..272a3079254 100644
--- a/go/vt/vtgate/planbuilder/builder.go
+++ b/go/vt/vtgate/planbuilder/builder.go
@@ -198,7 +198,7 @@ func buildDBDDLPlan(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Pr
 		if !dbDDL.IfNotExists && ksExists {
 			return nil, vterrors.Errorf(vtrpcpb.Code_ALREADY_EXISTS, "cannot create database '%s'; database exists", ksName)
 		}
-		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "create database not allowed")
+		return databaseCreator(dbDDL, vschema)
 	}
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unreachable code path: %s", sqlparser.String(dbDDLstmt))
 }
diff --git a/go/vt/vtgate/planbuilder/create_database_pluggable.go b/go/vt/vtgate/planbuilder/create_database_pluggable.go
new file mode 100644
index 00000000000..25c6da7bcf4
--- /dev/null
+++ b/go/vt/vtgate/planbuilder/create_database_pluggable.go
@@ -0,0 +1,36 @@
+/*
+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"
+	"vitess.io/vitess/go/vt/vtgate/engine"
+)
+
+// PlanFunc is the function signature that you need to implement to add a custom CREATE DATABASE handler
+type PlanFunc = func(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Primitive, error)
+
+var databaseCreator = defaultCreateDatabase
+
+//goland:noinspection GoVarAndConstTypeMayBeOmitted
+var _ PlanFunc = databaseCreator
+
+func defaultCreateDatabase(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Primitive, error) {
+	return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "create database not allowed")
+}
diff --git a/go/vt/vtgate/planbuilder/create_database_pluggable_test.go b/go/vt/vtgate/planbuilder/create_database_pluggable_test.go
new file mode 100644
index 00000000000..83e4e68e0e0
--- /dev/null
+++ b/go/vt/vtgate/planbuilder/create_database_pluggable_test.go
@@ -0,0 +1,59 @@
+/*
+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"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/engine"
+
+	"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
+	_, err := TestBuilder("create database test", vschema)
+	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
+	s := &engine.SingleRow{}
+	databaseCreator = func(stmt sqlparser.Statement, vschema ContextVSchema) (engine.Primitive, error) {
+		return s, nil
+	}
+
+	output, err := TestBuilder("create database test", vschema)
+	require.NoError(t, err)
+	assert.Same(t, s, output.Instructions)
+}