From 7dd48c0c504e9fdbf43d4974a10894858cd1d708 Mon Sep 17 00:00:00 2001
From: zhaoyupeng <zhaoyupeng@yizhisec.com>
Date: Tue, 17 Jun 2025 14:14:39 +0800
Subject: [PATCH] update: gorm session config

---
 database/db/db.go | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/database/db/db.go b/database/db/db.go
index a235b5d..73adcb3 100644
--- a/database/db/db.go
+++ b/database/db/db.go
@@ -5,8 +5,13 @@ import (
 	"gorm.io/gorm"
 )
 
+type Config struct {
+	Debug  bool
+	DryRun bool
+}
+
 type DB interface {
-	Session(ctx context.Context) *gorm.DB
+	Session(ctx context.Context, configs ...Config) *gorm.DB
 }
 
 type db struct {
@@ -17,6 +22,28 @@ var (
 	Default DB
 )
 
-func (db *db) Session(ctx context.Context) *gorm.DB {
-	return db.tx.Session(&gorm.Session{Context: ctx})
+func (db *db) Session(ctx context.Context, configs ...Config) *gorm.DB {
+	var (
+		sc      = &gorm.Session{Context: ctx}
+		session *gorm.DB
+	)
+
+	if len(configs) == 0 {
+		session = db.tx.Session(sc)
+		return session
+	}
+
+	cfg := configs[0]
+
+	if cfg.DryRun {
+		sc.DryRun = true
+	}
+
+	session = db.tx.Session(sc)
+
+	if cfg.Debug {
+		session = session.Debug()
+	}
+
+	return session
 }