初始化

2026-08-25 17:59:42 +08:00
commit 4b7380dd9b
408 changed files with 327400 additions and 0 deletions
@@ -0,0 +1,133 @@
+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
+}