check mysql db and table exist

Signed-off-by: yushiqie <[email protected]>
k3s-io · Jan 9, 2024 · dc1f6ea · dc1f6ea
1 parent f7ae7ce
commit dc1f6ea
Showing 1 changed file with 40 additions and 23 deletions.
diff --git a/pkg/drivers/mysql/mysql.go b/pkg/drivers/mysql/mysql.go
@@ -7,14 +7,15 @@ import (
 	"fmt"
 
 	"github.com/go-sql-driver/mysql"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/sirupsen/logrus"
+
 	"github.com/k3s-io/kine/pkg/drivers/generic"
 	"github.com/k3s-io/kine/pkg/logstructured"
 	"github.com/k3s-io/kine/pkg/logstructured/sqllog"
 	"github.com/k3s-io/kine/pkg/server"
 	"github.com/k3s-io/kine/pkg/tls"
 	"github.com/k3s-io/kine/pkg/util"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/sirupsen/logrus"
 )
 
 const (
@@ -24,7 +25,7 @@ const (
 
 var (
 	schema = []string{
-		`CREATE TABLE IF NOT EXISTS kine
+		`CREATE TABLE kine
 			(
 				id INTEGER AUTO_INCREMENT,
 				name VARCHAR(630) CHARACTER SET ascii,
@@ -43,7 +44,7 @@ var (
 		`CREATE INDEX kine_prev_revision_index ON kine (prev_revision)`,
 		`CREATE UNIQUE INDEX kine_name_prev_revision_uindex ON kine (name, prev_revision)`,
 	}
-	createDB = "CREATE DATABASE IF NOT EXISTS "
+	createDB = "CREATE DATABASE "
 )
 
 func New(ctx context.Context, dataSourceName string, tlsInfo tls.Config, connPoolConfig generic.ConnectionPoolConfig, metricsRegisterer prometheus.Registerer) (server.Backend, error) {
@@ -116,14 +117,21 @@ func New(ctx context.Context, dataSourceName string, tlsInfo tls.Config, connPoo
 }
 
 func setup(db *sql.DB) error {
-	logrus.Infof("Configuring database table schema and indexes, this may take a moment...")
-
-	for _, stmt := range schema {
-		logrus.Tracef("SETUP EXEC : %v", util.Stripped(stmt))
-		_, err := db.Exec(stmt)
-		if err != nil {
-			if mysqlError, ok := err.(*mysql.MySQLError); !ok || mysqlError.Number != 1061 {
-				return err
+	var exists bool
+	err := db.QueryRow("SELECT 1 FROM information_schema.TABLES WHERE table_schema = DATABASE() AND table_name = ?", "kine").Scan(&exists)
+	if err != nil && err != sql.ErrNoRows {
+		logrus.Warnf("failed to check existence of database table %s, going to attempt create: %v", "kine", err)
+	}
+
+	if !exists {
+		logrus.Infof("Configuring database table schema and indexes, this may take a moment...")
+		for _, stmt := range schema {
+			logrus.Tracef("SETUP EXEC : %v", util.Stripped(stmt))
+			_, err := db.Exec(stmt)
+			if err != nil {
+				if mysqlError, ok := err.(*mysql.MySQLError); !ok || mysqlError.Number != 1061 {
+					return err
+				}
 			}
 		}
 	}
@@ -143,19 +151,28 @@ func createDBIfNotExist(dataSourceName string) error {
 	if err != nil {
 		return err
 	}
-	_, err = db.Exec(createDB + dbName)
-	if err != nil {
-		if mysqlError, ok := err.(*mysql.MySQLError); !ok || mysqlError.Number != 1049 {
-			return err
-		}
-		config.DBName = ""
-		db, err = sql.Open("mysql", config.FormatDSN())
-		if err != nil {
-			return err
-		}
+
+	var exists bool
+	err = db.QueryRow("SELECT 1 FROM information_schema.SCHEMATA WHERE schema_name = ?", dbName).Scan(&exists)
+	if err != nil && err != sql.ErrNoRows {
+		logrus.Warnf("failed to check existence of database %s, going to attempt create: %v", dbName, err)
+	}
+
+	if !exists {
 		_, err = db.Exec(createDB + dbName)
 		if err != nil {
-			return err
+			if mysqlError, ok := err.(*mysql.MySQLError); !ok || mysqlError.Number != 1049 {
+				return err
+			}
+			config.DBName = ""
+			db, err = sql.Open("mysql", config.FormatDSN())
+			if err != nil {
+				return err
+			}
+			_, err = db.Exec(createDB + dbName)
+			if err != nil {
+				return err
+			}
 		}
 	}
 	return nil