You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
120 lines
3.9 KiB
120 lines
3.9 KiB
package db
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"time"
|
|
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"zestack.dev/env"
|
|
)
|
|
|
|
func Init() error {
|
|
var dialector gorm.Dialector
|
|
switch driver := env.String("DB_DRIVER", "postgres"); driver {
|
|
case "mysql":
|
|
schemaHelper = &mysqlSchemaHelper{}
|
|
dialector = openMysql()
|
|
case "postgres":
|
|
schemaHelper = &pgsqlSchemaHelper{}
|
|
dialector = openPostgres()
|
|
default:
|
|
return fmt.Errorf("unsupported database driver: %s", driver)
|
|
}
|
|
|
|
var err error
|
|
db, err = gorm.Open(dialector, &gorm.Config{
|
|
NamingStrategy: schema.NamingStrategy{
|
|
TablePrefix: env.String("DB_PREFIX"),
|
|
SingularTable: env.Bool("DB_SINGULAR_TABLE", false),
|
|
IdentifierMaxLength: env.Int("DB_IDENTIFIER_MAX_LENGTH", 0),
|
|
},
|
|
Logger: &dbLogger{200 * time.Millisecond},
|
|
// 在我们使用模型查询时，使用模型字段而不是通配符来查询。
|
|
// https://gorm.io/docs/advanced_query.html#Smart-Select-Fields
|
|
QueryFields: env.Bool("DB_QUERY_FIELDS", true),
|
|
// 在 AutoMigrate 或 CreateTable 时，GORM 会自动创建外键约束，
|
|
// 若要禁用该特性，可将其设置为 true 。
|
|
// https://gorm.io/docs/migration.html
|
|
DisableForeignKeyConstraintWhenMigrating: env.Bool("DB_DISABLE_FOREIGN_KEY_CONSTRAINT", false),
|
|
// 在 AutoMigrate 或 CreateTable 时，GORM 会自动创建索引，
|
|
// 若要禁用该特性，可将其设置为 true 。
|
|
IgnoreRelationshipsWhenMigrating: env.Bool("DB_IGNORE_RELATIONSHIPS", false),
|
|
// 开启方言错误转换，GORM 会将不同数据库的特定错误转换为常见的 GORM 错误类型，
|
|
// 这样，方便我们对错误进行统一处理。
|
|
// https://gorm.io/docs/error_handling.html#Dialect-Translated-Errors
|
|
// https://github.com/go-gorm/gorm/blob/master/errors.go
|
|
TranslateError: true,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 如果使用的是 MySQL 数据库，则启用 InnoDB 引擎。
|
|
if db.Dialector.Name() == "mysql" {
|
|
db.Set("gorm:table_options", "ENGINE=InnoDB")
|
|
}
|
|
|
|
// 配置连接池
|
|
var raw *sql.DB
|
|
if raw, err = db.DB(); err == nil {
|
|
err = configRawDB(raw)
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
func openMysql() gorm.Dialector {
|
|
return mysql.New(mysql.Config{
|
|
DSN: fmt.Sprintf(
|
|
"%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=%s",
|
|
env.String("DB_USER", "root"),
|
|
env.String("DB_AUTH", "password"),
|
|
env.String("DB_HOST", "localhost"),
|
|
env.Int("DB_PORT", 3306),
|
|
env.String("DB_NAME", "test"),
|
|
env.String("DB_CHARSET", "utf8mb4"),
|
|
time.Local.String(),
|
|
),
|
|
// string 类型字段的默认长度
|
|
DefaultStringSize: uint(env.Int("DB_STRING_SIZE", 256)),
|
|
// 禁用 datetime 精度，MySQL 5.6 之前的数据库不支持
|
|
DisableDatetimePrecision: true,
|
|
// 重命名索引时采用删除并新建的方式，MySQL 5.7 之前的数据库和 MariaDB 不支持重命名索引
|
|
DontSupportRenameIndex: true,
|
|
// 用 `change` 重命名列，MySQL 8 之前的数据库和 MariaDB 不支持重命名列
|
|
DontSupportRenameColumn: true,
|
|
})
|
|
}
|
|
|
|
func openPostgres() gorm.Dialector {
|
|
return postgres.Open(fmt.Sprintf(
|
|
"host=%s user=%s password=%s dbname=%s port=%d sslmode=%s TimeZone=%s",
|
|
env.String("DB_HOST", "localhost"),
|
|
env.String("DB_USER", "postgres"),
|
|
env.String("DB_AUTH", "password"),
|
|
env.String("DB_NAME", "postgres"),
|
|
env.Int("DB_PORT", 5432),
|
|
env.String("DB_SSLMODE", "disable"),
|
|
time.Local.String(),
|
|
))
|
|
}
|
|
|
|
func configRawDB(db *sql.DB) error {
|
|
// 用于设置连接池中空闲连接的最大数量。
|
|
if maxIdleConns := env.Int("DB_MAX_IDLE_CONNS", 0); maxIdleConns > 0 {
|
|
db.SetMaxIdleConns(maxIdleConns)
|
|
}
|
|
// 设置打开数据库连接的最大数量。
|
|
if maxOpenConns := env.Int("DB_MAX_OPEN_CONNS", 0); maxOpenConns > 0 {
|
|
db.SetMaxOpenConns(maxOpenConns)
|
|
}
|
|
// 设置了连接可复用的最大时间。
|
|
if connMaxLifetime := env.Duration("DB_CONN_MAX_LIFETIME", 0); connMaxLifetime > 0 {
|
|
db.SetConnMaxLifetime(connMaxLifetime)
|
|
}
|
|
return db.Ping()
|
|
}
|
|
|