From 2ae32ed2c213053cb7a937d47194aa8ca0895a85 Mon Sep 17 00:00:00 2001 From: John Guo Date: Mon, 11 Jan 2021 20:48:35 +0800 Subject: [PATCH] add QueryTimeout/ExecTimeout/TranTimeout/PrepareTimeout for package gdb --- database/gdb/gdb_core.go | 30 +++++++++++++++++++++++------- database/gdb/gdb_core_config.go | 9 ++++++--- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/database/gdb/gdb_core.go b/database/gdb/gdb_core.go index 033a9666d..12c0d25c1 100644 --- a/database/gdb/gdb_core.go +++ b/database/gdb/gdb_core.go @@ -1,4 +1,4 @@ -// Copyright GoFrame Author(https://github.com/gogf/gf). All Rights Reserved. +// Copyright GoFrame Author(https://goframe.org). 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, @@ -83,9 +83,13 @@ func (c *Core) Query(sql string, args ...interface{}) (rows *sql.Rows, err error func (c *Core) DoQuery(link Link, sql string, args ...interface{}) (rows *sql.Rows, err error) { sql, args = formatSql(sql, args) sql, args = c.DB.HandleSqlBeforeCommit(link, sql, args) + ctx := c.DB.GetCtx() + if c.GetConfig().QueryTimeout > 0 { + ctx, _ = context.WithTimeout(ctx, c.GetConfig().QueryTimeout) + } if c.DB.GetDebug() { mTime1 := gtime.TimestampMilli() - rows, err = link.QueryContext(c.DB.GetCtx(), sql, args...) + rows, err = link.QueryContext(ctx, sql, args...) mTime2 := gtime.TimestampMilli() s := &Sql{ Sql: sql, @@ -98,7 +102,7 @@ func (c *Core) DoQuery(link Link, sql string, args ...interface{}) (rows *sql.Ro } c.writeSqlToLogger(s) } else { - rows, err = link.QueryContext(c.DB.GetCtx(), sql, args...) + rows, err = link.QueryContext(ctx, sql, args...) } if err == nil { return rows, nil @@ -123,10 +127,14 @@ func (c *Core) Exec(sql string, args ...interface{}) (result sql.Result, err err func (c *Core) DoExec(link Link, sql string, args ...interface{}) (result sql.Result, err error) { sql, args = formatSql(sql, args) sql, args = c.DB.HandleSqlBeforeCommit(link, sql, args) + ctx := c.DB.GetCtx() + if c.GetConfig().ExecTimeout > 0 { + ctx, _ = context.WithTimeout(ctx, c.GetConfig().ExecTimeout) + } if c.DB.GetDebug() { mTime1 := gtime.TimestampMilli() if !c.DB.GetDryRun() { - result, err = link.ExecContext(c.DB.GetCtx(), sql, args...) + result, err = link.ExecContext(ctx, sql, args...) } else { result = new(SqlResult) } @@ -143,7 +151,7 @@ func (c *Core) DoExec(link Link, sql string, args ...interface{}) (result sql.Re c.writeSqlToLogger(s) } else { if !c.DB.GetDryRun() { - result, err = link.ExecContext(c.DB.GetCtx(), sql, args...) + result, err = link.ExecContext(ctx, sql, args...) } else { result = new(SqlResult) } @@ -178,7 +186,11 @@ func (c *Core) Prepare(sql string, execOnMaster ...bool) (*sql.Stmt, error) { // doPrepare calls prepare function on given link object and returns the statement object. func (c *Core) DoPrepare(link Link, sql string) (*sql.Stmt, error) { - return link.PrepareContext(c.DB.GetCtx(), sql) + ctx := c.DB.GetCtx() + if c.GetConfig().QueryTimeout > 0 { + ctx, _ = context.WithTimeout(ctx, c.GetConfig().QueryTimeout) + } + return link.PrepareContext(ctx, sql) } // GetAll queries and returns data records from database. @@ -320,7 +332,11 @@ func (c *Core) Begin() (*TX, error) { if master, err := c.DB.Master(); err != nil { return nil, err } else { - if tx, err := master.Begin(); err == nil { + ctx := c.DB.GetCtx() + if c.GetConfig().TranTimeout > 0 { + ctx, _ = context.WithTimeout(ctx, c.GetConfig().TranTimeout) + } + if tx, err := master.BeginTx(ctx, nil); err == nil { return &TX{ db: c.DB, tx: tx, diff --git a/database/gdb/gdb_core_config.go b/database/gdb/gdb_core_config.go index 85a9ee888..31b8106a1 100644 --- a/database/gdb/gdb_core_config.go +++ b/database/gdb/gdb_core_config.go @@ -1,4 +1,4 @@ -// Copyright GoFrame Author(https://github.com/gogf/gf). All Rights Reserved. +// Copyright GoFrame Author(https://goframe.org). 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, @@ -16,8 +16,7 @@ import ( ) const ( - DEFAULT_GROUP_NAME = "default" // Deprecated, use DefaultGroupName instead. - DefaultGroupName = "default" // Default group name. + DefaultGroupName = "default" // Default group name. ) // Config is the configuration management object. @@ -44,6 +43,10 @@ type ConfigNode struct { MaxIdleConnCount int `json:"maxidle"` // (Optional) Max idle connection configuration for underlying connection pool. MaxOpenConnCount int `json:"maxopen"` // (Optional) Max open connection configuration for underlying connection pool. MaxConnLifetime time.Duration `json:"maxlifetime"` // (Optional) Max connection TTL configuration for underlying connection pool. + QueryTimeout time.Duration // (Optional) Max query time for per dql. + ExecTimeout time.Duration // (Optional) Max exec time for dml. + TranTimeout time.Duration // (Optional) Max exec time time for a transaction. + PrepareTimeout time.Duration // (Optional) Max exec time time for prepare operation. CreatedAt string // (Optional) The filed name of table for automatic-filled created datetime. UpdatedAt string // (Optional) The filed name of table for automatic-filled updated datetime. DeletedAt string // (Optional) The filed name of table for automatic-filled updated datetime.