diff options
Diffstat (limited to 'common/context.go')
| -rw-r--r-- | common/context.go | 9 |
1 files changed, 4 insertions, 5 deletions
diff --git a/common/context.go b/common/context.go index 740e727..8bb02b2 100644 --- a/common/context.go +++ b/common/context.go @@ -2,8 +2,7 @@ package common import ( "context" - - "gorm.io/gorm" + "database/sql" ) type Key uint8 @@ -15,16 +14,16 @@ const ( KeyCopaingState Key = 3 ) -func SetDB(ctx context.Context, db *gorm.DB) context.Context { +func SetDB(ctx context.Context, db *sql.DB) context.Context { return context.WithValue(ctx, keyDB, db) } -func GetDB(ctx context.Context) *gorm.DB { +func GetDB(ctx context.Context) *sql.DB { raw := ctx.Value(keyDB) if raw == nil { return nil } - return raw.(*gorm.DB).WithContext(ctx) + return raw.(*sql.DB) } func SetDebug(ctx context.Context, b bool) context.Context { |