fix: fix context propagation in UseDB method (#1278)

* fix: context propagation in useDB method

Signed-off-by: hlts2 <[email protected]>

* fix: add nil check for statement

Signed-off-by: hlts2 <[email protected]>

---------

Signed-off-by: hlts2 <[email protected]>

Author: hlts2

do.go

@@ -50,7 +50,11 @@ var (
 
 // UseDB specify a db connection(*gorm.DB)
 func (d *DO) UseDB(db *gorm.DB, opts ...DOOption) {
-	db = db.Session(&gorm.Session{Context: context.Background()})
+	ctx := context.Background()
+	if db.Statement != nil && db.Statement.Context != nil {
+		ctx = db.Statement.Context
+	}
+	db = db.Session(&gorm.Session{Context: ctx})
 	d.db = db
 	config := &DOConfig{}
 	for _, opt := range opts {