Merge pull request #24 from TRON-US/support-db-timeout

Add ability to change db timeout on the fly and support passing new db

db/postgres/tx.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	"time"
 
 	"github.com/tron-us/go-common/v2/constant"
 
@@ -11,13 +12,18 @@ import (
 
 // RunInTransactionContext wraps around underlying go-pg's rollback-supported transaction execution
 // with our custom context so it can be easily passed down.
-func (db *TGPGDB) RunInTransactionContext(ctx context.Context, txFunc func(context.Context) error) error {
+func (db *TGPGDB) RunInTransactionContext(ctx context.Context, txFunc func(context.Context, *TGPGDB) error) error {
 	return db.DB.RunInTransaction(func(tx *pg.Tx) error {
 		// Pass ctx with tx object down to the transaction execution
-		return txFunc(context.WithValue(ctx, constant.PostgresTxContext, tx))
+		return txFunc(context.WithValue(ctx, constant.PostgresTxContext, tx), db)
 	})
 }
 
+// WithTimeout needs to create a new db instance in order to pass into the next chain of commands.
+func (db *TGPGDB) WithTimeout(timeout time.Duration) *TGPGDB {
+	return NewTGPGDB(db.DB.WithTimeout(timeout))
+}
+
 // The following functions are necessary to override include to support both transaction
 // and transaction-less queries through the ctx's tx existence.