mirror of
https://github.com/basketikun/infinite-canvas.git
synced 2026-08-05 17:04:27 +08:00
feat(database): add automatic database creation for MySQL and PostgreSQL
This commit is contained in:
+113
-2
@@ -1,6 +1,9 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -9,7 +12,10 @@ import (
|
||||
"github.com/basketikun/infinite-canvas/config"
|
||||
"github.com/basketikun/infinite-canvas/model"
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/driver/mysql"
|
||||
mysqldriver "github.com/go-sql-driver/mysql"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
gormmysql "gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -41,6 +47,18 @@ func DB() (*gorm.DB, error) {
|
||||
if driver == "sqlite" && dsn != ":memory:" {
|
||||
_ = os.MkdirAll(filepath.Dir(dsn), 0755)
|
||||
}
|
||||
if isPostgresDriver(driver) {
|
||||
dbErr = ensurePostgresDatabase(dsn)
|
||||
if dbErr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if driver == "mysql" {
|
||||
dbErr = ensureMySQLDatabase(dsn)
|
||||
if dbErr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
db, dbErr = gorm.Open(dialector(driver, dsn), &gorm.Config{})
|
||||
if dbErr != nil {
|
||||
return
|
||||
@@ -59,10 +77,103 @@ func DB() (*gorm.DB, error) {
|
||||
func dialector(driver string, dsn string) gorm.Dialector {
|
||||
switch driver {
|
||||
case "mysql":
|
||||
return mysql.Open(dsn)
|
||||
return gormmysql.Open(dsn)
|
||||
case "postgres", "postgresql":
|
||||
return postgres.Open(dsn)
|
||||
default:
|
||||
return sqlite.Open(dsn)
|
||||
}
|
||||
}
|
||||
|
||||
func isPostgresDriver(driver string) bool {
|
||||
return driver == "postgres" || driver == "postgresql"
|
||||
}
|
||||
|
||||
func ensureMySQLDatabase(dsn string) error {
|
||||
cfg, err := mysqldriver.ParseDSN(dsn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
target := strings.TrimSpace(cfg.DBName)
|
||||
if target == "" {
|
||||
return nil
|
||||
}
|
||||
ctx := context.Background()
|
||||
targetDB, err := sql.Open("mysql", dsn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = targetDB.PingContext(ctx)
|
||||
_ = targetDB.Close()
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if !isMySQLError(err, 1049) {
|
||||
return err
|
||||
}
|
||||
|
||||
maintenance := cfg.Clone()
|
||||
maintenance.DBName = ""
|
||||
serverDB, err := sql.Open("mysql", maintenance.FormatDSN())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer serverDB.Close()
|
||||
|
||||
_, err = serverDB.ExecContext(ctx, "CREATE DATABASE "+quoteMySQLIdentifier(target)+" CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
|
||||
if isMySQLError(err, 1007) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func ensurePostgresDatabase(dsn string) error {
|
||||
cfg, err := pgx.ParseConfig(dsn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
target := strings.TrimSpace(cfg.Database)
|
||||
if target == "" {
|
||||
return nil
|
||||
}
|
||||
ctx := context.Background()
|
||||
conn, err := pgx.ConnectConfig(ctx, cfg)
|
||||
if err == nil {
|
||||
_ = conn.Close(ctx)
|
||||
return nil
|
||||
}
|
||||
if !isPostgresError(err, "3D000") {
|
||||
return err
|
||||
}
|
||||
|
||||
maintenance := cfg.Copy()
|
||||
maintenance.Database = "postgres"
|
||||
if strings.EqualFold(target, "postgres") {
|
||||
maintenance.Database = "template1"
|
||||
}
|
||||
conn, err = pgx.ConnectConfig(ctx, maintenance)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close(ctx)
|
||||
|
||||
_, err = conn.Exec(ctx, "CREATE DATABASE "+pgx.Identifier{target}.Sanitize(), pgx.QueryExecModeExec)
|
||||
if isPostgresError(err, "42P04") {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func isMySQLError(err error, number uint16) bool {
|
||||
var mysqlErr *mysqldriver.MySQLError
|
||||
return errors.As(err, &mysqlErr) && mysqlErr.Number == number
|
||||
}
|
||||
|
||||
func isPostgresError(err error, code string) bool {
|
||||
var pgErr *pgconn.PgError
|
||||
return errors.As(err, &pgErr) && pgErr.Code == code
|
||||
}
|
||||
|
||||
func quoteMySQLIdentifier(name string) string {
|
||||
return "`" + strings.ReplaceAll(name, "`", "``") + "`"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user