From bd9c012184f404a05304a4e8365efa6599f1dfe1 Mon Sep 17 00:00:00 2001
From: Jason Mo <mohangjie1995@gmail.com>
Date: Tue, 6 Aug 2024 13:27:05 +0800
Subject: [PATCH 1/2] drop created schemas

---
 src/main.go | 42 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/main.go b/src/main.go
index 9ddaf45..c929a58 100644
--- a/src/main.go
+++ b/src/main.go
@@ -102,6 +102,8 @@ type tester struct {
 	mdb  *sql.DB
 	name string
 
+	originalSchemas map[string]struct{}
+
 	curr *Conn
 
 	buf bytes.Buffer
@@ -274,6 +276,20 @@ func (t *tester) preProcess() {
 		log.Fatalf("Open db err %v", err)
 	}
 
+	if !reserveSchema {
+		// store origial schemas
+		t.originalSchemas = make(map[string]struct{})
+		rows, err := mdb.Query("show databases")
+		if err != nil {
+			log.Errorf("failed to get databases: %s", err.Error())
+			return
+		}
+		for rows.Next() {
+			rows.Scan(&dbName)
+			t.originalSchemas[dbName] = struct{}{}
+		}
+	}
+
 	dbName = strings.ReplaceAll(t.name, "/", "__")
 	log.Debugf("Create new db `%s`", dbName)
 	if _, err = mdb.Exec(fmt.Sprintf("create database `%s`", dbName)); err != nil {
@@ -290,16 +306,30 @@ func (t *tester) preProcess() {
 }
 
 func (t *tester) postProcess() {
+	defer func() {
+		for _, v := range t.conn {
+			v.conn.Close()
+		}
+		t.mdb.Close()
+	}()
 	if !reserveSchema {
-		_, err := t.mdb.Exec(fmt.Sprintf("drop database `%s`", strings.ReplaceAll(t.name, "/", "__")))
+		rows, err := t.mdb.Query("show databases")
 		if err != nil {
-			log.Errorf("failed to drop database: %s", err.Error())
+			log.Errorf("failed to get databases: %s", err.Error())
+			return
+		}
+		var dbName string
+		for rows.Next() {
+			rows.Scan(&dbName)
+			if _, exists := t.originalSchemas[dbName]; !exists {
+				_, err := t.mdb.Exec(fmt.Sprintf("drop database `%s`", dbName))
+				if err != nil {
+					log.Errorf("failed to drop database: %s", err.Error())
+					return
+				}
+			}
 		}
 	}
-	for _, v := range t.conn {
-		v.conn.Close()
-	}
-	t.mdb.Close()
 }
 
 func (t *tester) addFailure(testSuite *XUnitTestSuite, err *error, cnt int) {

From 2625b51fbc9d491a9b78088b5f5672b83a50c2d2 Mon Sep 17 00:00:00 2001
From: Jason Mo <mohangjie1995@gmail.com>
Date: Tue, 6 Aug 2024 13:27:51 +0800
Subject: [PATCH 2/2] update

---
 src/main.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main.go b/src/main.go
index c929a58..b49d0fd 100644
--- a/src/main.go
+++ b/src/main.go
@@ -277,7 +277,7 @@ func (t *tester) preProcess() {
 	}
 
 	if !reserveSchema {
-		// store origial schemas
+		// store original schemas
 		t.originalSchemas = make(map[string]struct{})
 		rows, err := mdb.Query("show databases")
 		if err != nil {