77 lines
2.1 KiB
Go
77 lines
2.1 KiB
Go
package models
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var DBService *gorm.DB
|
|
|
|
// migrationAdvisoryLockKey scopes the PostgreSQL session lock to SenlinAI schema changes.
|
|
const migrationAdvisoryLockKey int64 = 0x53454e4c494e4149
|
|
|
|
func New(dsn string) error {
|
|
// 启用 Gorm 跨驱动错误翻译，业务层无需解析 Postgres 错误文本或 SQLSTATE。
|
|
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{TranslateError: true})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := AutoMigrate(db); err != nil {
|
|
return err
|
|
}
|
|
|
|
DBService = db
|
|
return nil
|
|
}
|
|
|
|
func AutoMigrate(database *gorm.DB) error {
|
|
if !migrationAdvisoryLockRequired(database.Dialector.Name()) {
|
|
return autoMigrateUnlocked(database)
|
|
}
|
|
// Connection pins every migration statement to one PostgreSQL session so the
|
|
// advisory lock covers both Gorm's schema work and the versioned data repairs.
|
|
return database.Connection(func(connection *gorm.DB) (resultErr error) {
|
|
if err := connection.Exec("SELECT pg_advisory_lock(?)", migrationAdvisoryLockKey).Error; err != nil {
|
|
return fmt.Errorf("acquire schema migration advisory lock: %w", err)
|
|
}
|
|
defer func() {
|
|
if err := connection.Exec("SELECT pg_advisory_unlock(?)", migrationAdvisoryLockKey).Error; resultErr == nil && err != nil {
|
|
resultErr = fmt.Errorf("release schema migration advisory lock: %w", err)
|
|
}
|
|
}()
|
|
return autoMigrateUnlocked(connection)
|
|
})
|
|
}
|
|
|
|
func migrationAdvisoryLockRequired(dialect string) bool {
|
|
return dialect == "postgres"
|
|
}
|
|
|
|
func autoMigrateUnlocked(database *gorm.DB) error {
|
|
if err := database.AutoMigrate(
|
|
&SenlinAgentSchemaMigration{},
|
|
&SenlinAgentUser{},
|
|
&SenlinAgentProject{},
|
|
&SenlinAgentInboxItem{},
|
|
&SenlinAgentInboxSuggestion{},
|
|
&SenlinAgentTask{},
|
|
&SenlinAgentNote{},
|
|
&SenlinAgentSource{},
|
|
&SenlinAgentAISession{},
|
|
&SenlinAgentTag{},
|
|
&SenlinAgentProjectChannel{},
|
|
&SenlinAgentCronPlan{},
|
|
&SenlinAgentProjectEvent{},
|
|
&SenlinAgentAIKey{},
|
|
&SenlinAgentAICallLog{},
|
|
&SenlinAgentAIRateBucket{},
|
|
&SenlinAgentTaskShare{},
|
|
); err != nil {
|
|
return err
|
|
}
|
|
return runVersionedMigrations(database)
|
|
}
|