gmbase/cmd/mbasedump/util.go

/*
 * Copyright 2026 Oleg Borodin <onborodin@gmail.com>
 */
package main

import (
	"bytes"
	"context"
	"io"
	"os"
	"path/filepath"

	"mbase/app/config"
	"mbase/app/database"
	"mbase/app/descriptor"
	"mbase/pkg/auxtool"
	"mbase/pkg/logger"

	"github.com/spf13/cobra"
	"go.yaml.in/yaml/v4"
)

type Util struct {
	rootCmd               *cobra.Command
	dumpDatabaseParams    dumpDatabaseParams
	restoreDatabaseParams restoreDatabaseParams
}

func NewUtil() *Util {
	return &Util{}
}

func (util *Util) GetRooCmd() *cobra.Command {
	return util.rootCmd
}

func (util *Util) Build() error {
	var err error
	execName := filepath.Base(os.Args[0])
	rootCmd := &cobra.Command{
		Use:          execName,
		Short:        "\nDump application database",
		SilenceUsage: true,
	}
	rootCmd.CompletionOptions.DisableDefaultCmd = true

	var dumpDatabaseCmd = &cobra.Command{
		Use:   "dump [filename|-]",
		Short: "Dump application database",
		Args:  cobra.ExactArgs(1),
		Run:   util.DumpDatabase,
	}
	rootCmd.AddCommand(dumpDatabaseCmd)

	var restoreDatabaseCmd = &cobra.Command{
		Use:   "restore [] [filename|-]",
		Short: "Restore application database",
		Args:  cobra.ExactArgs(1),
		Run:   util.RestoreDatabase,
	}
	restoreDatabaseCmd.Flags().BoolVarP(&util.restoreDatabaseParams.DeleteAllRecords, "clean", "C", false, "Clean all record")
	rootCmd.AddCommand(restoreDatabaseCmd)

	util.rootCmd = rootCmd
	return err
}

func (util *Util) Exec(args []string) error {
	var err error
	util.rootCmd.SetArgs(args)
	err = util.rootCmd.Execute()
	return err
}

func (util *Util) DumpDatabase(cmd *cobra.Command, args []string) {
	util.dumpDatabaseParams.Filename = args[0]
	res, err := util.dumpDatabase(util.dumpDatabaseParams)
	printResponse(res, err)
}

type dumpDatabaseParams struct {
	Filename string
}
type dumpDatabaseResult struct{}

func (util *Util) dumpDatabase(params dumpDatabaseParams) (dumpDatabaseResult, error) {
	var err error
	res := dumpDatabaseResult{}
	ctx := context.Background()

	conf := config.NewConfig()
	err = conf.ReadFile()
	if err != nil {
		return res, err
	}
	err = conf.ReadEnv()
	if err != nil {
		return res, err
	}
	db, err := database.NewDatabase(conf.DataDir)
	if err != nil {
		return res, err
	}
	err = db.OpenDatabase()
	if err != nil {
		return res, err
	}
	file := os.Stdout
	if params.Filename != "-" {
		file, err = os.OpenFile(params.Filename, os.O_CREATE|os.O_WRONLY, 0640)
		if err != nil {
			return res, err
		}
		defer file.Close()
	}
	listAccounts, err := db.CompletedListAccounts(ctx)
	if err != nil {
		return res, err
	}
	listGrants, err := db.ListGrants(ctx)
	if err != nil {
		return res, err
	}
	dump := descriptor.Dump{
		Timestamp: auxtool.TimeNow(),
		Accounts:  listAccounts,
		Grants:    listGrants,
	}
	dumpBytes, err := yaml.Marshal(dump)
	if err != nil {
		return res, err
	}
	_, err = file.Write(dumpBytes)
	if err != nil {
		return res, err
	}
	return res, err
}

func (util *Util) RestoreDatabase(cmd *cobra.Command, args []string) {
	util.restoreDatabaseParams.Filename = args[0]
	res, err := util.restoreDatabase(util.restoreDatabaseParams)
	printResponse(res, err)
}

type restoreDatabaseParams struct {
	Filename         string
	DeleteAllRecords bool
}
type restoreDatabaseResult struct{}

func (util *Util) restoreDatabase(params restoreDatabaseParams) (restoreDatabaseResult, error) {
	var err error
	res := restoreDatabaseResult{}
	ctx := context.Background()
	log := logger.NewLogger("restore")

	conf := config.NewConfig()
	err = conf.ReadFile()
	if err != nil {
		return res, err
	}
	err = conf.ReadEnv()
	if err != nil {
		return res, err
	}

	db, err := database.NewDatabase(conf.DataDir)
	if err != nil {
		return res, err
	}
	err = db.OpenDatabase()
	if err != nil {
		return res, err
	}
	file := os.Stdin
	if params.Filename != "-" {
		file, err = os.Open(params.Filename)
		if err != nil {
			return res, err
		}
		defer file.Close()
	}
	buffer := bytes.NewBuffer(nil)
	_, err = io.Copy(buffer, file)
	if err != nil {
		return res, err
	}
	dump := descriptor.Dump{}
	err = yaml.Unmarshal(buffer.Bytes(), &dump)
	if err != nil {
		return res, err
	}
	if params.DeleteAllRecords {
		err = db.CleanDatabase(ctx)
		if err != nil {
			return res, err
		}
	}
	for _, account := range dump.Accounts {
		log.Infof("Insert account %s", account.Username)
		err = db.InsertAccount(ctx, &account)
		if err != nil {
			log.Errorf("Insert account error: %v", err)
		}
	}
	for _, grant := range dump.Grants {
		log.Infof("Insert grant %s for account %d", grant.Operation, grant.AccountID)
		err = db.InsertGrant(ctx, &grant)
		if err != nil {
			log.Errorf("Insert account error: %v", err)
		}
	}
	return res, err
}