add request id to sql query log

pkg/datastore/database.go

@@ -1,15 +1,29 @@
 package datastore
 
-import "xorm.io/xorm"
+import (
+	"xorm.io/xorm"
+
+	"github.com/mayswind/ezbookkeeping/pkg/core"
+)
 
 // Database represents a database instance
 type Database struct {
-	*xorm.EngineGroup
+	engineGroup *xorm.EngineGroup
+}
+
+// NewSession starts a new session with the specified context
+func (db *Database) NewSession(c *core.Context) *xorm.Session {
+	return db.engineGroup.Context(NewXOrmContextAdapter(c))
 }
 
 // DoTransaction runs a new database transaction
-func (db *Database) DoTransaction(fn func(sess *xorm.Session) error) (err error) {
-	sess := db.NewSession()
+func (db *Database) DoTransaction(c *core.Context, fn func(sess *xorm.Session) error) (err error) {
+	sess := db.engineGroup.NewSession()
+
+	if c != nil {
+		sess.Context(NewXOrmContextAdapter(c))
+	}
+
 	defer sess.Close()
 
 	if err = sess.Begin(); err != nil {

pkg/datastore/datastore.go

@@ -3,6 +3,7 @@ package datastore
 import (
 	"xorm.io/xorm"
 
+	"github.com/mayswind/ezbookkeeping/pkg/core"
 	"github.com/mayswind/ezbookkeeping/pkg/errs"
 )
 
@@ -17,13 +18,13 @@ func (s *DataStore) Choose(key int64) *Database {
 }
 
 // Query returns a new database session in a specific database by sharding key
-func (s *DataStore) Query(key int64) *xorm.Session {
-	return s.Choose(key).NewSession()
+func (s *DataStore) Query(c *core.Context, key int64) *xorm.Session {
+	return s.Choose(key).NewSession(c)
 }
 
 // DoTransaction runs a new database transaction in a specific database by sharding key
-func (s *DataStore) DoTransaction(key int64, fn func(sess *xorm.Session) error) (err error) {
-	return s.Choose(key).DoTransaction(fn)
+func (s *DataStore) DoTransaction(key int64, c *core.Context, fn func(sess *xorm.Session) error) (err error) {
+	return s.Choose(key).DoTransaction(c, fn)
 }
 
 // SyncStructs updates database structs by database models
@@ -31,7 +32,7 @@ func (s *DataStore) SyncStructs(beans ...interface{}) error {
 	var err error
 
 	for i := 0; i < len(s.databases); i++ {
-		err = s.databases[i].Sync2(beans...)
+		err = s.databases[i].engineGroup.Sync2(beans...)
 
 		if err != nil {
 			return err

pkg/datastore/datastore_container.go

@@ -104,14 +104,14 @@ func initializeDatabase(dbConfig *settings.DatabaseConfig) (*Database, error) {
 	engineGroup.SetConnMaxLifetime(time.Duration(dbConfig.ConnectionMaxLifeTime) * time.Second)
 
 	return &Database{
-		EngineGroup: engineGroup,
+		engineGroup: engineGroup,
 	}, nil
 }
 
 func setDatabaseLogger(database *Database, config *settings.Config) {
 	if config.EnableQueryLog {
-		database.SetLogger(NewXOrmLoggerAdapter(config.EnableQueryLog, config.LogLevel))
-		database.ShowSQL(true)
+		database.engineGroup.SetLogger(NewXOrmLoggerAdapter(config.EnableQueryLog, config.LogLevel))
+		database.engineGroup.ShowSQL(true)
 	}
 }
 

pkg/datastore/query_context.go

@@ -0,0 +1,50 @@
+package datastore
+
+import (
+	"fmt"
+	"time"
+
+	"xorm.io/xorm/log"
+
+	"github.com/mayswind/ezbookkeeping/pkg/core"
+)
+
+// XOrmContextAdapter represents the context adapter for xorm
+type XOrmContextAdapter struct {
+	requestId string
+}
+
+// Deadline does nothing
+func (c *XOrmContextAdapter) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+// Done always returns nil
+func (c *XOrmContextAdapter) Done() <-chan struct{} {
+	return nil
+}
+
+// Err always returns nil
+func (c *XOrmContextAdapter) Err() error {
+	return nil
+}
+
+// Value returns the value associated with this context for key, or nil
+// if no value is associated with key.
+func (c *XOrmContextAdapter) Value(key any) any {
+	if key == log.SessionIDKey && c.requestId != "" {
+		return fmt.Sprintf("r=%s", c.requestId)
+	}
+
+	return nil
+}
+
+func NewXOrmContextAdapter(c *core.Context) *XOrmContextAdapter {
+	if c != nil {
+		return &XOrmContextAdapter{
+			requestId: c.GetRequestId(),
+		}
+	}
+
+	return &XOrmContextAdapter{}
+}