goyco/internal/database/monitoring_plugin.go

package database

import (
	"context"
	"time"

	"gorm.io/gorm"
	"goyco/internal/middleware"
)

type contextKey string

const gormOperationStartKey contextKey = "gorm_operation_start"

type GormDBMonitor struct {
	monitor middleware.DBMonitor
}

func NewGormDBMonitor(monitor middleware.DBMonitor) *GormDBMonitor {
	return &GormDBMonitor{
		monitor: monitor,
	}
}

func (g *GormDBMonitor) Name() string {
	return "db_monitor"
}

func (g *GormDBMonitor) Initialize(db *gorm.DB) error {
	if err := db.Callback().Create().Before("gorm:create").Register("db_monitor:before_create", g.beforeCreate); err != nil {
		return err
	}
	if err := db.Callback().Create().After("gorm:create").Register("db_monitor:after_create", g.afterCreate); err != nil {
		return err
	}

	if err := db.Callback().Query().Before("gorm:query").Register("db_monitor:before_query", g.beforeQuery); err != nil {
		return err
	}
	if err := db.Callback().Query().After("gorm:query").Register("db_monitor:after_query", g.afterQuery); err != nil {
		return err
	}

	if err := db.Callback().Update().Before("gorm:update").Register("db_monitor:before_update", g.beforeUpdate); err != nil {
		return err
	}
	if err := db.Callback().Update().After("gorm:update").Register("db_monitor:after_update", g.afterUpdate); err != nil {
		return err
	}

	if err := db.Callback().Delete().Before("gorm:delete").Register("db_monitor:before_delete", g.beforeDelete); err != nil {
		return err
	}
	if err := db.Callback().Delete().After("gorm:delete").Register("db_monitor:after_delete", g.afterDelete); err != nil {
		return err
	}

	if err := db.Callback().Row().Before("gorm:row").Register("db_monitor:before_row", g.beforeRow); err != nil {
		return err
	}
	if err := db.Callback().Row().After("gorm:row").Register("db_monitor:after_row", g.afterRow); err != nil {
		return err
	}

	if err := db.Callback().Raw().Before("gorm:raw").Register("db_monitor:before_raw", g.beforeRaw); err != nil {
		return err
	}
	if err := db.Callback().Raw().After("gorm:raw").Register("db_monitor:after_raw", g.afterRaw); err != nil {
		return err
	}

	return nil
}

func (g *GormDBMonitor) beforeCreate(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterCreate(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "CREATE")
}

func (g *GormDBMonitor) beforeQuery(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterQuery(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "SELECT")
}

func (g *GormDBMonitor) beforeUpdate(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterUpdate(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "UPDATE")
}

func (g *GormDBMonitor) beforeDelete(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterDelete(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "DELETE")
}

func (g *GormDBMonitor) beforeRow(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterRow(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "ROW")
}

func (g *GormDBMonitor) beforeRaw(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	ctx := context.WithValue(db.Statement.Context, gormOperationStartKey, time.Now())
	db.Statement.Context = ctx
}

func (g *GormDBMonitor) afterRaw(db *gorm.DB) {
	if g.monitor == nil {
		return
	}

	g.logOperation(db, "RAW")
}

func (g *GormDBMonitor) logOperation(db *gorm.DB, operation string) {
	if g.monitor == nil {
		return
	}

	startTime, ok := db.Statement.Context.Value(gormOperationStartKey).(time.Time)
	if !ok {
		return
	}

	duration := time.Since(startTime)

	query := g.buildQueryString(db, operation)

	g.monitor.LogQuery(query, duration, db.Error)
}

func (g *GormDBMonitor) buildQueryString(db *gorm.DB, operation string) string {
	if db.Statement.SQL.String() != "" {
		return db.Statement.SQL.String()
	}

	query := operation

	if db.Statement.Table != "" {
		query += " FROM " + db.Statement.Table
	}

	if db.Statement.Model != nil {

		if stmt := db.Statement; stmt.Schema != nil {
			query = operation + " " + stmt.Schema.Table
		}
	}

	return query
}