fix sqlite init (#2344)

Co-authored-by: Yening Qin <710leo@gmail.com>

models/builtin_component.go

@@ -20,6 +20,17 @@ type BuiltinComponent struct {
 	UpdatedBy string `json:"updated_by" gorm:"type:varchar(191);not null;default:'';comment:'updater'"`
 }
 
+type PostgresBuiltinComponent struct {
+	ID        uint64 `json:"id" gorm:"primaryKey;type:bigint;autoIncrement;comment:'unique identifier'"`
+	Ident     string `json:"ident" gorm:"type:varchar(191);not null;uniqueIndex:idx_ident,sort:asc;comment:'identifier of component'"`
+	Logo      string `json:"logo" gorm:"type:text;comment:'logo of component'"`
+	Readme    string `json:"readme" gorm:"type:text;not null;comment:'readme of component'"`
+	CreatedAt int64  `json:"created_at" gorm:"type:bigint;not null;default:0;comment:'create time'"`
+	CreatedBy string `json:"created_by" gorm:"type:varchar(191);not null;default:'';comment:'creator'"`
+	UpdatedAt int64  `json:"updated_at" gorm:"type:bigint;not null;default:0;comment:'update time'"`
+	UpdatedBy string `json:"updated_by" gorm:"type:varchar(191);not null;default:'';comment:'updater'"`
+}
+
 func (bc *BuiltinComponent) TableName() string {
 	return "builtin_components"
 }

models/migrate/migrate.go

@@ -38,13 +38,22 @@ func MigrateIbexTables(db *gorm.DB) {
 
 	for i := 0; i < 100; i++ {
 		tableName := fmt.Sprintf("task_host_%d", i)
-		err := db.Table(tableName).AutoMigrate(&imodels.TaskHost{})
-		if err != nil {
-			logger.Errorf("failed to migrate table:%s %v", tableName, err)
+		exists := db.Migrator().HasTable(tableName)
+		if exists {
+			continue
+		} else {
+			err := db.Table(tableName).AutoMigrate(&imodels.TaskHost{})
+			if err != nil {
+				logger.Errorf("failed to migrate table:%s %v", tableName, err)
+			}
 		}
 	}
 }
 
+func isPostgres(db *gorm.DB) bool {
+	dialect := db.Dialector.Name()
+	return dialect == "postgres"
+}
 func MigrateTables(db *gorm.DB) error {
 	var tableOptions string
 	switch db.Dialector.(type) {
@@ -54,13 +63,18 @@ func MigrateTables(db *gorm.DB) error {
 	if tableOptions != "" {
 		db = db.Set("gorm:table_options", tableOptions)
 	}
-
 	dts := []interface{}{&RecordingRule{}, &AlertRule{}, &AlertSubscribe{}, &AlertMute{},
 		&TaskRecord{}, &ChartShare{}, &Target{}, &Configs{}, &Datasource{}, &NotifyTpl{},
 		&Board{}, &BoardBusigroup{}, &Users{}, &SsoConfig{}, &models.BuiltinMetric{},
-		&models.MetricFilter{}, &models.BuiltinComponent{}, &models.NotificaitonRecord{},
+		&models.MetricFilter{}, &models.NotificaitonRecord{},
 		&models.TargetBusiGroup{}}
 
+	if isPostgres(db) {
+		dts = append(dts, &models.PostgresBuiltinComponent{})
+	} else {
+		dts = append(dts, &models.BuiltinComponent{})
+	}
+
 	if !db.Migrator().HasColumn(&imodels.TaskSchedulerHealth{}, "scheduler") {
 		dts = append(dts, &imodels.TaskSchedulerHealth{})
 	}

pkg/ormx/database_init_test.go

@@ -0,0 +1,60 @@
+package ormx
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+func TestDataBaseInit(t *testing.T) {
+	tests := []struct {
+		name   string
+		config DBConfig
+	}{
+		{
+			name: "MySQL",
+			config: DBConfig{
+				DBType: "mysql",
+				DSN:    "root:1234@tcp(127.0.0.1:3306)/test?charset=utf8mb4&parseTime=True&loc=Local&allowNativePasswords=true",
+			},
+		},
+		{
+			name: "Postgres",
+			config: DBConfig{
+				DBType: "postgres",
+				DSN:    "host=127.0.0.1 port=5432 user=postgres dbname=test password=1234 sslmode=disable",
+			},
+		},
+		{
+			name: "SQLite",
+			config: DBConfig{
+				DBType: "sqlite",
+				DSN:    "./test.db",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := createDatabase(tt.config, &gorm.Config{})
+			assert.NoError(t, err)
+			var dialector gorm.Dialector
+			switch tt.config.DBType {
+			case "mysql":
+				dialector = mysql.Open(tt.config.DSN)
+			case "postgres":
+				dialector = postgres.Open(tt.config.DSN)
+			case "sqlite":
+				dialector = sqlite.Open(tt.config.DSN)
+			}
+			db, err := gorm.Open(dialector, &gorm.Config{})
+			assert.NoError(t, err)
+			err = DataBaseInit(tt.config, db)
+			assert.NoError(t, err)
+		})
+	}
+}

pkg/ormx/ormx.go

@@ -2,6 +2,7 @@ package ormx
 
 import (
 	"fmt"
+	"os"
 	"reflect"
 	"strings"
 	"time"
@@ -70,6 +71,234 @@ func (l *TKitLogger) Printf(s string, i ...interface{}) {
 	}
 }
 
+func createDatabase(c DBConfig, gconfig *gorm.Config) error {
+	switch strings.ToLower(c.DBType) {
+	case "mysql":
+		return createMysqlDatabase(c.DSN, gconfig)
+	case "postgres":
+		return createPostgresDatabase(c.DSN, gconfig)
+	case "sqlite":
+		return createSqliteDatabase(c.DSN, gconfig)
+	default:
+		return fmt.Errorf("dialector(%s) not supported", c.DBType)
+	}
+}
+
+func createSqliteDatabase(dsn string, gconfig *gorm.Config) error {
+	tempDialector := sqlite.Open(dsn)
+
+	_, err := gorm.Open(tempDialector, gconfig)
+	if err != nil {
+		return fmt.Errorf("failed to open temporary connection: %v", err)
+	}
+	
+	fmt.Println("sqlite file created")
+
+	return nil
+}
+
+func createPostgresDatabase(dsn string, gconfig *gorm.Config) error {
+	dsnParts := strings.Split(dsn, " ")
+	dbName := ""
+	connectionWithoutDB := ""
+	for _, part := range dsnParts {
+		if strings.HasPrefix(part, "dbname=") {
+			dbName = part[strings.Index(part, "=")+1:]
+		} else {
+			connectionWithoutDB += part
+			connectionWithoutDB += " "
+		}
+	}
+
+	createDBQuery := fmt.Sprintf("CREATE DATABASE %s ENCODING='UTF8' LC_COLLATE='en_US.UTF-8' LC_CTYPE='en_US.UTF-8';", dbName)
+
+	tempDialector := postgres.Open(connectionWithoutDB)
+
+	tempDB, err := gorm.Open(tempDialector, gconfig)
+	if err != nil {
+		return fmt.Errorf("failed to open temporary connection: %v", err)
+	}
+
+	result := tempDB.Exec(createDBQuery)
+	if result.Error != nil {
+		return fmt.Errorf("failed to execute create database query: %v", result.Error)
+	}
+
+	return nil
+}
+
+func createMysqlDatabase(dsn string, gconfig *gorm.Config) error {
+	dsnParts := strings.SplitN(dsn, "/", 2)
+	if len(dsnParts) != 2 {
+		return fmt.Errorf("failed to parse DSN: %s", dsn)
+	}
+
+	connectionInfo := dsnParts[0]
+	dbInfo := dsnParts[1]
+	dbName := dbInfo
+
+	queryIndex := strings.Index(dbInfo, "?")
+	if queryIndex != -1 {
+		dbName = dbInfo[:queryIndex]
+	} else {
+		return fmt.Errorf("failed to parse database name from DSN: %s", dsn)
+	}
+
+	connectionWithoutDB := connectionInfo + "/?" + dbInfo[queryIndex+1:]
+	createDBQuery := fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s CHARACTER SET utf8mb4", dbName)
+
+	tempDialector := mysql.Open(connectionWithoutDB)
+
+	tempDB, err := gorm.Open(tempDialector, gconfig)
+	if err != nil {
+		return fmt.Errorf("failed to open temporary connection: %v", err)
+	}
+
+	result := tempDB.Exec(createDBQuery)
+	if result.Error != nil {
+		return fmt.Errorf("failed to execute create database query: %v", result.Error)
+	}
+
+	return nil
+}
+
+func checkDatabaseExist(c DBConfig) (bool, error) {
+	switch strings.ToLower(c.DBType) {
+	case "mysql":
+		return checkMysqlDatabaseExist(c)
+	case "postgres":
+		return checkPostgresDatabaseExist(c)
+	case "sqlite":
+		return checkSqliteDatabaseExist(c)
+	default:
+		return false, fmt.Errorf("dialector(%s) not supported", c.DBType)
+	}
+
+}
+
+func checkSqliteDatabaseExist(c DBConfig) (bool, error) {
+	if _, err := os.Stat(c.DSN); os.IsNotExist(err) {
+		fmt.Printf("sqlite file not exists: %s\n", c.DSN)
+		return false, nil
+	} else {
+		return true, nil
+	}
+}
+
+func checkPostgresDatabaseExist(c DBConfig) (bool, error) {
+	dsnParts := strings.Split(c.DSN, " ")
+	dbName := ""
+	connectionWithoutDB := ""
+	for _, part := range dsnParts {
+		if strings.HasPrefix(part, "dbname=") {
+			dbName = part[strings.Index(part, "=")+1:]
+		} else {
+			connectionWithoutDB += part
+			connectionWithoutDB += " "
+		}
+	}
+
+	dialector := postgres.Open(connectionWithoutDB)
+
+	gconfig := &gorm.Config{
+		NamingStrategy: schema.NamingStrategy{
+			TablePrefix:   c.TablePrefix,
+			SingularTable: true,
+		},
+		Logger: gormLogger,
+	}
+
+	db, err := gorm.Open(dialector, gconfig)
+	if err != nil {
+		return false, fmt.Errorf("failed to open database: %v", err)
+	}
+
+    var databases []string
+    query := genQuery(c)
+    if err := db.Raw(query).Scan(&databases).Error; err != nil {
+        return false, fmt.Errorf("failed to query: %v", err)
+    }
+
+    for _, database := range databases {
+        if database == dbName {
+			fmt.Println("Database exist")
+            return true, nil
+        }
+    }
+
+    return false, nil
+}
+
+func checkMysqlDatabaseExist(c DBConfig) (bool, error) {
+	dsnParts := strings.SplitN(c.DSN, "/", 2)
+    if len(dsnParts) != 2 {
+        return false, fmt.Errorf("failed to parse DSN: %s", c.DSN)
+    }
+
+    connectionInfo := dsnParts[0]
+    dbInfo := dsnParts[1]
+    dbName := dbInfo
+
+    queryIndex := strings.Index(dbInfo, "?")
+    if queryIndex != -1 {
+        dbName = dbInfo[:queryIndex]
+    } else {
+        return false, fmt.Errorf("failed to parse database name from DSN: %s", c.DSN)
+    }
+
+    connectionWithoutDB := connectionInfo + "/?" + dbInfo[queryIndex+1:]
+
+    var dialector gorm.Dialector
+    switch strings.ToLower(c.DBType) {
+    case "mysql":
+        dialector = mysql.Open(connectionWithoutDB)
+    case "postgres":
+        dialector = postgres.Open(connectionWithoutDB)
+    default:
+        return false, fmt.Errorf("unsupported database type: %s", c.DBType)
+    }
+
+    gconfig := &gorm.Config{
+        NamingStrategy: schema.NamingStrategy{
+            TablePrefix:   c.TablePrefix,
+            SingularTable: true,
+        },
+        Logger: gormLogger,
+    }
+
+    db, err := gorm.Open(dialector, gconfig)
+    if err != nil {
+        return false, fmt.Errorf("failed to open database: %v", err)
+    }
+
+    var databases []string
+    query := genQuery(c)
+    if err := db.Raw(query).Scan(&databases).Error; err != nil {
+        return false, fmt.Errorf("failed to query: %v", err)
+    }
+
+    for _, database := range databases {
+        if database == dbName {
+            return true, nil
+        }
+    }
+
+    return false, nil
+}
+
+func genQuery(c DBConfig) string {
+	switch strings.ToLower(c.DBType) {
+    case "mysql":
+        return "SHOW DATABASES"
+    case "postgres":
+        return "SELECT datname FROM pg_database"
+    case "sqlite":
+        return ""
+    default:
+        return ""
+	}
+}
+
 // New Create gorm.DB instance
 func New(c DBConfig) (*gorm.DB, error) {
 	var dialector gorm.Dialector
@@ -95,9 +324,30 @@ func New(c DBConfig) (*gorm.DB, error) {
 		Logger: gormLogger,
 	}
 
+	dbExist, checkErr := checkDatabaseExist(c)
+	if checkErr != nil {
+		return nil, checkErr
+	}
+	if !dbExist {
+		fmt.Println("Database not exist, trying to create it")
+		createErr := createDatabase(c, gconfig)
+		if createErr != nil {
+			return nil, fmt.Errorf("failed to create database: %v", createErr)
+		}
+
+		db, err := gorm.Open(dialector, gconfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to reopen database after creation: %v", err)
+		}
+		err = DataBaseInit(c, db)
+		if err != nil {
+			return nil, fmt.Errorf("failed to init database: %v", err)
+		}
+	}
+
 	db, err := gorm.Open(dialector, gconfig)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to open database: %v", err)
 	}
 
 	if c.Debug {