package database import ( "context" "database/sql" "strings" "time" "juhe-factory/api/internal/config" "gorm.io/driver/postgres" "gorm.io/gorm" "gorm.io/gorm/logger" ) func Open(cfg config.Config) (*gorm.DB, *sql.DB, error) { logLevel := logger.Warn if cfg.Debug { logLevel = logger.Info } databaseURL := cfg.DatabaseURL if !strings.Contains(strings.ToLower(databaseURL), "timezone=") { separator := "?" if strings.Contains(databaseURL, " ") { separator = " " } if strings.Contains(databaseURL, "?") { separator = "&" } databaseURL += separator + "TimeZone=Asia%2FShanghai" } db, err := gorm.Open(postgres.Open(databaseURL), &gorm.Config{ Logger: logger.Default.LogMode(logLevel), }) if err != nil { return nil, nil, err } sqlDB, err := db.DB() if err != nil { return nil, nil, err } sqlDB.SetMaxOpenConns(20) sqlDB.SetMaxIdleConns(5) sqlDB.SetConnMaxLifetime(30 * time.Minute) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := sqlDB.PingContext(ctx); err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`SET TIME ZONE 'Asia/Shanghai'`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := bootstrapEmptyDatabase(db); err != nil { sqlDB.Close() return nil, nil, err } // Keep older installations compatible when migrations are applied manually. if err := db.Exec(`ALTER TABLE project_model_configs ADD COLUMN IF NOT EXISTS settings jsonb NOT NULL DEFAULT '{}'::jsonb`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`ALTER TABLE creative_projects ADD COLUMN IF NOT EXISTS video_resolution varchar(8) NOT NULL DEFAULT '720p'`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`ALTER TABLE episode_storyboards ADD COLUMN IF NOT EXISTS image_prompt text`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`ALTER TABLE prompts ADD COLUMN IF NOT EXISTS type varchar(50); ALTER TABLE prompts ADD COLUMN IF NOT EXISTS content text; ALTER TABLE prompts ADD COLUMN IF NOT EXISTS scope varchar(20) NOT NULL DEFAULT 'system'; ALTER TABLE prompts ADD COLUMN IF NOT EXISTS owner_user_id uuid REFERENCES web_users(id);`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`ALTER TABLE web_users ADD COLUMN IF NOT EXISTS account citext; UPDATE web_users SET account=username WHERE account IS NULL; CREATE UNIQUE INDEX IF NOT EXISTS idx_web_users_account ON web_users(account);`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`ALTER TABLE redemption_codes ADD COLUMN IF NOT EXISTS code character varying(40);`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`CREATE TABLE IF NOT EXISTS user_model_configs ( id uuid PRIMARY KEY DEFAULT gen_random_uuid(), user_id uuid NOT NULL REFERENCES web_users(id), purpose varchar(32) NOT NULL CHECK (purpose IN ('prompt_reverse','image_generation','video_generation')), model_id uuid NOT NULL REFERENCES models(id), prompt_id uuid REFERENCES prompts(id), settings jsonb NOT NULL DEFAULT '{}'::jsonb, created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(user_id, purpose) ); CREATE INDEX IF NOT EXISTS idx_user_model_configs_user ON user_model_configs(user_id); DO $$ BEGIN IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='user_model_configs' AND column_name='model_type') THEN INSERT INTO user_model_configs(id,user_id,project_type,model_type,purpose,model_id,prompt_id,settings) SELECT DISTINCT ON (project.user_id, config.model_type) gen_random_uuid(), project.user_id, 'video_redraw', config.model_type, config.purpose, config.model_id, config.prompt_id, config.settings FROM project_model_configs config JOIN creative_projects project ON project.id=config.project_id WHERE project.deleted_at IS NULL AND project.project_type='video_redraw' ORDER BY project.user_id, config.model_type, project.updated_at DESC, config.updated_at DESC ON CONFLICT DO NOTHING; ELSE INSERT INTO user_model_configs(id,user_id,purpose,model_id,prompt_id,settings) SELECT DISTINCT ON (project.user_id, config.purpose) gen_random_uuid(), project.user_id, config.purpose, config.model_id, config.prompt_id, config.settings FROM project_model_configs config JOIN creative_projects project ON project.id=config.project_id WHERE project.deleted_at IS NULL ORDER BY project.user_id, config.purpose, project.updated_at DESC, config.updated_at DESC ON CONFLICT DO NOTHING; END IF; END $$;`).Error; err != nil { sqlDB.Close() return nil, nil, err } if err := db.Exec(`CREATE TABLE IF NOT EXISTS user_prompt_preferences ( id uuid PRIMARY KEY DEFAULT gen_random_uuid(), user_id uuid NOT NULL REFERENCES web_users(id) ON DELETE CASCADE, prompt_id uuid NOT NULL REFERENCES prompts(id) ON DELETE CASCADE, prompt_type varchar(64) NOT NULL, created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(user_id,prompt_type) ); CREATE INDEX IF NOT EXISTS idx_user_prompt_preferences_user ON user_prompt_preferences(user_id);`).Error; err != nil { sqlDB.Close() return nil, nil, err } return db, sqlDB, nil }