feat: support dbName with operators in name (minus, ...)

Signed-off-by: Lucas Fernando Cardoso Nunes <lucasfc.nunes@gmail.com>
k3s-io · Sep 19, 2024 · 2cd50d0 · 2cd50d0
1 parent 70a99f8
commit 2cd50d0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/pkg/drivers/mysql/mysql.go b/pkg/drivers/mysql/mysql.go
@@ -49,7 +49,7 @@ var (
 	schemaMigrations = []string{
 		`ALTER TABLE kine MODIFY COLUMN id BIGINT UNSIGNED AUTO_INCREMENT NOT NULL UNIQUE, MODIFY COLUMN create_revision BIGINT UNSIGNED, MODIFY COLUMN prev_revision BIGINT UNSIGNED`,
 	}
-	createDB = "CREATE DATABASE IF NOT EXISTS "
+	createDB = "CREATE DATABASE IF NOT EXISTS `%s`;"
 )
 
 func New(ctx context.Context, dataSourceName string, tlsInfo tls.Config, connPoolConfig generic.ConnectionPoolConfig, metricsRegisterer prometheus.Registerer) (server.Backend, error) {
@@ -178,8 +178,11 @@ func createDBIfNotExist(dataSourceName string) error {
 		logrus.Warnf("failed to check existence of database %s, going to attempt create: %v", dbName, err)
 	}
 
+
 	if !exists {
-		if _, err = db.Exec(createDB + dbName); err != nil {
+		stmt := fmt.Sprintf(createDB, dbName)
+		logrus.Tracef("SETUP EXEC : %v", util.Stripped(stmt))
+		if _, err = db.Exec(stmt); err != nil {
 			if mysqlError, ok := err.(*mysql.MySQLError); !ok || mysqlError.Number != 1049 {
 				return err
 			}
@@ -188,7 +191,7 @@ func createDBIfNotExist(dataSourceName string) error {
 			if err != nil {
 				return err
 			}
-			if _, err = db.Exec(createDB + dbName); err != nil {
+			if _, err = db.Exec(stmt); err != nil {
 				return err
 			}
 		}

diff --git a/pkg/drivers/pgsql/pgsql.go b/pkg/drivers/pgsql/pgsql.go
@@ -50,7 +50,7 @@ var (
 	schemaMigrations = []string{
 		`ALTER TABLE kine ALTER COLUMN id SET DATA TYPE BIGINT, ALTER COLUMN create_revision SET DATA TYPE BIGINT, ALTER COLUMN prev_revision SET DATA TYPE BIGINT; ALTER SEQUENCE kine_id_seq AS BIGINT`,
 	}
-	createDB = "CREATE DATABASE "
+	createDB = `CREATE DATABASE "%s";`
 )
 
 func New(ctx context.Context, dataSourceName string, tlsInfo tls.Config, connPoolConfig generic.ConnectionPoolConfig, metricsRegisterer prometheus.Registerer) (server.Backend, error) {
@@ -165,9 +165,9 @@ func createDBIfNotExist(dataSourceName string) error {
 		logrus.Warnf("failed to check existence of database %s, going to attempt create: %v", dbName, err)
 	}
 
-	stmt := createDB + dbName + ";"
-
+
 	if !exists {
+		stmt := fmt.Sprintf(createDB, dbName)
 		logrus.Tracef("SETUP EXEC : %v", util.Stripped(stmt))
 		if _, err = db.Exec(stmt); err != nil {
 			logrus.Warnf("failed to create database %s: %v", dbName, err)