agent/backend/internal/initdb/dataset.go

package initdb

import (
	"errors"
	"fmt"
	"strings"

	"gorm.io/gorm"
	"senlinai-agent/backend/internal/models"
)

type defaultDatasetSource struct {
	Key        string
	Name       string
	URL        string
	IconURL    string
	LegacyURLs []string
}

var defaultDatasetSources = []defaultDatasetSource{
	{Key: "buzzing-bloomberg", Name: "彭博社最新报道", URL: "https://bbg.buzzing.cc/feed.json", IconURL: "https://bbg.buzzing.cc/apple-touch-icon.png"},
	{Key: "buzzing-tech", Name: "国外科技头条", URL: "https://tech.buzzing.cc/feed.json", IconURL: "https://tech.buzzing.cc/apple-touch-icon.png"},
	{Key: "buzzing-finance", Name: "国外财经新闻", URL: "https://finance.buzzing.cc/feed.json", IconURL: "https://finance.buzzing.cc/apple-touch-icon.png"},
	{Key: "buzzing-wsj", Name: "华尔街日报热门", URL: "https://wsj.buzzing.cc/feed.json", IconURL: "https://wsj.buzzing.cc/apple-touch-icon.png"},
	{Key: "buzzing-product-hunt", Name: "Product Hunt", URL: "https://ph.buzzing.cc/feed.json", IconURL: "https://ph.buzzing.cc/apple-touch-icon.png"},
	{Key: "buzzing-devto", Name: "Dev.to", URL: "https://dev.buzzing.cc/feed.json", IconURL: "https://dev.buzzing.cc/apple-touch-icon.png"},
	{
		Key: "rsshub-cls-hot", Name: "财联社-热门",
		URL:        "https://rsshub.ktachibana.party/cls/hot",
		LegacyURLs: []string{"https://rsshub.app/cls/hot"},
	},
	{
		Key: "rsshub-jin10", Name: "金十数据-快讯",
		URL:        "https://rsshub.ktachibana.party/jin10",
		LegacyURLs: []string{"https://rsshub.app/jin10"},
	},
}

// InitDataset inserts or repairs root's built-in dataset sources without touching user-created sources.
func InitDataset(database *gorm.DB) error {
	var root models.SaUser
	if err := database.Where("email = ? AND role = ?", rootUsername, rootRole).First(&root).Error; err != nil {
		return fmt.Errorf("find root user for default dataset sources: %w", err)
	}

	return database.Transaction(func(tx *gorm.DB) error {
		if err := prepareDatasetSourceOwners(tx); err != nil {
			return err
		}
		if err := prepareDatasetItems(tx); err != nil {
			return err
		}
		for _, definition := range defaultDatasetSources {
			if err := upsertDefaultDatasetSource(tx, root, definition); err != nil {
				return err
			}
		}
		return nil
	})
}

func prepareDatasetSourceOwners(tx *gorm.DB) error {
	hasCreatedBy, err := hasDatasetSourceColumn(tx, "created_by")
	if err != nil {
		return err
	}
	if hasCreatedBy {
		if err := tx.Exec(
			"UPDATE sa_dataset_sources SET owner_id = created_by WHERE owner_id = 0",
		).Error; err != nil {
			return fmt.Errorf("backfill dataset source owners: %w", err)
		}
	}

	var sources []models.SaDatasetSource
	if err := tx.Where("owner_id = ? OR owner_identity = ?", 0, "").Find(&sources).Error; err != nil {
		return err
	}
	for _, source := range sources {
		var owner models.SaUser
		if err := tx.Select("identity").First(&owner, source.OwnerID).Error; err != nil {
			return err
		}
		if err := tx.Model(&source).Updates(map[string]any{
			"owner_identity": owner.Identity,
		}).Error; err != nil {
			return err
		}
	}

	if tx.Migrator().HasIndex(&models.SaDatasetSource{}, "idx_sa_dataset_source_seed") {
		if err := tx.Migrator().DropIndex(&models.SaDatasetSource{}, "idx_sa_dataset_source_seed"); err != nil {
			return fmt.Errorf("drop legacy dataset source seed index: %w", err)
		}
	}
	for _, index := range []string{
		"idx_sa_dataset_sources_created_by",
		"idx_sa_dataset_sources_created_by_identity",
	} {
		if tx.Migrator().HasIndex(&models.SaDatasetSource{}, index) {
			if err := tx.Migrator().DropIndex(&models.SaDatasetSource{}, index); err != nil {
				return fmt.Errorf("drop legacy dataset source creator index %s: %w", index, err)
			}
		}
	}
	hasCreatedByIdentity, err := hasDatasetSourceColumn(tx, "created_by_identity")
	if err != nil {
		return err
	}
	if hasCreatedByIdentity {
		if err := tx.Exec(
			"ALTER TABLE sa_dataset_sources DROP COLUMN created_by_identity",
		).Error; err != nil {
			return fmt.Errorf("drop legacy dataset source creator identity: %w", err)
		}
	}
	hasCreatedBy, err = hasDatasetSourceColumn(tx, "created_by")
	if err != nil {
		return err
	}
	if hasCreatedBy {
		if err := tx.Exec(
			"ALTER TABLE sa_dataset_sources DROP COLUMN created_by",
		).Error; err != nil {
			return fmt.Errorf("drop legacy dataset source creator: %w", err)
		}
	}
	return nil
}

func hasDatasetSourceColumn(tx *gorm.DB, name string) (bool, error) {
	columns, err := tx.Migrator().ColumnTypes(&models.SaDatasetSource{})
	if err != nil {
		return false, fmt.Errorf("inspect dataset source columns: %w", err)
	}
	for _, column := range columns {
		if strings.EqualFold(column.Name(), name) {
			return true, nil
		}
	}
	return false, nil
}

func prepareDatasetItems(tx *gorm.DB) error {
	if tx.Migrator().HasIndex(&models.SaDatasetItem{}, "idx_sa_dataset_item_source_external") {
		return nil
	}
	if err := tx.Exec(
		"UPDATE sa_dataset_items SET external_id = NULL WHERE external_id = ''",
	).Error; err != nil {
		return err
	}
	if err := tx.Exec(`
		DELETE FROM sa_dataset_items
		WHERE external_id IS NOT NULL
		  AND EXISTS (
			SELECT 1
			FROM sa_dataset_items AS earlier
			WHERE earlier.source_id = sa_dataset_items.source_id
			  AND earlier.external_id = sa_dataset_items.external_id
			  AND earlier.id < sa_dataset_items.id
		  )
	`).Error; err != nil {
		return err
	}
	return tx.Exec(`
		CREATE UNIQUE INDEX IF NOT EXISTS idx_sa_dataset_item_source_external
		ON sa_dataset_items (source_id, external_id)
	`).Error
}

func upsertDefaultDatasetSource(tx *gorm.DB, root models.SaUser, definition defaultDatasetSource) error {
	var source models.SaDatasetSource
	err := tx.Where("owner_id = ? AND seed_key = ?", root.ID, definition.Key).First(&source).Error
	if err == nil {
		return tx.Model(&source).Updates(map[string]any{
			"url": definition.URL, "icon_url": definition.IconURL, "kind": "rss",
		}).Error
	}
	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return err
	}

	knownURLs := append([]string{definition.URL}, definition.LegacyURLs...)
	err = tx.Where("owner_id = ? AND url IN ?", root.ID, knownURLs).First(&source).Error
	if err == nil {
		return tx.Model(&source).Updates(map[string]any{
			"seed_key": definition.Key, "url": definition.URL,
			"icon_url": definition.IconURL, "kind": "rss",
		}).Error
	}
	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return err
	}

	seedKey := definition.Key
	return tx.Create(&models.SaDatasetSource{
		OwnerID: root.ID, SeedKey: &seedKey, Name: definition.Name,
		Kind: "rss", URL: definition.URL, IconURL: definition.IconURL,
		Description: "默认资讯订阅", Enabled: true,
	}).Error
}