JuYou/API/internal/database/postgres.go

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
}