add context feature for package gdb

database/gdb/gdb.go

@@ -8,6 +8,7 @@
 package gdb
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -42,6 +43,12 @@ type DB interface {
 	// Note that it is not recommended using the this function manually.
 	Open(config *ConfigNode) (*sql.DB, error)
 
+	// Ctx is a chaining function, which creates and returns a new DB that is a shallow copy
+	// of current DB object and with given context in it.
+	// Note that this returned DB object can be used only once, so do not assign it to
+	// a global or package variable for long using.
+	Ctx(ctx context.Context) DB
+
 	// ===========================================================================
 	// Query APIs.
 	// ===========================================================================
@@ -137,6 +144,7 @@ type DB interface {
 	// Utility methods.
 	// ===========================================================================
 
+	GetCtx() context.Context
 	GetChars() (charLeft string, charRight string)
 	GetMaster(schema ...string) (*sql.DB, error)
 	GetSlave(schema ...string) (*sql.DB, error)
@@ -154,23 +162,24 @@ type DB interface {
 	HandleSqlBeforeCommit(link Link, sql string, args []interface{}) (string, []interface{})
 
 	// ===========================================================================
-	// Internal methods.
+	// Internal methods, for internal usage purpose, you do not need consider it.
 	// ===========================================================================
 
 	mappingAndFilterData(schema, table string, data map[string]interface{}, filter bool) (map[string]interface{}, error)
-	convertDatabaseValueToLocalValue(fieldValue interface{}, fieldType string) interface{}
-	rowsToResult(rows *sql.Rows) (Result, error)
+	convertFieldValueToLocalValue(fieldValue interface{}, fieldType string) interface{}
+	convertRowsToResult(rows *sql.Rows) (Result, error)
 }
 
 // Core is the base struct for database management.
 type Core struct {
-	DB     DB            // DB interface object.
-	group  string        // Configuration group name.
-	debug  *gtype.Bool   // Enable debug mode for the database, which can be changed in runtime.
-	cache  *gcache.Cache // Cache manager, SQL result cache only.
-	schema *gtype.String // Custom schema for this object.
-	logger *glog.Logger  // Logger.
-	config *ConfigNode   // Current config node.
+	DB     DB              // DB interface object.
+	group  string          // Configuration group name.
+	debug  *gtype.Bool     // Enable debug mode for the database, which can be changed in runtime.
+	cache  *gcache.Cache   // Cache manager, SQL result cache only.
+	schema *gtype.String   // Custom schema for this object.
+	logger *glog.Logger    // Logger.
+	config *ConfigNode     // Current config node.
+	ctx    context.Context // Context for chaining operation only.
 }
 
 // Driver is the interface for integrating sql drivers into package gdb.

database/gdb/gdb_core.go

@@ -8,6 +8,7 @@
 package gdb
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -23,6 +24,35 @@ import (
 	"github.com/gogf/gf/util/gconv"
 )
 
+// Ctx is a chaining function, which creates and returns a new DB that is a shallow copy
+// of current DB object and with given context in it.
+// Note that this returned DB object can be used only once, so do not assign it to
+// a global or package variable for long using.
+func (c *Core) Ctx(ctx context.Context) DB {
+	if ctx == nil {
+		return c.DB
+	}
+	var (
+		err        error
+		newCore    = &Core{}
+		configNode = c.DB.GetConfig()
+	)
+	*newCore = *c
+	newCore.ctx = ctx
+	newCore.DB, err = driverMap[configNode.Type].New(newCore, configNode)
+	// Seldom error, just log it.
+	if err != nil {
+		c.DB.GetLogger().Ctx(ctx).Error(err)
+	}
+	return newCore.DB
+}
+
+// GetCtx returns the context for current DB.
+// Note that it might be nil.
+func (c *Core) GetCtx() context.Context {
+	return c.ctx
+}
+
 // Master creates and returns a connection from master node if master-slave configured.
 // It returns the default connection if master-slave not configured.
 func (c *Core) Master() (*sql.DB, error) {
@@ -164,7 +194,7 @@ func (c *Core) DoGetAll(link Link, sql string, args ...interface{}) (result Resu
 		return nil, err
 	}
 	defer rows.Close()
-	return c.DB.rowsToResult(rows)
+	return c.DB.convertRowsToResult(rows)
 }
 
 // GetOne queries and returns one record from database.
@@ -767,8 +797,8 @@ func (c *Core) DoDelete(link Link, table string, condition string, args ...inter
 	return c.DB.DoExec(link, fmt.Sprintf("DELETE FROM %s%s", table, condition), args...)
 }
 
-// rowsToResult converts underlying data record type sql.Rows to Result type.
-func (c *Core) rowsToResult(rows *sql.Rows) (Result, error) {
+// convertRowsToResult converts underlying data record type sql.Rows to Result type.
+func (c *Core) convertRowsToResult(rows *sql.Rows) (Result, error) {
 	if !rows.Next() {
 		return nil, nil
 	}
@@ -800,7 +830,7 @@ func (c *Core) rowsToResult(rows *sql.Rows) (Result, error) {
 			if value == nil {
 				row[columnNames[i]] = gvar.New(nil)
 			} else {
-				row[columnNames[i]] = gvar.New(c.DB.convertDatabaseValueToLocalValue(value, columnTypes[i]))
+				row[columnNames[i]] = gvar.New(c.DB.convertFieldValueToLocalValue(value, columnTypes[i]))
 			}
 		}
 		records = append(records, row)
@@ -821,14 +851,14 @@ func (c *Core) MarshalJSON() ([]byte, error) {
 }
 
 // writeSqlToLogger outputs the sql object to logger.
-// It is enabled when configuration "debug" is true.
+// It is enabled only if configuration "debug" is true.
 func (c *Core) writeSqlToLogger(v *Sql) {
 	s := fmt.Sprintf("[%3d ms] [%s] %s", v.End-v.Start, v.Group, v.Format)
 	if v.Error != nil {
 		s += "\nError: " + v.Error.Error()
-		c.logger.Error(s)
+		c.logger.Ctx(c.DB.GetCtx()).Error(s)
 	} else {
-		c.logger.Debug(s)
+		c.logger.Ctx(c.DB.GetCtx()).Debug(s)
 	}
 }
 

database/gdb/gdb_core_structure.go

@@ -22,9 +22,9 @@ import (
 	"github.com/gogf/gf/util/gconv"
 )
 
-// convertDatabaseValueToLocalValue automatically checks and converts field value from database type
+// convertFieldValueToLocalValue automatically checks and converts field value from database type
 // to golang variable type.
-func (c *Core) convertDatabaseValueToLocalValue(fieldValue interface{}, fieldType string) interface{} {
+func (c *Core) convertFieldValueToLocalValue(fieldValue interface{}, fieldType string) interface{} {
 	// If there's no type retrieved, it returns the <fieldValue> directly
 	// to use its original data type, as <fieldValue> is type of interface{}.
 	if fieldType == "" {
@@ -146,7 +146,8 @@ func (c *Core) convertDatabaseValueToLocalValue(fieldValue interface{}, fieldTyp
 	}
 }
 
-// filterFields removes all key-value pairs which are not the field of given table.
+// mappingAndFilterData automatically mappings the map key to table field and removes
+// all key-value pairs that are not the field of given table.
 func (c *Core) mappingAndFilterData(schema, table string, data map[string]interface{}, filter bool) (map[string]interface{}, error) {
 	if fieldsMap, err := c.DB.TableFields(table, schema); err == nil {
 		fieldsKeyMap := make(map[string]interface{}, len(fieldsMap))

database/gdb/gdb_model.go

@@ -7,6 +7,7 @@
 package gdb
 
 import (
+	"context"
 	"fmt"
 	"github.com/gogf/gf/text/gregex"
 	"time"
@@ -52,13 +53,14 @@ type whereHolder struct {
 }
 
 const (
-	gLINK_TYPE_MASTER   = 1
-	gLINK_TYPE_SLAVE    = 2
-	gWHERE_HOLDER_WHERE = 1
-	gWHERE_HOLDER_AND   = 2
-	gWHERE_HOLDER_OR    = 3
-	OPTION_OMITEMPTY    = 1 << iota
-	OPTION_ALLOWEMPTY
+	OPTION_OMITEMPTY  = 1
+	OPTION_ALLOWEMPTY = 2
+
+	linkTypeMaster   = 1
+	linkTypeMSlave   = 2
+	whereHolderWhere = 1
+	whereHolderAnd   = 2
+	whereHolderOr    = 3
 )
 
 // Table creates and returns a new ORM model from given schema.
@@ -112,6 +114,16 @@ func (tx *TX) Model(table ...string) *Model {
 	return tx.Table(table...)
 }
 
+//
+func (m *Model) Ctx(ctx context.Context) *Model {
+	if ctx == nil {
+		return m
+	}
+	model := m.getModel()
+	model.db = model.db.Ctx(ctx)
+	return model
+}
+
 // As sets an alias name for current table.
 func (m *Model) As(as string) *Model {
 	if m.tables != "" {
@@ -178,7 +190,7 @@ func (m *Model) Clone() *Model {
 // Master marks the following operation on master node.
 func (m *Model) Master() *Model {
 	model := m.getModel()
-	model.linkType = gLINK_TYPE_MASTER
+	model.linkType = linkTypeMaster
 	return model
 }
 
@@ -186,7 +198,7 @@ func (m *Model) Master() *Model {
 // Note that it makes sense only if there's any slave node configured.
 func (m *Model) Slave() *Model {
 	model := m.getModel()
-	model.linkType = gLINK_TYPE_SLAVE
+	model.linkType = linkTypeMSlave
 	return model
 }
 

database/gdb/gdb_model_condition.go

@@ -27,7 +27,7 @@ func (m *Model) Where(where interface{}, args ...interface{}) *Model {
 		model.whereHolder = make([]*whereHolder, 0)
 	}
 	model.whereHolder = append(model.whereHolder, &whereHolder{
-		operator: gWHERE_HOLDER_WHERE,
+		operator: whereHolderWhere,
 		where:    where,
 		args:     args,
 	})
@@ -65,7 +65,7 @@ func (m *Model) And(where interface{}, args ...interface{}) *Model {
 		model.whereHolder = make([]*whereHolder, 0)
 	}
 	model.whereHolder = append(model.whereHolder, &whereHolder{
-		operator: gWHERE_HOLDER_AND,
+		operator: whereHolderAnd,
 		where:    where,
 		args:     args,
 	})
@@ -79,7 +79,7 @@ func (m *Model) Or(where interface{}, args ...interface{}) *Model {
 		model.whereHolder = make([]*whereHolder, 0)
 	}
 	model.whereHolder = append(model.whereHolder, &whereHolder{
-		operator: gWHERE_HOLDER_OR,
+		operator: whereHolderOr,
 		where:    where,
 		args:     args,
 	})

database/gdb/gdb_model_utility.go

@@ -149,19 +149,19 @@ func (m *Model) getLink(master bool) Link {
 	linkType := m.linkType
 	if linkType == 0 {
 		if master {
-			linkType = gLINK_TYPE_MASTER
+			linkType = linkTypeMaster
 		} else {
-			linkType = gLINK_TYPE_SLAVE
+			linkType = linkTypeMSlave
 		}
 	}
 	switch linkType {
-	case gLINK_TYPE_MASTER:
+	case linkTypeMaster:
 		link, err := m.db.GetMaster(m.schema)
 		if err != nil {
 			panic(err)
 		}
 		return link
-	case gLINK_TYPE_SLAVE:
+	case linkTypeMSlave:
 		link, err := m.db.GetSlave(m.schema)
 		if err != nil {
 			panic(err)
@@ -196,7 +196,7 @@ func (m *Model) formatCondition(limit1 bool, isCountStatement bool) (conditionWh
 	if len(m.whereHolder) > 0 {
 		for _, v := range m.whereHolder {
 			switch v.operator {
-			case gWHERE_HOLDER_WHERE:
+			case whereHolderWhere:
 				if conditionWhere == "" {
 					newWhere, newArgs := formatWhere(
 						m.db, v.where, v.args, m.option&OPTION_OMITEMPTY > 0,
@@ -209,7 +209,7 @@ func (m *Model) formatCondition(limit1 bool, isCountStatement bool) (conditionWh
 				}
 				fallthrough
 
-			case gWHERE_HOLDER_AND:
+			case whereHolderAnd:
 				newWhere, newArgs := formatWhere(
 					m.db, v.where, v.args, m.option&OPTION_OMITEMPTY > 0,
 				)
@@ -224,7 +224,7 @@ func (m *Model) formatCondition(limit1 bool, isCountStatement bool) (conditionWh
 					conditionArgs = append(conditionArgs, newArgs...)
 				}
 
-			case gWHERE_HOLDER_OR:
+			case whereHolderOr:
 				newWhere, newArgs := formatWhere(
 					m.db, v.where, v.args, m.option&OPTION_OMITEMPTY > 0,
 				)

database/gdb/gdb_transaction.go

@@ -59,7 +59,7 @@ func (tx *TX) GetAll(sql string, args ...interface{}) (Result, error) {
 		return nil, err
 	}
 	defer rows.Close()
-	return tx.db.rowsToResult(rows)
+	return tx.db.convertRowsToResult(rows)
 }
 
 // GetOne queries and returns one record from database.

database/gdb/gdb_z_mysql_ctx_test.go

@@ -0,0 +1,41 @@
+// Copyright 2019 gf Author(https://github.com/gogf/gf). All Rights Reserved.
+//
+// This Source Code Form is subject to the terms of the MIT License.
+// If a copy of the MIT was not distributed with this file,
+// You can obtain one at https://github.com/gogf/gf.
+
+package gdb_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/gogf/gf/database/gdb"
+	"github.com/gogf/gf/test/gtest"
+)
+
+func Test_Ctx(t *testing.T) {
+	gtest.C(t, func(t *gtest.T) {
+		db, err := gdb.Instance()
+		t.Assert(err, nil)
+
+		err1 := db.PingMaster()
+		err2 := db.PingSlave()
+		t.Assert(err1, nil)
+		t.Assert(err2, nil)
+
+		newDb := db.Ctx(context.Background())
+		t.AssertNE(newDb, nil)
+	})
+}
+
+func Test_Ctx_Query(t *testing.T) {
+	db.GetLogger().SetCtxKeys("SpanId", "TraceId")
+	gtest.C(t, func(t *gtest.T) {
+		db.SetDebug(true)
+		defer db.SetDebug(false)
+		ctx := context.WithValue(context.Background(), "TraceId", "12345678")
+		ctx = context.WithValue(ctx, "SpanId", "0.1")
+		db.Ctx(ctx).Query("select 1")
+	})
+}

os/glog/glog_logger_chaining.go

@@ -17,6 +17,9 @@ import (
 // Ctx is a chaining function,
 // which sets the context for current logging.
 func (l *Logger) Ctx(ctx context.Context, keys ...interface{}) *Logger {
+	if ctx == nil {
+		return l
+	}
 	logger := (*Logger)(nil)
 	if l.parent == nil {
 		logger = l.Clone()