fix(mariadb): Fix restore when db name has special characters

clevyr · Nov 22, 2023 · 72ae32d · 72ae32d
1 parent b721420
commit 72ae32d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/internal/database/dialect/mariadb.go b/internal/database/dialect/mariadb.go
@@ -46,7 +46,8 @@ func (MariaDB) DefaultUser() string {
 	return "root"
 }
 
-func (MariaDB) DropDatabaseQuery(database string) string {
+func (db MariaDB) DropDatabaseQuery(database string) string {
+	database = db.quoteIdentifier(database)
 	return "set FOREIGN_KEY_CHECKS=0; create or replace database " + database + "; set FOREIGN_KEY_CHECKS=1; use " + database + ";"
 }
 
@@ -160,3 +161,9 @@ func (db MariaDB) DumpExtension(format sqlformat.Format) string {
 	}
 	return ""
 }
+
+func (db MariaDB) quoteIdentifier(param string) string {
+	param = strings.ReplaceAll(param, "`", "``")
+	param = "`" + param + "`"
+	return param
+}
diff --git a/internal/database/dialect/mariadb_test.go b/internal/database/dialect/mariadb_test.go
@@ -21,7 +21,7 @@ func TestMariaDB_DropDatabaseQuery(t *testing.T) {
 		args args
 		want string
 	}{
-		{"database", args{"database"}, "set FOREIGN_KEY_CHECKS=0; create or replace database database; set FOREIGN_KEY_CHECKS=1; use database;"},
+		{"database", args{"database"}, "set FOREIGN_KEY_CHECKS=0; create or replace database `database`; set FOREIGN_KEY_CHECKS=1; use `database`;"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -264,3 +264,23 @@ func TestMariaDB_FormatFromFilename(t *testing.T) {
 		})
 	}
 }
+
+func TestMariaDB_quoteIdentifier(t *testing.T) {
+	type args struct {
+		param string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{"simple", args{"table"}, "`table`"},
+		{"escaped", args{"T`able"}, "`T``able`"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db := MariaDB{}
+			assert.Equal(t, tt.want, db.quoteIdentifier(tt.args.param))
+		})
+	}
+}