初始化

2026-08-25 17:59:42 +08:00
commit 4b7380dd9b
408 changed files with 327400 additions and 0 deletions
@@ -0,0 +1,76 @@
+package database
+
+import (
+	_ "embed"
+	"fmt"
+
+	"juhe-factory/api/internal/model"
+
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+//go:embed schema.sql
+var baselineSchema string
+
+func bootstrapEmptyDatabase(db *gorm.DB) error {
+	var businessTableCount int64
+	if err := db.Raw(`SELECT count(*) FROM pg_tables
+		WHERE schemaname = 'public' AND tablename NOT IN ('jcf_schema_migrations', 'schema_migrations')`).
+		Scan(&businessTableCount).Error; err != nil {
+		return fmt.Errorf("inspect database schema: %w", err)
+	}
+	if businessTableCount == 0 {
+		baselineDB := db.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
+		if err := baselineDB.Transaction(func(tx *gorm.DB) error {
+			return tx.Exec(baselineSchema).Error
+		}); err != nil {
+			return fmt.Errorf("initialize database baseline: %w", err)
+		}
+		if err := db.Exec(`SET search_path TO public`).Error; err != nil {
+			return fmt.Errorf("restore database search path: %w", err)
+		}
+	}
+
+	if err := autoMigrateMissingTables(db, autoMigrateModels()...); err != nil {
+		return fmt.Errorf("auto migrate database models: %w", err)
+	}
+	return nil
+}
+
+func autoMigrateMissingTables(db *gorm.DB, tableModels ...any) error {
+	for _, tableModel := range tableModels {
+		if db.Migrator().HasTable(tableModel) {
+			continue
+		}
+		if err := db.AutoMigrate(tableModel); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func autoMigrateModels() []any {
+	return []any{
+		&model.AdminUser{},
+		&model.AdminRefreshToken{},
+		&model.AdminAuditLog{},
+		&model.WebUser{},
+		&model.WebRefreshToken{},
+		&model.MediaAsset{},
+		&model.CreativeProject{},
+		&model.ProjectModelConfig{},
+		&model.UserModelConfig{},
+		&model.ScriptAnalysis{},
+		&model.ScriptAnalysisCharacter{},
+		&model.ProjectEpisode{},
+		&model.ProjectAsset{},
+		&model.EpisodeStoryboard{},
+		&model.EpisodeSource{},
+		&model.DramaImportSession{},
+		&model.DramaParseBatch{},
+		&model.DramaParseTask{},
+		&model.GenerationTask{},
+		&model.GenerationOutput{},
+	}
+}
@@ -0,0 +1,59 @@
+package database
+
+import (
+	"os"
+	"testing"
+
+	"juhe-factory/api/internal/config"
+)
+
+type autoMigrateProbe struct {
+	ID uint `gorm:"primaryKey"`
+}
+
+func (autoMigrateProbe) TableName() string { return "auto_migrate_probe" }
+
+func TestBootstrapEmptyDatabase(t *testing.T) {
+	databaseURL := os.Getenv("JCF_BOOTSTRAP_TEST_DATABASE_URL")
+	if databaseURL == "" {
+		t.Skip("JCF_BOOTSTRAP_TEST_DATABASE_URL is not configured")
+	}
+
+	db, sqlDB, err := Open(config.Config{DatabaseURL: databaseURL})
+	if err != nil {
+		t.Fatalf("open empty database: %v", err)
+	}
+	defer sqlDB.Close()
+
+	var tableCount int64
+	if err := db.Raw(`SELECT count(*) FROM pg_tables WHERE schemaname = 'public'`).Scan(&tableCount).Error; err != nil {
+		t.Fatalf("count initialized tables: %v", err)
+	}
+	if tableCount != 33 {
+		t.Fatalf("expected 33 initialized tables, got %d", tableCount)
+	}
+
+	var migrationCount int64
+	if err := db.Table("jcf_schema_migrations").Count(&migrationCount).Error; err != nil {
+		t.Fatalf("count baseline migrations: %v", err)
+	}
+	if migrationCount != 0 {
+		t.Fatalf("expected no baseline migrations, got %d", migrationCount)
+	}
+
+	if err := autoMigrateMissingTables(db, &autoMigrateProbe{}); err != nil {
+		t.Fatalf("auto migrate missing table: %v", err)
+	}
+	if !db.Migrator().HasTable(&autoMigrateProbe{}) {
+		t.Fatal("expected AutoMigrate to create missing model table")
+	}
+	if err := db.Migrator().DropTable(&autoMigrateProbe{}); err != nil {
+		t.Fatalf("drop AutoMigrate probe table: %v", err)
+	}
+
+	_, reopenedSQLDB, err := Open(config.Config{DatabaseURL: databaseURL})
+	if err != nil {
+		t.Fatalf("reopen initialized database: %v", err)
+	}
+	defer reopenedSQLDB.Close()
+}
@@ -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
+}