kra-new/internal/data/config_store.go

package data

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"

	"kra/internal/conf"

	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"gopkg.in/yaml.v3"
)

func protoMap(message proto.Message) (map[string]any, error) {
	raw, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(message)
	if err != nil {
		return nil, err
	}
	var value map[string]any
	if err = json.Unmarshal(raw, &value); err != nil {
		return nil, err
	}
	delete(value, "config_path")
	return value, nil
}

func setYAMLMapping(node *yaml.Node, key string, value any) error {
	if node.Kind == yaml.DocumentNode {
		node = node.Content[0]
	}
	if node.Kind != yaml.MappingNode {
		return fmt.Errorf("configuration root is not a mapping")
	}
	raw, err := yaml.Marshal(value)
	if err != nil {
		return err
	}
	var replacement yaml.Node
	if err = yaml.Unmarshal(raw, &replacement); err != nil {
		return err
	}
	for i := 0; i < len(node.Content); i += 2 {
		if node.Content[i].Value == key {
			node.Content[i+1] = replacement.Content[0]
			return nil
		}
	}
	node.Content = append(node.Content, &yaml.Node{Kind: yaml.ScalarNode, Value: key}, replacement.Content[0])
	return nil
}

func (d *Data) persistConfig() error {
	return d.persistConfigValues(d.config, d.admin)
}

func (d *Data) persistConfigValues(dataConfig *conf.Data, adminConfig *conf.AdminBackend) error {
	if d.admin == nil || d.admin.ConfigPath == "" {
		return nil
	}
	raw, err := os.ReadFile(d.admin.ConfigPath)
	if err != nil {
		return err
	}
	var document yaml.Node
	if err = yaml.Unmarshal(raw, &document); err != nil {
		return err
	}
	dataValue, err := protoMap(dataConfig)
	if err != nil {
		return err
	}
	adminValue, err := protoMap(adminConfig)
	if err != nil {
		return err
	}
	if err = setYAMLMapping(&document, "data", dataValue); err != nil {
		return err
	}
	if err = setYAMLMapping(&document, "admin", adminValue); err != nil {
		return err
	}
	output, err := yaml.Marshal(&document)
	if err != nil {
		return err
	}
	if err = os.MkdirAll(filepath.Dir(d.admin.ConfigPath), 0o755); err != nil {
		return err
	}
	temporary, err := os.CreateTemp(filepath.Dir(d.admin.ConfigPath), ".kra-config-*.yaml")
	if err != nil {
		return err
	}
	tempName := temporary.Name()
	defer os.Remove(tempName)
	if _, err = temporary.Write(output); err != nil {
		_ = temporary.Close()
		return err
	}
	if err = temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return err
	}
	if err = temporary.Close(); err != nil {
		return err
	}
	return os.Rename(tempName, d.admin.ConfigPath)
}

func (d *Data) reloadConfig(ctx context.Context) error {
	if d.admin == nil || d.admin.ConfigPath == "" {
		return fmt.Errorf("configuration path is not set")
	}
	raw, err := os.ReadFile(d.admin.ConfigPath)
	if err != nil {
		return err
	}
	var plain map[string]any
	if err = yaml.Unmarshal(raw, &plain); err != nil {
		return err
	}
	encoded, err := json.Marshal(plain)
	if err != nil {
		return err
	}
	var next conf.Bootstrap
	if err = (protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(encoded, &next); err != nil {
		return err
	}
	if next.Data == nil || next.Data.Database == nil || next.Admin == nil {
		return fmt.Errorf("data.database and admin configuration are required")
	}
	next.Admin.ConfigPath = d.admin.ConfigPath
	candidateDB, err := openDatabase(next.Data.Database, false, "")
	if err != nil {
		return fmt.Errorf("reload database: %w", err)
	}
	closeCandidate := true
	defer func() {
		if closeCandidate {
			if sqlDB, closeErr := candidateDB.DB(); closeErr == nil {
				_ = sqlDB.Close()
			}
		}
	}()
	if sqlDB, dbErr := candidateDB.DB(); dbErr != nil {
		return dbErr
	} else if err = sqlDB.PingContext(ctx); err != nil {
		return fmt.Errorf("reload database: %w", err)
	}
	if err = migrateAll(candidateDB.WithContext(ctx)); err != nil {
		return fmt.Errorf("reload database migrations: %w", err)
	}
	candidateStorage, err := buildFileStorage(next.Admin)
	if err != nil {
		return fmt.Errorf("reload storage: %w", err)
	}
	candidateRedis := openRedis(next.Data.Redis)

	d.mu.Lock()
	oldDB, oldRedis := d.gormDB, d.redis
	d.gormDB, d.redis = candidateDB, candidateRedis
	proto.Reset(d.config)
	proto.Merge(d.config, next.Data)
	proto.Reset(d.admin)
	proto.Merge(d.admin, next.Admin)
	d.database = d.config.Database
	d.mu.Unlock()
	if d.storage != nil {
		d.storage.replace(candidateStorage)
	}
	closeCandidate = false
	if oldDB != nil {
		if sqlDB, closeErr := oldDB.DB(); closeErr == nil {
			_ = sqlDB.Close()
		}
	}
	if oldRedis != nil {
		_ = oldRedis.Close()
	}
	return nil
}