From b020b41352e5854953320147397c509baceec640 Mon Sep 17 00:00:00 2001
From: newborn22 <953950914@qq.com>
Date: Mon, 16 Dec 2024 00:57:59 +0800
Subject: [PATCH] fix: lock when vscham delete and add

---
 go/vt/vtgate/engine/dbddl.go             |  6 +++++-
 go/vt/vtgate/engine/fake_vcursor_test.go | 16 ++++++++++++++++
 go/vt/vtgate/engine/primitive.go         |  8 ++++++--
 go/vt/vtgate/executor.go                 | 14 ++++++++++++++
 go/vt/vtgate/vcursor_impl.go             | 10 ++++++++++
 5 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/go/vt/vtgate/engine/dbddl.go b/go/vt/vtgate/engine/dbddl.go
index 6d76ba4a0b..12e8a9136d 100644
--- a/go/vt/vtgate/engine/dbddl.go
+++ b/go/vt/vtgate/engine/dbddl.go
@@ -164,7 +164,8 @@ func (c *DBDDL) createDatabase(ctx context.Context, vcursor VCursor, plugin DBDD
 	}
 
 	newKSSchema := &vindexes.KeyspaceSchema{Keyspace: &vindexes.Keyspace{Name: c.name, Sharded: false}, Tables: make(map[string]*vindexes.Table)}
-	vcursor.GetExecutorVSchema().Keyspaces[c.name] = newKSSchema
+	vcursor.VSchemaAddKeyspaceIfNotExists(c.name, newKSSchema)
+	//vcursor.GetExecutorVSchema().Keyspaces[c.name] = newKSSchema
 	//vcursor.GetVSchema().Keyspaces[c.name] = newKSSchema
 
 	return &sqltypes.Result{RowsAffected: 1}, nil
@@ -183,6 +184,9 @@ func (c *DBDDL) dropDatabase(ctx context.Context, vcursor VCursor, plugin DBDDLP
 		}
 	}
 
+	vcursor.VSchemaDeleteKeyspace(c.name)
+	//delete(vcursor.GetExecutorVSchema().Keyspaces, c.name)
+
 	return &sqltypes.Result{StatusFlags: sqltypes.ServerStatusDbDropped}, nil
 }
 
diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go
index bd9b3229b3..022d41d2ee 100644
--- a/go/vt/vtgate/engine/fake_vcursor_test.go
+++ b/go/vt/vtgate/engine/fake_vcursor_test.go
@@ -110,6 +110,14 @@ func (t *noopVCursor) GetExecutorVSchema() *vindexes.VSchema {
 	panic("implement me")
 }
 
+func (t *noopVCursor) VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema) {
+	panic("implement me")
+}
+
+func (t *noopVCursor) VSchemaDeleteKeyspace(name string) {
+	panic("implement me")
+}
+
 func (t *noopVCursor) GetVSchema() *vindexes.VSchema {
 	// TODO implement me
 	panic("implement me")
@@ -836,6 +844,14 @@ func (f *loggingVCursor) GetExecutorVSchema() *vindexes.VSchema {
 	panic("implement me")
 }
 
+func (f *loggingVCursor) VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema) {
+	panic("implement me")
+}
+
+func (f *loggingVCursor) VSchemaDeleteKeyspace(name string) {
+	panic("implement me")
+}
+
 func (f *loggingVCursor) GetVSchema() {
 	panic("implement me")
 }
diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
index ef6ee1d6e1..43431921a0 100644
--- a/go/vt/vtgate/engine/primitive.go
+++ b/go/vt/vtgate/engine/primitive.go
@@ -129,9 +129,13 @@ type (
 
 		FindHealthyPrimaryTablet() (*discovery.TabletHealth, error)
 
-		GetExecutorVSchema() *vindexes.VSchema
+		//GetExecutorVSchema() *vindexes.VSchema
 
-		GetVSchema() *vindexes.VSchema
+		VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema)
+
+		VSchemaDeleteKeyspace(name string)
+
+		//GetVSchema() *vindexes.VSchema
 	}
 
 	// SessionActions gives primitives ability to interact with the session state
diff --git a/go/vt/vtgate/executor.go b/go/vt/vtgate/executor.go
index 663a5ecada..5a032063f9 100644
--- a/go/vt/vtgate/executor.go
+++ b/go/vt/vtgate/executor.go
@@ -1168,6 +1168,20 @@ func (e *Executor) VSchema() *vindexes.VSchema {
 	return e.vschema
 }
 
+func (e *Executor) VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	if _, ok := e.vschema.Keyspaces[name]; !ok {
+		e.vschema.Keyspaces[name] = KeyspaceSchema
+	}
+}
+
+func (e *Executor) VSchemaDeleteKeyspace(name string) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	delete(e.vschema.Keyspaces, name)
+}
+
 // SaveVSchema updates the vschema and stats
 func (e *Executor) SaveVSchema(vschema *vindexes.VSchema, stats *VSchemaStats) {
 	e.mu.Lock()
diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go
index 4aedda7819..1a32146f0a 100644
--- a/go/vt/vtgate/vcursor_impl.go
+++ b/go/vt/vtgate/vcursor_impl.go
@@ -93,6 +93,8 @@ type iExecute interface {
 	ParseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.Destination, error)
 	reloadExec(ctx context.Context, reloadType *sqlparser.ReloadType) error
 	VSchema() *vindexes.VSchema
+	VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema)
+	VSchemaDeleteKeyspace(name string)
 	SetFailPoint(command string, key string, value string) error
 	SubmitDMLJob(command, sql, uuid, tableSchema, timePeriodStart, timePeriodEnd, timePeriodTimeZone string, timeGapInMs, batchSize int64, postponeLaunch bool, failPolicy, throttleDuration, throttleRatio string) (*sqltypes.Result, error)
 	ShowDMLJob(uuid string, showDetail bool) (*sqltypes.Result, error)
@@ -1220,6 +1222,14 @@ func (vc *vcursorImpl) GetExecutorVSchema() *vindexes.VSchema {
 	return vc.executor.VSchema()
 }
 
+func (vc *vcursorImpl) VSchemaAddKeyspaceIfNotExists(name string, KeyspaceSchema *vindexes.KeyspaceSchema) {
+	vc.executor.VSchemaAddKeyspaceIfNotExists(name, KeyspaceSchema)
+}
+
+func (vc *vcursorImpl) VSchemaDeleteKeyspace(name string) {
+	vc.executor.VSchemaDeleteKeyspace(name)
+}
+
 func (vc *vcursorImpl) GetVSchema() *vindexes.VSchema {
 	return vc.vschema
 }