kra-new/internal/data/data.go

package data

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

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

var ProviderSet = wire.NewSet(NewData, NewRuntimeSettings, NewTokenIssuer, NewUserRepo, NewInitializationRepo, NewAuthorityAccessRepo, NewAPIRepo, NewPermissionRepo, NewMenuRepo, NewDepartmentRepo, NewPositionRepo, NewDictionaryRepo, NewParameterRepo, NewAPITokenRepo, NewSecurityRepo, NewVersionRepo, NewExportRepo, NewAuditRepo, NewAuditRecorderRepo, NewLogFileRepo, NewTaskRepo, NewMediaRepo, NewAnnouncementRepo, NewEmailRepo, NewCache, NewFileStorage)

type Data struct {
	initMu        sync.Mutex
	configMu      sync.Mutex
	databaseReady atomic.Bool
	gormDB        *reloadableDB
	redis         *reloadableRedis
	mongo         *reloadableMongo
	runtime       *conf.Runtime
	storage       *reloadableStorage
	dbListMu      sync.RWMutex
	dbList        map[string]*gorm.DB
	appLogger     *slog.Logger
	auditLog      *dataScopeAuditWriter
}

func (d *Data) logger() *slog.Logger {
	if d != nil && d.appLogger != nil {
		return d.appLogger
	}
	return slog.Default()
}

func openDatabaseList(configs []*conf.Data_Database, appLogger ...*slog.Logger) (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, "", appLogger...)
		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)
		}
		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, appLogger *slog.Logger) (*Data, func(), error) {
	if appLogger == nil {
		appLogger = slog.Default()
	}
	c := runtime.Data()
	if c == nil {
		c = &conf.Data{}
	}
	if c.Database == nil {
		// An empty database block is the Kratos equivalent of GVA's empty
		// Dbname: the service starts on the bootstrap database so /init/checkdb
		// and /init/initdb remain available.
		c.Database = &conf.Data_Database{}
	}
	d := &Data{runtime: runtime, appLogger: appLogger}
	usingFallback := !databaseConnectionConfigured(c.Database)
	var db *gorm.DB
	var err error
	if !usingFallback {
		db, err = openDatabase(c.Database, false, "", appLogger)
	}
	if usingFallback || err != nil {
		// The initialization endpoint must remain available when the configured
		// target database has not been created yet.
		if err != nil {
			appLogger.Warn("configured database unavailable before initialization", "mod", "system", "error", err)
		}
		db, err = openFallbackDatabase(appLogger)
		if err != nil {
			return nil, nil, fmt.Errorf("open bootstrap database: %w", err)
		}
		usingFallback = true
	}
	d.databaseReady.Store(!usingFallback)
	d.gormDB = newReloadableDB(db, d.enqueueDataScopeAudit)
	d.auditLog = newDataScopeAuditWriter(d, appLogger)
	d.dbList, err = openDatabaseList(c.DatabaseList, appLogger)
	if err != nil {
		d.auditLog.Close()
		d.gormDB.close()
		return nil, nil, err
	}
	for _, item := range d.dbList {
		registerDataScopeCallbacks(item, d.enqueueDataScopeAudit)
	}
	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, appLogger))
	useMongo := admin != nil && admin.System != nil && admin.System.UseMongo
	mongoClient, err := openMongo(c.Mongo, useMongo)
	if err != nil {
		appLogger.Error("mongo unavailable", "mod", "mongo", "error", err)
		mongoClient = nil
	}
	d.mongo = newReloadableMongo(mongoClient)
	stopConfigWatcher := d.watchConfig()
	cleanup := func() {
		stopConfigWatcher()
		d.auditLog.Close()
		d.gormDB.close()
		closeDatabaseList(d.dbList)
		d.redis.close()
		d.mongo.close()
	}
	return d, cleanup, nil
}

func openRedis(config *conf.Data_Redis, enabled bool, appLogger ...*slog.Logger) 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 := slog.Default()
		if len(appLogger) > 0 && appLogger[0] != nil {
			log = appLogger[0]
		}
		log.Warn("redis unavailable, using in-memory cache", "mod", "redis", "error", err)
		_ = candidate.Close()
		return nil
	}
	return candidate
}

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