kra-oa/internal/data/data.go

package data

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/google/wire"
	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
	"kra/internal/conf"
)

var ProviderSet = wire.NewSet(NewData, NewUserRepo, NewInitializationRepo, NewAccessRepo, NewMenuRepo, NewOrganizationRepo, NewDictionaryRepo, NewParameterRepo, NewAPITokenRepo, NewSecurityRepo, NewVersionRepo, NewExportRepo, NewAuditRepo, NewTaskRepo, NewMediaRepo, NewAnnouncementRepo, NewEmailRepo, NewCache, NewFileStorage)

type Data struct {
	initMu   sync.Mutex
	configMu sync.Mutex
	gormDB   *reloadableDB
	redis    *reloadableRedis
	mongo    *reloadableMongo
	runtime  *conf.Runtime
	storage  *reloadableStorage
	dbListMu sync.RWMutex
	dbList   map[string]*gorm.DB
}

func openDatabaseList(configs []*conf.Data_Database) (map[string]*gorm.DB, error) {
	items := make(map[string]*gorm.DB)
	for _, config := range configs {
		if config == nil || config.Disable || config.AliasName == "" {
			continue
		}
		db, err := openDatabase(config, false, "")
		if err != nil {
			for _, opened := range items {
				if sqlDB, dbErr := opened.DB(); dbErr == nil {
					_ = sqlDB.Close()
				}
			}
			return nil, fmt.Errorf("open database %q: %w", config.AliasName, err)
		}
		registerDataScopeCallbacks(db)
		items[config.AliasName] = db
	}
	return items, nil
}

func closeDatabaseList(items map[string]*gorm.DB) {
	for _, db := range items {
		if sqlDB, err := db.DB(); err == nil {
			_ = sqlDB.Close()
		}
	}
}

func (d *Data) replaceDatabaseList(items map[string]*gorm.DB) {
	d.dbListMu.Lock()
	old := d.dbList
	d.dbList = items
	d.dbListMu.Unlock()
	closeDatabaseList(old)
}

func (d *Data) database(name string) (*gorm.DB, error) {
	if name == "" {
		return d.gormDB.DB(), nil
	}
	d.dbListMu.RLock()
	db := d.dbList[name]
	d.dbListMu.RUnlock()
	if db == nil {
		return nil, fmt.Errorf("database %q not found", name)
	}
	return db, nil
}

func NewData(runtime *conf.Runtime) (*Data, func(), error) {
	c := runtime.Data()
	if c == nil || c.Database == nil {
		return nil, nil, fmt.Errorf("database configuration is required")
	}
	d := &Data{runtime: runtime}
	db, err := openDatabase(c.Database, false, "")
	usingFallback := false
	if err != nil {
		// The initialization endpoint must remain available when the configured
		// target database has not been created yet.
		log.Printf("configured database unavailable before initialization: %v", err)
		db, err = openFallbackDatabase()
		if err != nil {
			return nil, nil, fmt.Errorf("open bootstrap database: %w", err)
		}
		usingFallback = true
	}
	d.gormDB = newReloadableDB(db)
	d.dbList, err = openDatabaseList(c.DatabaseList)
	if err != nil {
		d.gormDB.close()
		return nil, nil, err
	}
	admin := runtime.Admin()
	disableAutoMigrate := admin != nil && admin.System != nil && admin.System.DisableAutoMigrate
	if !usingFallback && !disableAutoMigrate {
		if err = migrateAll(db); err != nil {
			return nil, nil, fmt.Errorf("migrate tables: %w", err)
		}
	}
	useRedis := admin != nil && admin.System != nil && admin.System.UseRedis
	d.redis = newReloadableRedis(openRedis(c.Redis, useRedis))
	useMongo := admin != nil && admin.System != nil && admin.System.UseMongo
	mongoClient, err := openMongo(c.Mongo, useMongo)
	if err != nil {
		log.Printf("mongo unavailable: %v", err)
		mongoClient = nil
	}
	d.mongo = newReloadableMongo(mongoClient)
	stopConfigWatcher := d.watchConfig()
	cleanup := func() {
		stopConfigWatcher()
		d.gormDB.close()
		closeDatabaseList(d.dbList)
		d.redis.close()
		d.mongo.close()
	}
	return d, cleanup, nil
}

func openRedis(config *conf.Data_Redis, enabled bool) redis.UniversalClient {
	if !enabled || config == nil || (config.Addr == "" && len(config.ClusterAddrs) == 0) {
		return nil
	}
	var candidate redis.UniversalClient
	if config.UseCluster {
		addresses := config.ClusterAddrs
		if len(addresses) == 0 && config.Addr != "" {
			addresses = []string{config.Addr}
		}
		candidate = redis.NewClusterClient(&redis.ClusterOptions{Addrs: addresses, Password: config.Password})
	} else {
		options := &redis.Options{Addr: config.Addr, Network: config.Network, Password: config.Password, DB: int(config.Db)}
		if config.ReadTimeout != nil {
			options.ReadTimeout = config.ReadTimeout.AsDuration()
		}
		if config.WriteTimeout != nil {
			options.WriteTimeout = config.WriteTimeout.AsDuration()
		}
		candidate = redis.NewClient(options)
	}
	pingCtx, cancel := context.WithTimeout(context.Background(), 800*time.Millisecond)
	defer cancel()
	if err := candidate.Ping(pingCtx).Err(); err != nil {
		log.Printf("redis unavailable, using in-memory cache: %v", err)
		_ = candidate.Close()
		return nil
	}
	return candidate
}

func (d *Data) activateDatabase(db *gorm.DB, config *conf.Data_Database) {
	d.gormDB.replace(db)
	d.runtime.UpdateDatabase(config)
}