database/module.go

package database

import (
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"log/slog"
	"netgarden.dev/maf/maf"
	"time"
)

func NewModule() *Module {
	return NewModuleWithConfig(NewConfig())
}

func NewModuleWithConfig(config *Config) *Module {
	return &Module{
		config: config,
	}
}

type Module struct {
	manager *maf.Manager
	config  *Config
	db      *gorm.DB
}

func (m *Module) GetID() string {
	return "database"
}

func (m *Module) GetName() string {
	return "database"
}

func (m *Module) SetManager(manager *maf.Manager) {
	m.manager = manager
}

func (m *Module) CreateConfig() interface{} {
	if m.config != nil {
		return m.config
	}
	return NewConfig()
}

func (m *Module) SetConfig(cfg interface{}) {
	m.config = cfg.(*Config)
}

func (m *Module) PreInitialize() error {

	var err error

	m.db, err = m.connect()
	if err != nil {
		return err
	}

	err = m.autoMigrate()
	if err != nil {
		return err
	}

	m.notifyConsumers()

	return nil
}

func (m *Module) connect() (*gorm.DB, error) {

	gormPgConfig := postgres.Config{
		DSN: m.config.DSN,
	}

	gormConfig := &gorm.Config{}
	if m.config.ShowSql {
		gormConfig.Logger = logger.Default.LogMode(logger.Info)
	}

	var db *gorm.DB
	var err error

	db, err = gorm.Open(postgres.New(gormPgConfig), gormConfig)
	if err != nil {
		return nil, err
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}
	sqlDB.SetMaxIdleConns(m.config.MaxIdleConns)
	sqlDB.SetMaxOpenConns(m.config.MaxOpenConns)
	sqlDB.SetConnMaxLifetime(time.Second * time.Duration(m.config.ConnMaxLifetime))

	err = sqlDB.Ping()
	if err != nil {
		return nil, err
	}

	return db, err
}

func (m *Module) autoMigrate() error {

	if !m.config.AutoMigrate {
		return nil
	}

	var err error

	for _, module := range m.manager.GetModulesList() {
		if preMigrationModule, ok := module.(PreMigrationConsumer); ok {
			err = preMigrationModule.DBPreMigration(m.db)
			if err != nil {
				return err
			}
		}
	}

	err = m.autoMigrateTables()
	if err != nil {
		return err
	}

	return nil
}

func (m *Module) autoMigrateTables() error {

	slog.Info("Storage tables auto migration starting ...")

	entities := make([]interface{}, 0)
	for _, module := range m.manager.GetModulesList() {
		if consumer, ok := module.(EntitiesProvider); ok {
			consumerEntities := consumer.GetDBEntities()
			if consumerEntities != nil && len(consumerEntities) > 0 {
				entities = append(entities, consumerEntities...)
			}
		}
	}

	var err error

	err = m.db.AutoMigrate(entities...)

	slog.Info("Storage tables auto migration finished ...")

	return err
}

func (m *Module) notifyConsumers() {
	for _, module := range m.manager.GetModulesList() {
		if consumer, ok := module.(Consumer); ok {
			consumer.SetDB(m.db)
		}
	}
}

func (m *Module) Stop() error {

	if m.db == nil {
		return nil
	}

	sqlDB, err := m.db.DB()
	if err != nil {
		return err
	}

	err = sqlDB.Close()
	m.db = nil

	return err
}

func (m *Module) GetDB() *gorm.DB {
	return m.db
}