Revert "🐞 fix(gen): pg gen of insert (#1591)" (#1598)

This reverts commit cc4c4928e0.

ROADMAP.md

@@ -20,9 +20,9 @@ We hope that the items listed below will inspire further engagement from the com
 - [x] Support `goctl bug` to report bugs conveniently
 
 ## 2022
-- [ ] Support `goctl mock` command to start a mocking server with given `.api` file
+- [x] Support `context` in redis related methods for timeout and tracing
+- [x] Support `context` in sql related methods for timeout and tracing
+- [ ] Support `context` in mongodb related methods for timeout and tracing
 - [ ] Add `httpx.Client` with governance, like circuit breaker etc.
 - [ ] Support `goctl doctor` command to report potential issues for given service
-- [ ] Support `context` in redis related methods for timeout and tracing
+- [ ] Support `goctl mock` command to start a mocking server with given `.api` file
-- [ ] Support `context` in sql related methods for timeout and tracing
-- [ ] Support `context` in mongodb related methods for timeout and tracing

core/collection/fifo_test.go

@@ -61,3 +61,41 @@ func TestPutMore(t *testing.T) {
 		assert.Equal(t, string(element), string(body.([]byte)))
 	}
 }
+
+func TestPutMoreWithHeaderNotZero(t *testing.T) {
+	elements := [][]byte{
+		[]byte("hello"),
+		[]byte("world"),
+		[]byte("again"),
+	}
+	queue := NewQueue(4)
+	for i := range elements {
+		queue.Put(elements[i])
+	}
+
+	// take 1
+	body, ok := queue.Take()
+	assert.True(t, ok)
+	element, ok := body.([]byte)
+	assert.True(t, ok)
+	assert.Equal(t, element, []byte("hello"))
+
+	// put more
+	queue.Put([]byte("b4"))
+	queue.Put([]byte("b5")) // will store in elements[0]
+	queue.Put([]byte("b6")) // cause expansion
+
+	results := [][]byte{
+		[]byte("world"),
+		[]byte("again"),
+		[]byte("b4"),
+		[]byte("b5"),
+		[]byte("b6"),
+	}
+
+	for _, element := range results {
+		body, ok := queue.Take()
+		assert.True(t, ok)
+		assert.Equal(t, string(element), string(body.([]byte)))
+	}
+}

core/errorx/batcherror.go

@@ -11,10 +11,12 @@ type (
 	errorArray []error
 )
 
-// Add adds err to be.
+// Add adds errs to be, nil errors are ignored.
-func (be *BatchError) Add(err error) {
+func (be *BatchError) Add(errs ...error) {
-	if err != nil {
+	for _, err := range errs {
-		be.errs = append(be.errs, err)
+		if err != nil {
+			be.errs = append(be.errs, err)
+		}
 	}
 }
 

core/logx/durationlogger.go

@@ -3,6 +3,7 @@ package logx
 import (
 	"fmt"
 	"io"
+	"sync/atomic"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/timex"
@@ -79,7 +80,7 @@ func (l *durationLogger) WithDuration(duration time.Duration) Logger {
 }
 
 func (l *durationLogger) write(writer io.Writer, level string, val interface{}) {
-	switch encoding {
+	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
 		writePlainAny(writer, level, val, l.Duration)
 	default:

core/logx/durationlogger_test.go

@@ -3,6 +3,7 @@ package logx
 import (
 	"log"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -38,10 +39,10 @@ func TestWithDurationInfo(t *testing.T) {
 }
 
 func TestWithDurationInfoConsole(t *testing.T) {
-	old := encoding
+	old := atomic.LoadUint32(&encoding)
-	encoding = plainEncodingType
+	atomic.StoreUint32(&encoding, plainEncodingType)
 	defer func() {
-		encoding = old
+		atomic.StoreUint32(&encoding, old)
 	}()
 
 	var builder strings.Builder

core/logx/logs.go

@@ -75,7 +75,7 @@ var (
 	timeFormat   = "2006-01-02T15:04:05.000Z07:00"
 	writeConsole bool
 	logLevel     uint32
-	encoding     = jsonEncodingType
+	encoding     uint32 = jsonEncodingType
 	// use uint32 for atomic operations
 	disableStat uint32
 	infoLog     io.WriteCloser
@@ -137,9 +137,9 @@ func SetUp(c LogConf) error {
 	}
 	switch c.Encoding {
 	case plainEncoding:
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 	default:
-		encoding = jsonEncodingType
+		atomic.StoreUint32(&encoding, jsonEncodingType)
 	}
 
 	switch c.Mode {
@@ -424,7 +424,7 @@ func infoTextSync(msg string) {
 }
 
 func outputAny(writer io.Writer, level string, val interface{}) {
-	switch encoding {
+	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
 		writePlainAny(writer, level, val)
 	default:
@@ -438,7 +438,7 @@ func outputAny(writer io.Writer, level string, val interface{}) {
 }
 
 func outputText(writer io.Writer, level, msg string) {
-	switch encoding {
+	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
 		writePlainText(writer, level, msg)
 	default:

core/logx/logs_test.go

@@ -145,10 +145,10 @@ func TestStructedLogInfoConsoleAny(t *testing.T) {
 	doTestStructedLogConsole(t, func(writer io.WriteCloser) {
 		infoLog = writer
 	}, func(v ...interface{}) {
-		old := encoding
+		old := atomic.LoadUint32(&encoding)
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 		defer func() {
-			encoding = old
+			atomic.StoreUint32(&encoding, old)
 		}()
 
 		Infov(v)
@@ -159,10 +159,10 @@ func TestStructedLogInfoConsoleAnyString(t *testing.T) {
 	doTestStructedLogConsole(t, func(writer io.WriteCloser) {
 		infoLog = writer
 	}, func(v ...interface{}) {
-		old := encoding
+		old := atomic.LoadUint32(&encoding)
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 		defer func() {
-			encoding = old
+			atomic.StoreUint32(&encoding, old)
 		}()
 
 		Infov(fmt.Sprint(v...))
@@ -173,10 +173,10 @@ func TestStructedLogInfoConsoleAnyError(t *testing.T) {
 	doTestStructedLogConsole(t, func(writer io.WriteCloser) {
 		infoLog = writer
 	}, func(v ...interface{}) {
-		old := encoding
+		old := atomic.LoadUint32(&encoding)
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 		defer func() {
-			encoding = old
+			atomic.StoreUint32(&encoding, old)
 		}()
 
 		Infov(errors.New(fmt.Sprint(v...)))
@@ -187,10 +187,10 @@ func TestStructedLogInfoConsoleAnyStringer(t *testing.T) {
 	doTestStructedLogConsole(t, func(writer io.WriteCloser) {
 		infoLog = writer
 	}, func(v ...interface{}) {
-		old := encoding
+		old := atomic.LoadUint32(&encoding)
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 		defer func() {
-			encoding = old
+			atomic.StoreUint32(&encoding, old)
 		}()
 
 		Infov(ValStringer{
@@ -203,10 +203,10 @@ func TestStructedLogInfoConsoleText(t *testing.T) {
 	doTestStructedLogConsole(t, func(writer io.WriteCloser) {
 		infoLog = writer
 	}, func(v ...interface{}) {
-		old := encoding
+		old := atomic.LoadUint32(&encoding)
-		encoding = plainEncodingType
+		atomic.StoreUint32(&encoding, plainEncodingType)
 		defer func() {
-			encoding = old
+			atomic.StoreUint32(&encoding, old)
 		}()
 
 		Info(fmt.Sprint(v...))

core/logx/syslog_test.go

@@ -29,9 +29,9 @@ func TestRedirector(t *testing.T) {
 }
 
 func captureOutput(f func()) string {
-	atomic.StoreUint32(&initialized, 1)
 	writer := new(mockWriter)
 	infoLog = writer
+	atomic.StoreUint32(&initialized, 1)
 
 	prevLevel := atomic.LoadUint32(&logLevel)
 	SetLevel(InfoLevel)
@@ -44,5 +44,9 @@ func captureOutput(f func()) string {
 func getContent(jsonStr string) string {
 	var entry logEntry
 	json.Unmarshal([]byte(jsonStr), &entry)
-	return entry.Content.(string)
+	val, ok := entry.Content.(string)
+	if ok {
+		return val
+	}
+	return ""
 }

core/logx/tracelogger.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"sync/atomic"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/timex"
@@ -80,7 +81,7 @@ func (l *traceLogger) write(writer io.Writer, level string, val interface{}) {
 	traceID := traceIdFromContext(l.ctx)
 	spanID := spanIdFromContext(l.ctx)
 
-	switch encoding {
+	switch atomic.LoadUint32(&encoding) {
 	case plainEncodingType:
 		writePlainAny(writer, level, val, l.Duration, traceID, spanID)
 	default:

core/logx/tracelogger_test.go

@@ -83,10 +83,10 @@ func TestTraceInfo(t *testing.T) {
 }
 
 func TestTraceInfoConsole(t *testing.T) {
-	old := encoding
+	old := atomic.LoadUint32(&encoding)
-	encoding = plainEncodingType
+	atomic.StoreUint32(&encoding, jsonEncodingType)
 	defer func() {
-		encoding = old
+		atomic.StoreUint32(&encoding, old)
 	}()
 
 	var buf mockWriter

core/mapping/yamlunmarshaler_test.go

@@ -937,7 +937,6 @@ func TestUnmarshalYamlReaderError(t *testing.T) {
 	reader = strings.NewReader("chenquan")
 	err = UnmarshalYamlReader(reader, &v)
 	assert.ErrorIs(t, err, ErrUnsupportedType)
-
 }
 
 func TestUnmarshalYamlBadReader(t *testing.T) {

core/stores/redis/redis.go

@@ -2358,7 +2358,7 @@ func WithTLS() Option {
 }
 
 func acceptable(err error) bool {
-	return err == nil || err == red.Nil
+	return err == nil || err == red.Nil || err == context.Canceled
 }
 
 func getRedis(r *Redis) (RedisNode, error) {

core/stores/sqlc/cachedsql.go

@@ -1,6 +1,7 @@
 package sqlc
 
 import (
+	"context"
 	"database/sql"
 	"time"
 
@@ -18,19 +19,27 @@ var (
 	ErrNotFound = sqlx.ErrNotFound
 
 	// can't use one SingleFlight per conn, because multiple conns may share the same cache key.
-	exclusiveCalls = syncx.NewSingleFlight()
+	singleFlights = syncx.NewSingleFlight()
-	stats          = cache.NewStat("sqlc")
+	stats         = cache.NewStat("sqlc")
 )
 
 type (
 	// ExecFn defines the sql exec method.
 	ExecFn func(conn sqlx.SqlConn) (sql.Result, error)
+	// ExecCtxFn defines the sql exec method.
+	ExecCtxFn func(ctx context.Context, conn sqlx.SqlConn) (sql.Result, error)
 	// IndexQueryFn defines the query method that based on unique indexes.
 	IndexQueryFn func(conn sqlx.SqlConn, v interface{}) (interface{}, error)
+	// IndexQueryCtxFn defines the query method that based on unique indexes.
+	IndexQueryCtxFn func(ctx context.Context, conn sqlx.SqlConn, v interface{}) (interface{}, error)
 	// PrimaryQueryFn defines the query method that based on primary keys.
 	PrimaryQueryFn func(conn sqlx.SqlConn, v, primary interface{}) error
+	// PrimaryQueryCtxFn defines the query method that based on primary keys.
+	PrimaryQueryCtxFn func(ctx context.Context, conn sqlx.SqlConn, v, primary interface{}) error
 	// QueryFn defines the query method.
 	QueryFn func(conn sqlx.SqlConn, v interface{}) error
+	// QueryCtxFn defines the query method.
+	QueryCtxFn func(ctx context.Context, conn sqlx.SqlConn, v interface{}) error
 
 	// A CachedConn is a DB connection with cache capability.
 	CachedConn struct {
@@ -41,7 +50,7 @@ type (
 
 // NewConn returns a CachedConn with a redis cluster cache.
 func NewConn(db sqlx.SqlConn, c cache.CacheConf, opts ...cache.Option) CachedConn {
-	cc := cache.New(c, exclusiveCalls, stats, sql.ErrNoRows, opts...)
+	cc := cache.New(c, singleFlights, stats, sql.ErrNoRows, opts...)
 	return NewConnWithCache(db, cc)
 }
 
@@ -55,28 +64,46 @@ func NewConnWithCache(db sqlx.SqlConn, c cache.Cache) CachedConn {
 
 // NewNodeConn returns a CachedConn with a redis node cache.
 func NewNodeConn(db sqlx.SqlConn, rds *redis.Redis, opts ...cache.Option) CachedConn {
-	c := cache.NewNode(rds, exclusiveCalls, stats, sql.ErrNoRows, opts...)
+	c := cache.NewNode(rds, singleFlights, stats, sql.ErrNoRows, opts...)
 	return NewConnWithCache(db, c)
 }
 
 // DelCache deletes cache with keys.
 func (cc CachedConn) DelCache(keys ...string) error {
-	return cc.cache.Del(keys...)
+	return cc.DelCacheCtx(context.Background(), keys...)
+}
+
+// DelCacheCtx deletes cache with keys.
+func (cc CachedConn) DelCacheCtx(ctx context.Context, keys ...string) error {
+	return cc.cache.DelCtx(ctx, keys...)
 }
 
 // GetCache unmarshals cache with given key into v.
 func (cc CachedConn) GetCache(key string, v interface{}) error {
-	return cc.cache.Get(key, v)
+	return cc.GetCacheCtx(context.Background(), key, v)
+}
+
+// GetCacheCtx unmarshals cache with given key into v.
+func (cc CachedConn) GetCacheCtx(ctx context.Context, key string, v interface{}) error {
+	return cc.cache.GetCtx(ctx, key, v)
 }
 
 // Exec runs given exec on given keys, and returns execution result.
 func (cc CachedConn) Exec(exec ExecFn, keys ...string) (sql.Result, error) {
-	res, err := exec(cc.db)
+	execCtx := func(_ context.Context, conn sqlx.SqlConn) (sql.Result, error) {
+		return exec(conn)
+	}
+	return cc.ExecCtx(context.Background(), execCtx, keys...)
+}
+
+// ExecCtx runs given exec on given keys, and returns execution result.
+func (cc CachedConn) ExecCtx(ctx context.Context, exec ExecCtxFn, keys ...string) (sql.Result, error) {
+	res, err := exec(ctx, cc.db)
 	if err != nil {
 		return nil, err
 	}
 
-	if err := cc.DelCache(keys...); err != nil {
+	if err := cc.DelCacheCtx(ctx, keys...); err != nil {
 		return nil, err
 	}
 
@@ -85,31 +112,61 @@ func (cc CachedConn) Exec(exec ExecFn, keys ...string) (sql.Result, error) {
 
 // ExecNoCache runs exec with given sql statement, without affecting cache.
 func (cc CachedConn) ExecNoCache(q string, args ...interface{}) (sql.Result, error) {
-	return cc.db.Exec(q, args...)
+	return cc.ExecNoCacheCtx(context.Background(), q, args...)
+}
+
+// ExecNoCacheCtx runs exec with given sql statement, without affecting cache.
+func (cc CachedConn) ExecNoCacheCtx(ctx context.Context, q string, args ...interface{}) (
+	sql.Result, error) {
+	return cc.db.ExecCtx(ctx, q, args...)
 }
 
 // QueryRow unmarshals into v with given key and query func.
 func (cc CachedConn) QueryRow(v interface{}, key string, query QueryFn) error {
-	return cc.cache.Take(v, key, func(v interface{}) error {
+	queryCtx := func(_ context.Context, conn sqlx.SqlConn, v interface{}) error {
-		return query(cc.db, v)
+		return query(conn, v)
+	}
+	return cc.QueryRowCtx(context.Background(), v, key, queryCtx)
+}
+
+// QueryRowCtx unmarshals into v with given key and query func.
+func (cc CachedConn) QueryRowCtx(ctx context.Context, v interface{}, key string, query QueryCtxFn) error {
+	return cc.cache.TakeCtx(ctx, v, key, func(v interface{}) error {
+		return query(ctx, cc.db, v)
 	})
 }
 
 // QueryRowIndex unmarshals into v with given key.
 func (cc CachedConn) QueryRowIndex(v interface{}, key string, keyer func(primary interface{}) string,
 	indexQuery IndexQueryFn, primaryQuery PrimaryQueryFn) error {
+	indexQueryCtx := func(_ context.Context, conn sqlx.SqlConn, v interface{}) (interface{}, error) {
+		return indexQuery(conn, v)
+	}
+	primaryQueryCtx := func(_ context.Context, conn sqlx.SqlConn, v, primary interface{}) error {
+		return primaryQuery(conn, v, primary)
+	}
+
+	return cc.QueryRowIndexCtx(context.Background(), v, key, keyer, indexQueryCtx, primaryQueryCtx)
+}
+
+// QueryRowIndexCtx unmarshals into v with given key.
+func (cc CachedConn) QueryRowIndexCtx(ctx context.Context, v interface{}, key string,
+	keyer func(primary interface{}) string, indexQuery IndexQueryCtxFn,
+	primaryQuery PrimaryQueryCtxFn) error {
 	var primaryKey interface{}
 	var found bool
 
-	if err := cc.cache.TakeWithExpire(&primaryKey, key, func(val interface{}, expire time.Duration) (err error) {
+	if err := cc.cache.TakeWithExpireCtx(ctx, &primaryKey, key,
-		primaryKey, err = indexQuery(cc.db, v)
+		func(val interface{}, expire time.Duration) (err error) {
-		if err != nil {
+			primaryKey, err = indexQuery(ctx, cc.db, v)
-			return
+			if err != nil {
-		}
+				return
+			}
 
-		found = true
+			found = true
-		return cc.cache.SetWithExpire(keyer(primaryKey), v, expire+cacheSafeGapBetweenIndexAndPrimary)
+			return cc.cache.SetWithExpireCtx(ctx, keyer(primaryKey), v,
-	}); err != nil {
+				expire+cacheSafeGapBetweenIndexAndPrimary)
+		}); err != nil {
 		return err
 	}
 
@@ -117,28 +174,54 @@ func (cc CachedConn) QueryRowIndex(v interface{}, key string, keyer func(primary
 		return nil
 	}
 
-	return cc.cache.Take(v, keyer(primaryKey), func(v interface{}) error {
+	return cc.cache.TakeCtx(ctx, v, keyer(primaryKey), func(v interface{}) error {
-		return primaryQuery(cc.db, v, primaryKey)
+		return primaryQuery(ctx, cc.db, v, primaryKey)
 	})
 }
 
 // QueryRowNoCache unmarshals into v with given statement.
 func (cc CachedConn) QueryRowNoCache(v interface{}, q string, args ...interface{}) error {
-	return cc.db.QueryRow(v, q, args...)
+	return cc.QueryRowNoCacheCtx(context.Background(), v, q, args...)
+}
+
+// QueryRowNoCacheCtx unmarshals into v with given statement.
+func (cc CachedConn) QueryRowNoCacheCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return cc.db.QueryRowCtx(ctx, v, q, args...)
 }
 
 // QueryRowsNoCache unmarshals into v with given statement.
 // It doesn't use cache, because it might cause consistency problem.
 func (cc CachedConn) QueryRowsNoCache(v interface{}, q string, args ...interface{}) error {
-	return cc.db.QueryRows(v, q, args...)
+	return cc.QueryRowsNoCacheCtx(context.Background(), v, q, args...)
+}
+
+// QueryRowsNoCacheCtx unmarshals into v with given statement.
+// It doesn't use cache, because it might cause consistency problem.
+func (cc CachedConn) QueryRowsNoCacheCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return cc.db.QueryRowsCtx(ctx, v, q, args...)
 }
 
 // SetCache sets v into cache with given key.
-func (cc CachedConn) SetCache(key string, v interface{}) error {
+func (cc CachedConn) SetCache(key string, val interface{}) error {
-	return cc.cache.Set(key, v)
+	return cc.SetCacheCtx(context.Background(), key, val)
+}
+
+// SetCacheCtx sets v into cache with given key.
+func (cc CachedConn) SetCacheCtx(ctx context.Context, key string, val interface{}) error {
+	return cc.cache.SetCtx(ctx, key, val)
 }
 
 // Transact runs given fn in transaction mode.
 func (cc CachedConn) Transact(fn func(sqlx.Session) error) error {
-	return cc.db.Transact(fn)
+	fnCtx := func(_ context.Context, session sqlx.Session) error {
+		return fn(session)
+	}
+	return cc.TransactCtx(context.Background(), fnCtx)
+}
+
+// TransactCtx runs given fn in transaction mode.
+func (cc CachedConn) TransactCtx(ctx context.Context, fn func(context.Context, sqlx.Session) error) error {
+	return cc.db.TransactCtx(ctx, fn)
 }

core/stores/sqlc/cachedsql_test.go

@@ -1,6 +1,7 @@
 package sqlc
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"errors"
@@ -568,7 +569,7 @@ func TestNewConnWithCache(t *testing.T) {
 	defer clean()
 
 	var conn trackedConn
-	c := NewConnWithCache(&conn, cache.NewNode(r, exclusiveCalls, stats, sql.ErrNoRows))
+	c := NewConnWithCache(&conn, cache.NewNode(r, singleFlights, stats, sql.ErrNoRows))
 	_, err = c.ExecNoCache("delete from user_table where id='kevin'")
 	assert.Nil(t, err)
 	assert.True(t, conn.execValue)
@@ -585,6 +586,30 @@ type dummySqlConn struct {
 	queryRow func(interface{}, string, ...interface{}) error
 }
 
+func (d dummySqlConn) ExecCtx(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	return nil, nil
+}
+
+func (d dummySqlConn) PrepareCtx(ctx context.Context, query string) (sqlx.StmtSession, error) {
+	return nil, nil
+}
+
+func (d dummySqlConn) QueryRowPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
+func (d dummySqlConn) QueryRowsCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
+func (d dummySqlConn) QueryRowsPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
+func (d dummySqlConn) TransactCtx(ctx context.Context, fn func(context.Context, sqlx.Session) error) error {
+	return nil
+}
+
 func (d dummySqlConn) Exec(query string, args ...interface{}) (sql.Result, error) {
 	return nil, nil
 }
@@ -594,6 +619,10 @@ func (d dummySqlConn) Prepare(query string) (sqlx.StmtSession, error) {
 }
 
 func (d dummySqlConn) QueryRow(v interface{}, query string, args ...interface{}) error {
+	return d.QueryRowCtx(context.Background(), v, query, args...)
+}
+
+func (d dummySqlConn) QueryRowCtx(_ context.Context, v interface{}, query string, args ...interface{}) error {
 	if d.queryRow != nil {
 		return d.queryRow(v, query, args...)
 	}
@@ -628,13 +657,21 @@ type trackedConn struct {
 }
 
 func (c *trackedConn) Exec(query string, args ...interface{}) (sql.Result, error) {
+	return c.ExecCtx(context.Background(), query, args...)
+}
+
+func (c *trackedConn) ExecCtx(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
 	c.execValue = true
-	return c.dummySqlConn.Exec(query, args...)
+	return c.dummySqlConn.ExecCtx(ctx, query, args...)
 }
 
 func (c *trackedConn) QueryRows(v interface{}, query string, args ...interface{}) error {
+	return c.QueryRowsCtx(context.Background(), v, query, args...)
+}
+
+func (c *trackedConn) QueryRowsCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
 	c.queryRowsValue = true
-	return c.dummySqlConn.QueryRows(v, query, args...)
+	return c.dummySqlConn.QueryRowsCtx(ctx, v, query, args...)
 }
 
 func (c *trackedConn) RawDB() (*sql.DB, error) {
@@ -642,6 +679,12 @@ func (c *trackedConn) RawDB() (*sql.DB, error) {
 }
 
 func (c *trackedConn) Transact(fn func(session sqlx.Session) error) error {
-	c.transactValue = true
+	return c.TransactCtx(context.Background(), func(_ context.Context, session sqlx.Session) error {
-	return c.dummySqlConn.Transact(fn)
+		return fn(session)
+	})
+}
+
+func (c *trackedConn) TransactCtx(ctx context.Context, fn func(context.Context, sqlx.Session) error) error {
+	c.transactValue = true
+	return c.dummySqlConn.TransactCtx(ctx, fn)
 }

core/stores/sqlx/bulkinserter_test.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"strconv"
@@ -17,12 +18,40 @@ type mockedConn struct {
 	execErr error
 }
 
-func (c *mockedConn) Exec(query string, args ...interface{}) (sql.Result, error) {
+func (c *mockedConn) ExecCtx(_ context.Context, query string, args ...interface{}) (sql.Result, error) {
 	c.query = query
 	c.args = args
 	return nil, c.execErr
 }
 
+func (c *mockedConn) PrepareCtx(ctx context.Context, query string) (StmtSession, error) {
+	panic("implement me")
+}
+
+func (c *mockedConn) QueryRowCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	panic("implement me")
+}
+
+func (c *mockedConn) QueryRowPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	panic("implement me")
+}
+
+func (c *mockedConn) QueryRowsCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	panic("implement me")
+}
+
+func (c *mockedConn) QueryRowsPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	panic("implement me")
+}
+
+func (c *mockedConn) TransactCtx(ctx context.Context, fn func(context.Context, Session) error) error {
+	panic("should not called")
+}
+
+func (c *mockedConn) Exec(query string, args ...interface{}) (sql.Result, error) {
+	return c.ExecCtx(context.Background(), query, args...)
+}
+
 func (c *mockedConn) Prepare(query string) (StmtSession, error) {
 	panic("should not called")
 }

core/stores/sqlx/orm_test.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"testing"
@@ -16,7 +17,7 @@ func TestUnmarshalRowBool(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value bool
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.True(t, value)
@@ -29,7 +30,7 @@ func TestUnmarshalRowBoolNotSettable(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value bool
-		assert.NotNil(t, query(db, func(rows *sql.Rows) error {
+		assert.NotNil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 	})
@@ -41,7 +42,7 @@ func TestUnmarshalRowInt(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value int
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, 2, value)
@@ -54,7 +55,7 @@ func TestUnmarshalRowInt8(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value int8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, int8(3), value)
@@ -67,7 +68,7 @@ func TestUnmarshalRowInt16(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value int16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.Equal(t, int16(4), value)
@@ -80,7 +81,7 @@ func TestUnmarshalRowInt32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value int32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.Equal(t, int32(5), value)
@@ -93,7 +94,7 @@ func TestUnmarshalRowInt64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value int64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, int64(6), value)
@@ -106,7 +107,7 @@ func TestUnmarshalRowUint(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value uint
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, uint(2), value)
@@ -119,7 +120,7 @@ func TestUnmarshalRowUint8(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value uint8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, uint8(3), value)
@@ -132,7 +133,7 @@ func TestUnmarshalRowUint16(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value uint16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, uint16(4), value)
@@ -145,7 +146,7 @@ func TestUnmarshalRowUint32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value uint32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, uint32(5), value)
@@ -158,7 +159,7 @@ func TestUnmarshalRowUint64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value uint64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, uint16(6), value)
@@ -171,7 +172,7 @@ func TestUnmarshalRowFloat32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value float32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, float32(7), value)
@@ -184,7 +185,7 @@ func TestUnmarshalRowFloat64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value float64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, float64(8), value)
@@ -198,7 +199,7 @@ func TestUnmarshalRowString(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value string
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -215,7 +216,7 @@ func TestUnmarshalRowStruct(t *testing.T) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("liao,5")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 		assert.Equal(t, "liao", value.Name)
@@ -233,7 +234,7 @@ func TestUnmarshalRowStructWithTags(t *testing.T) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("liao,5")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 		assert.Equal(t, "liao", value.Name)
@@ -251,7 +252,7 @@ func TestUnmarshalRowStructWithTagsWrongColumns(t *testing.T) {
 		rs := sqlmock.NewRows([]string{"name"}).FromCSVString("liao")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
-		assert.NotNil(t, query(db, func(rows *sql.Rows) error {
+		assert.NotNil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 	})
@@ -264,7 +265,7 @@ func TestUnmarshalRowsBool(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []bool
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -278,7 +279,7 @@ func TestUnmarshalRowsInt(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []int
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -292,7 +293,7 @@ func TestUnmarshalRowsInt8(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []int8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -306,7 +307,7 @@ func TestUnmarshalRowsInt16(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []int16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -320,7 +321,7 @@ func TestUnmarshalRowsInt32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []int32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -334,7 +335,7 @@ func TestUnmarshalRowsInt64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []int64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -348,7 +349,7 @@ func TestUnmarshalRowsUint(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []uint
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -362,7 +363,7 @@ func TestUnmarshalRowsUint8(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []uint8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -376,7 +377,7 @@ func TestUnmarshalRowsUint16(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []uint16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -390,7 +391,7 @@ func TestUnmarshalRowsUint32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []uint32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -404,7 +405,7 @@ func TestUnmarshalRowsUint64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []uint64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -418,7 +419,7 @@ func TestUnmarshalRowsFloat32(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []float32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -432,7 +433,7 @@ func TestUnmarshalRowsFloat64(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []float64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -446,7 +447,7 @@ func TestUnmarshalRowsString(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []string
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -462,7 +463,7 @@ func TestUnmarshalRowsBoolPtr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*bool
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -478,7 +479,7 @@ func TestUnmarshalRowsIntPtr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*int
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -494,7 +495,7 @@ func TestUnmarshalRowsInt8Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*int8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -510,7 +511,7 @@ func TestUnmarshalRowsInt16Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*int16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -526,7 +527,7 @@ func TestUnmarshalRowsInt32Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*int32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -542,7 +543,7 @@ func TestUnmarshalRowsInt64Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*int64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -558,7 +559,7 @@ func TestUnmarshalRowsUintPtr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*uint
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -574,7 +575,7 @@ func TestUnmarshalRowsUint8Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*uint8
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -590,7 +591,7 @@ func TestUnmarshalRowsUint16Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*uint16
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -606,7 +607,7 @@ func TestUnmarshalRowsUint32Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*uint32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -622,7 +623,7 @@ func TestUnmarshalRowsUint64Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*uint64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -638,7 +639,7 @@ func TestUnmarshalRowsFloat32Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*float32
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -654,7 +655,7 @@ func TestUnmarshalRowsFloat64Ptr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*float64
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -670,7 +671,7 @@ func TestUnmarshalRowsStringPtr(t *testing.T) {
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
 
 		var value []*string
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select value from users where user=?", "anyone"))
 		assert.EqualValues(t, expect, value)
@@ -699,7 +700,7 @@ func TestUnmarshalRowsStruct(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("first,2\nsecond,3")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -739,7 +740,7 @@ func TestUnmarshalRowsStructWithNullStringType(t *testing.T) {
 		rs := sqlmock.NewRows([]string{"name", "value"}).AddRow(
 			"first", "firstnullstring").AddRow("second", nil)
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -773,7 +774,7 @@ func TestUnmarshalRowsStructWithTags(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("first,2\nsecond,3")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -814,7 +815,7 @@ func TestUnmarshalRowsStructAndEmbeddedAnonymousStructWithTags(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age", "value"}).FromCSVString("first,2,3\nsecond,3,4")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age, value from users where user=?", "anyone"))
 
@@ -856,7 +857,7 @@ func TestUnmarshalRowsStructAndEmbeddedStructPtrAnonymousWithTags(t *testing.T)
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age", "value"}).FromCSVString("first,2,3\nsecond,3,4")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age, value from users where user=?", "anyone"))
 
@@ -890,7 +891,7 @@ func TestUnmarshalRowsStructPtr(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("first,2\nsecond,3")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -923,7 +924,7 @@ func TestUnmarshalRowsStructWithTagsPtr(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("first,2\nsecond,3")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -956,7 +957,7 @@ func TestUnmarshalRowsStructWithTagsPtrWithInnerPtr(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("first,2\nsecond,3")
 		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRows(&value, rows, true)
 		}, "select name, age from users where user=?", "anyone"))
 
@@ -976,7 +977,7 @@ func TestCommonSqlConn_QueryRowOptional(t *testing.T) {
 			User string `db:"user"`
 			Age  int    `db:"age"`
 		}
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
+		assert.Nil(t, query(context.Background(), db, func(rows *sql.Rows) error {
 			return unmarshalRow(&r, rows, false)
 		}, "select age from users where user=?", "anyone"))
 		assert.Empty(t, r.User)
@@ -1027,7 +1028,7 @@ func TestUnmarshalRowError(t *testing.T) {
 					User string `db:"user"`
 					Age  int    `db:"age"`
 				}
-				test.validate(query(db, func(rows *sql.Rows) error {
+				test.validate(query(context.Background(), db, func(rows *sql.Rows) error {
 					scanner := mockedScanner{
 						colErr:  test.colErr,
 						scanErr: test.scanErr,

core/stores/sqlx/sqlconn.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 
 	"github.com/zeromicro/go-zero/core/breaker"
@@ -14,11 +15,17 @@ type (
 	// Session stands for raw connections or transaction sessions
 	Session interface {
 		Exec(query string, args ...interface{}) (sql.Result, error)
+		ExecCtx(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
 		Prepare(query string) (StmtSession, error)
+		PrepareCtx(ctx context.Context, query string) (StmtSession, error)
 		QueryRow(v interface{}, query string, args ...interface{}) error
+		QueryRowCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error
 		QueryRowPartial(v interface{}, query string, args ...interface{}) error
+		QueryRowPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error
 		QueryRows(v interface{}, query string, args ...interface{}) error
+		QueryRowsCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error
 		QueryRowsPartial(v interface{}, query string, args ...interface{}) error
+		QueryRowsPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error
 	}
 
 	// SqlConn only stands for raw connections, so Transact method can be called.
@@ -27,7 +34,8 @@ type (
 		// RawDB is for other ORM to operate with, use it with caution.
 		// Notice: don't close it.
 		RawDB() (*sql.DB, error)
-		Transact(func(session Session) error) error
+		Transact(fn func(Session) error) error
+		TransactCtx(ctx context.Context, fn func(context.Context, Session) error) error
 	}
 
 	// SqlOption defines the method to customize a sql connection.
@@ -37,10 +45,15 @@ type (
 	StmtSession interface {
 		Close() error
 		Exec(args ...interface{}) (sql.Result, error)
+		ExecCtx(ctx context.Context, args ...interface{}) (sql.Result, error)
 		QueryRow(v interface{}, args ...interface{}) error
+		QueryRowCtx(ctx context.Context, v interface{}, args ...interface{}) error
 		QueryRowPartial(v interface{}, args ...interface{}) error
+		QueryRowPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error
 		QueryRows(v interface{}, args ...interface{}) error
+		QueryRowsCtx(ctx context.Context, v interface{}, args ...interface{}) error
 		QueryRowsPartial(v interface{}, args ...interface{}) error
+		QueryRowsPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error
 	}
 
 	// thread-safe
@@ -58,7 +71,9 @@ type (
 
 	sessionConn interface {
 		Exec(query string, args ...interface{}) (sql.Result, error)
+		ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
 		Query(query string, args ...interface{}) (*sql.Rows, error)
+		QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
 	}
 
 	statement struct {
@@ -68,7 +83,9 @@ type (
 
 	stmtConn interface {
 		Exec(args ...interface{}) (sql.Result, error)
+		ExecContext(ctx context.Context, args ...interface{}) (sql.Result, error)
 		Query(args ...interface{}) (*sql.Rows, error)
+		QueryContext(ctx context.Context, args ...interface{}) (*sql.Rows, error)
 	}
 )
 
@@ -112,6 +129,11 @@ func NewSqlConnFromDB(db *sql.DB, opts ...SqlOption) SqlConn {
 }
 
 func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result, err error) {
+	return db.ExecCtx(context.Background(), q, args...)
+}
+
+func (db *commonSqlConn) ExecCtx(ctx context.Context, q string, args ...interface{}) (
+	result sql.Result, err error) {
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
@@ -120,7 +142,7 @@ func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result,
 			return err
 		}
 
-		result, err = exec(conn, q, args...)
+		result, err = exec(ctx, conn, q, args...)
 		return err
 	}, db.acceptable)
 
@@ -128,6 +150,10 @@ func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result,
 }
 
 func (db *commonSqlConn) Prepare(query string) (stmt StmtSession, err error) {
+	return db.PrepareCtx(context.Background(), query)
+}
+
+func (db *commonSqlConn) PrepareCtx(ctx context.Context, query string) (stmt StmtSession, err error) {
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
@@ -136,7 +162,7 @@ func (db *commonSqlConn) Prepare(query string) (stmt StmtSession, err error) {
 			return err
 		}
 
-		st, err := conn.Prepare(query)
+		st, err := conn.PrepareContext(ctx, query)
 		if err != nil {
 			return err
 		}
@@ -152,25 +178,45 @@ func (db *commonSqlConn) Prepare(query string) (stmt StmtSession, err error) {
 }
 
 func (db *commonSqlConn) QueryRow(v interface{}, q string, args ...interface{}) error {
-	return db.queryRows(func(rows *sql.Rows) error {
+	return db.QueryRowCtx(context.Background(), v, q, args...)
+}
+
+func (db *commonSqlConn) QueryRowCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, q, args...)
 }
 
 func (db *commonSqlConn) QueryRowPartial(v interface{}, q string, args ...interface{}) error {
-	return db.queryRows(func(rows *sql.Rows) error {
+	return db.QueryRowPartialCtx(context.Background(), v, q, args...)
+}
+
+func (db *commonSqlConn) QueryRowPartialCtx(ctx context.Context, v interface{},
+	q string, args ...interface{}) error {
+	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, q, args...)
 }
 
 func (db *commonSqlConn) QueryRows(v interface{}, q string, args ...interface{}) error {
-	return db.queryRows(func(rows *sql.Rows) error {
+	return db.QueryRowsCtx(context.Background(), v, q, args...)
+}
+
+func (db *commonSqlConn) QueryRowsCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, q, args...)
 }
 
 func (db *commonSqlConn) QueryRowsPartial(v interface{}, q string, args ...interface{}) error {
-	return db.queryRows(func(rows *sql.Rows) error {
+	return db.QueryRowsPartialCtx(context.Background(), v, q, args...)
+}
+
+func (db *commonSqlConn) QueryRowsPartialCtx(ctx context.Context, v interface{},
+	q string, args ...interface{}) error {
+	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, q, args...)
 }
@@ -180,13 +226,19 @@ func (db *commonSqlConn) RawDB() (*sql.DB, error) {
 }
 
 func (db *commonSqlConn) Transact(fn func(Session) error) error {
+	return db.TransactCtx(context.Background(), func(_ context.Context, session Session) error {
+		return fn(session)
+	})
+}
+
+func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Context, Session) error) error {
 	return db.brk.DoWithAcceptable(func() error {
-		return transact(db, db.beginTx, fn)
+		return transact(ctx, db, db.beginTx, fn)
 	}, db.acceptable)
 }
 
 func (db *commonSqlConn) acceptable(err error) bool {
-	ok := err == nil || err == sql.ErrNoRows || err == sql.ErrTxDone
+	ok := err == nil || err == sql.ErrNoRows || err == sql.ErrTxDone || err == context.Canceled
 	if db.accept == nil {
 		return ok
 	}
@@ -194,7 +246,8 @@ func (db *commonSqlConn) acceptable(err error) bool {
 	return ok || db.accept(err)
 }
 
-func (db *commonSqlConn) queryRows(scanner func(*sql.Rows) error, q string, args ...interface{}) error {
+func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows) error,
+	q string, args ...interface{}) error {
 	var qerr error
 	return db.brk.DoWithAcceptable(func() error {
 		conn, err := db.connProv()
@@ -203,7 +256,7 @@ func (db *commonSqlConn) queryRows(scanner func(*sql.Rows) error, q string, args
 			return err
 		}
 
-		return query(conn, func(rows *sql.Rows) error {
+		return query(ctx, conn, func(rows *sql.Rows) error {
 			qerr = scanner(rows)
 			return qerr
 		}, q, args...)
@@ -217,29 +270,49 @@ func (s statement) Close() error {
 }
 
 func (s statement) Exec(args ...interface{}) (sql.Result, error) {
-	return execStmt(s.stmt, s.query, args...)
+	return s.ExecCtx(context.Background(), args...)
+}
+
+func (s statement) ExecCtx(ctx context.Context, args ...interface{}) (sql.Result, error) {
+	return execStmt(ctx, s.stmt, s.query, args...)
 }
 
 func (s statement) QueryRow(v interface{}, args ...interface{}) error {
-	return queryStmt(s.stmt, func(rows *sql.Rows) error {
+	return s.QueryRowCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, s.query, args...)
 }
 
 func (s statement) QueryRowPartial(v interface{}, args ...interface{}) error {
-	return queryStmt(s.stmt, func(rows *sql.Rows) error {
+	return s.QueryRowPartialCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, s.query, args...)
 }
 
 func (s statement) QueryRows(v interface{}, args ...interface{}) error {
-	return queryStmt(s.stmt, func(rows *sql.Rows) error {
+	return s.QueryRowsCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowsCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, s.query, args...)
 }
 
 func (s statement) QueryRowsPartial(v interface{}, args ...interface{}) error {
-	return queryStmt(s.stmt, func(rows *sql.Rows) error {
+	return s.QueryRowsPartialCtx(context.Background(), v, args...)
+}
+
+func (s statement) QueryRowsPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, s.query, args...)
 }

core/stores/sqlx/stmt.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"time"
 
@@ -18,64 +19,65 @@ func SetSlowThreshold(threshold time.Duration) {
 	slowThreshold.Set(threshold)
 }
 
-func exec(conn sessionConn, q string, args ...interface{}) (sql.Result, error) {
+func exec(ctx context.Context, conn sessionConn, q string, args ...interface{}) (sql.Result, error) {
 	stmt, err := format(q, args...)
 	if err != nil {
 		return nil, err
 	}
 
 	startTime := timex.Now()
-	result, err := conn.Exec(q, args...)
+	result, err := conn.ExecContext(ctx, q, args...)
 	duration := timex.Since(startTime)
 	if duration > slowThreshold.Load() {
-		logx.WithDuration(duration).Slowf("[SQL] exec: slowcall - %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] exec: slowcall - %s", stmt)
 	} else {
-		logx.WithDuration(duration).Infof("sql exec: %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Infof("sql exec: %s", stmt)
 	}
 	if err != nil {
-		logSqlError(stmt, err)
+		logSqlError(ctx, stmt, err)
 	}
 
 	return result, err
 }
 
-func execStmt(conn stmtConn, q string, args ...interface{}) (sql.Result, error) {
+func execStmt(ctx context.Context, conn stmtConn, q string, args ...interface{}) (sql.Result, error) {
 	stmt, err := format(q, args...)
 	if err != nil {
 		return nil, err
 	}
 
 	startTime := timex.Now()
-	result, err := conn.Exec(args...)
+	result, err := conn.ExecContext(ctx, args...)
 	duration := timex.Since(startTime)
 	if duration > slowThreshold.Load() {
-		logx.WithDuration(duration).Slowf("[SQL] execStmt: slowcall - %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] execStmt: slowcall - %s", stmt)
 	} else {
-		logx.WithDuration(duration).Infof("sql execStmt: %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Infof("sql execStmt: %s", stmt)
 	}
 	if err != nil {
-		logSqlError(stmt, err)
+		logSqlError(ctx, stmt, err)
 	}
 
 	return result, err
 }
 
-func query(conn sessionConn, scanner func(*sql.Rows) error, q string, args ...interface{}) error {
+func query(ctx context.Context, conn sessionConn, scanner func(*sql.Rows) error,
+	q string, args ...interface{}) error {
 	stmt, err := format(q, args...)
 	if err != nil {
 		return err
 	}
 
 	startTime := timex.Now()
-	rows, err := conn.Query(q, args...)
+	rows, err := conn.QueryContext(ctx, q, args...)
 	duration := timex.Since(startTime)
 	if duration > slowThreshold.Load() {
-		logx.WithDuration(duration).Slowf("[SQL] query: slowcall - %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] query: slowcall - %s", stmt)
 	} else {
-		logx.WithDuration(duration).Infof("sql query: %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Infof("sql query: %s", stmt)
 	}
 	if err != nil {
-		logSqlError(stmt, err)
+		logSqlError(ctx, stmt, err)
 		return err
 	}
 	defer rows.Close()
@@ -83,22 +85,23 @@ func query(conn sessionConn, scanner func(*sql.Rows) error, q string, args ...in
 	return scanner(rows)
 }
 
-func queryStmt(conn stmtConn, scanner func(*sql.Rows) error, q string, args ...interface{}) error {
+func queryStmt(ctx context.Context, conn stmtConn, scanner func(*sql.Rows) error,
+	q string, args ...interface{}) error {
 	stmt, err := format(q, args...)
 	if err != nil {
 		return err
 	}
 
 	startTime := timex.Now()
-	rows, err := conn.Query(args...)
+	rows, err := conn.QueryContext(ctx, args...)
 	duration := timex.Since(startTime)
 	if duration > slowThreshold.Load() {
-		logx.WithDuration(duration).Slowf("[SQL] queryStmt: slowcall - %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] queryStmt: slowcall - %s", stmt)
 	} else {
-		logx.WithDuration(duration).Infof("sql queryStmt: %s", stmt)
+		logx.WithContext(ctx).WithDuration(duration).Infof("sql queryStmt: %s", stmt)
 	}
 	if err != nil {
-		logSqlError(stmt, err)
+		logSqlError(ctx, stmt, err)
 		return err
 	}
 	defer rows.Close()

core/stores/sqlx/stmt_test.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"testing"
@@ -57,7 +58,7 @@ func TestStmt_exec(t *testing.T) {
 		test := test
 		fns := []func(args ...interface{}) (sql.Result, error){
 			func(args ...interface{}) (sql.Result, error) {
-				return exec(&mockedSessionConn{
+				return exec(context.Background(), &mockedSessionConn{
 					lastInsertId: test.lastInsertId,
 					rowsAffected: test.rowsAffected,
 					err:          test.err,
@@ -65,7 +66,7 @@ func TestStmt_exec(t *testing.T) {
 				}, test.query, args...)
 			},
 			func(args ...interface{}) (sql.Result, error) {
-				return execStmt(&mockedStmtConn{
+				return execStmt(context.Background(), &mockedStmtConn{
 					lastInsertId: test.lastInsertId,
 					rowsAffected: test.rowsAffected,
 					err:          test.err,
@@ -137,7 +138,7 @@ func TestStmt_query(t *testing.T) {
 		test := test
 		fns := []func(args ...interface{}) error{
 			func(args ...interface{}) error {
-				return query(&mockedSessionConn{
+				return query(context.Background(), &mockedSessionConn{
 					err:   test.err,
 					delay: test.delay,
 				}, func(rows *sql.Rows) error {
@@ -145,7 +146,7 @@ func TestStmt_query(t *testing.T) {
 				}, test.query, args...)
 			},
 			func(args ...interface{}) error {
-				return queryStmt(&mockedStmtConn{
+				return queryStmt(context.Background(), &mockedStmtConn{
 					err:   test.err,
 					delay: test.delay,
 				}, func(rows *sql.Rows) error {
@@ -185,6 +186,10 @@ type mockedSessionConn struct {
 }
 
 func (m *mockedSessionConn) Exec(query string, args ...interface{}) (sql.Result, error) {
+	return m.ExecContext(context.Background(), query, args...)
+}
+
+func (m *mockedSessionConn) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
 	if m.delay {
 		time.Sleep(defaultSlowThreshold + time.Millisecond)
 	}
@@ -195,6 +200,10 @@ func (m *mockedSessionConn) Exec(query string, args ...interface{}) (sql.Result,
 }
 
 func (m *mockedSessionConn) Query(query string, args ...interface{}) (*sql.Rows, error) {
+	return m.QueryContext(context.Background(), query, args...)
+}
+
+func (m *mockedSessionConn) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
 	if m.delay {
 		time.Sleep(defaultSlowThreshold + time.Millisecond)
 	}
@@ -214,6 +223,10 @@ type mockedStmtConn struct {
 }
 
 func (m *mockedStmtConn) Exec(args ...interface{}) (sql.Result, error) {
+	return m.ExecContext(context.Background(), args...)
+}
+
+func (m *mockedStmtConn) ExecContext(_ context.Context, _ ...interface{}) (sql.Result, error) {
 	if m.delay {
 		time.Sleep(defaultSlowThreshold + time.Millisecond)
 	}
@@ -224,6 +237,10 @@ func (m *mockedStmtConn) Exec(args ...interface{}) (sql.Result, error) {
 }
 
 func (m *mockedStmtConn) Query(args ...interface{}) (*sql.Rows, error) {
+	return m.QueryContext(context.Background(), args...)
+}
+
+func (m *mockedStmtConn) QueryContext(_ context.Context, _ ...interface{}) (*sql.Rows, error) {
 	if m.delay {
 		time.Sleep(defaultSlowThreshold + time.Millisecond)
 	}

core/stores/sqlx/tx.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 )
@@ -26,11 +27,19 @@ func NewSessionFromTx(tx *sql.Tx) Session {
 }
 
 func (t txSession) Exec(q string, args ...interface{}) (sql.Result, error) {
-	return exec(t.Tx, q, args...)
+	return t.ExecCtx(context.Background(), q, args...)
+}
+
+func (t txSession) ExecCtx(ctx context.Context, q string, args ...interface{}) (sql.Result, error) {
+	return exec(ctx, t.Tx, q, args...)
 }
 
 func (t txSession) Prepare(q string) (StmtSession, error) {
-	stmt, err := t.Tx.Prepare(q)
+	return t.PrepareCtx(context.Background(), q)
+}
+
+func (t txSession) PrepareCtx(ctx context.Context, q string) (StmtSession, error) {
+	stmt, err := t.Tx.PrepareContext(ctx, q)
 	if err != nil {
 		return nil, err
 	}
@@ -42,25 +51,43 @@ func (t txSession) Prepare(q string) (StmtSession, error) {
 }
 
 func (t txSession) QueryRow(v interface{}, q string, args ...interface{}) error {
-	return query(t.Tx, func(rows *sql.Rows) error {
+	return t.QueryRowCtx(context.Background(), v, q, args...)
+}
+
+func (t txSession) QueryRowCtx(ctx context.Context, v interface{}, q string, args ...interface{}) error {
+	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, q, args...)
 }
 
 func (t txSession) QueryRowPartial(v interface{}, q string, args ...interface{}) error {
-	return query(t.Tx, func(rows *sql.Rows) error {
+	return t.QueryRowPartialCtx(context.Background(), v, q, args...)
+}
+
+func (t txSession) QueryRowPartialCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, q, args...)
 }
 
 func (t txSession) QueryRows(v interface{}, q string, args ...interface{}) error {
-	return query(t.Tx, func(rows *sql.Rows) error {
+	return t.QueryRowsCtx(context.Background(), v, q, args...)
+}
+
+func (t txSession) QueryRowsCtx(ctx context.Context, v interface{}, q string, args ...interface{}) error {
+	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, q, args...)
 }
 
 func (t txSession) QueryRowsPartial(v interface{}, q string, args ...interface{}) error {
-	return query(t.Tx, func(rows *sql.Rows) error {
+	return t.QueryRowsPartialCtx(context.Background(), v, q, args...)
+}
+
+func (t txSession) QueryRowsPartialCtx(ctx context.Context, v interface{}, q string,
+	args ...interface{}) error {
+	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, q, args...)
 }
@@ -76,17 +103,19 @@ func begin(db *sql.DB) (trans, error) {
 	}, nil
 }
 
-func transact(db *commonSqlConn, b beginnable, fn func(Session) error) (err error) {
+func transact(ctx context.Context, db *commonSqlConn, b beginnable,
+	fn func(context.Context, Session) error) (err error) {
 	conn, err := db.connProv()
 	if err != nil {
 		db.onError(err)
 		return err
 	}
 
-	return transactOnConn(conn, b, fn)
+	return transactOnConn(ctx, conn, b, fn)
 }
 
-func transactOnConn(conn *sql.DB, b beginnable, fn func(Session) error) (err error) {
+func transactOnConn(ctx context.Context, conn *sql.DB, b beginnable,
+	fn func(context.Context, Session) error) (err error) {
 	var tx trans
 	tx, err = b(conn)
 	if err != nil {
@@ -96,18 +125,18 @@ func transactOnConn(conn *sql.DB, b beginnable, fn func(Session) error) (err err
 	defer func() {
 		if p := recover(); p != nil {
 			if e := tx.Rollback(); e != nil {
-				err = fmt.Errorf("recover from %#v, rollback failed: %s", p, e)
+				err = fmt.Errorf("recover from %#v, rollback failed: %w", p, e)
 			} else {
 				err = fmt.Errorf("recoveer from %#v", p)
 			}
 		} else if err != nil {
 			if e := tx.Rollback(); e != nil {
-				err = fmt.Errorf("transaction failed: %s, rollback failed: %s", err, e)
+				err = fmt.Errorf("transaction failed: %s, rollback failed: %w", err, e)
 			}
 		} else {
 			err = tx.Commit()
 		}
 	}()
 
-	return fn(tx)
+	return fn(ctx, tx)
 }

core/stores/sqlx/tx_test.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"testing"
@@ -26,26 +27,50 @@ func (mt *mockTx) Exec(q string, args ...interface{}) (sql.Result, error) {
 	return nil, nil
 }
 
+func (mt *mockTx) ExecCtx(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	return nil, nil
+}
+
 func (mt *mockTx) Prepare(query string) (StmtSession, error) {
 	return nil, nil
 }
 
+func (mt *mockTx) PrepareCtx(ctx context.Context, query string) (StmtSession, error) {
+	return nil, nil
+}
+
 func (mt *mockTx) QueryRow(v interface{}, q string, args ...interface{}) error {
 	return nil
 }
 
+func (mt *mockTx) QueryRowCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
 func (mt *mockTx) QueryRowPartial(v interface{}, q string, args ...interface{}) error {
 	return nil
 }
 
+func (mt *mockTx) QueryRowPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
 func (mt *mockTx) QueryRows(v interface{}, q string, args ...interface{}) error {
 	return nil
 }
 
+func (mt *mockTx) QueryRowsCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
 func (mt *mockTx) QueryRowsPartial(v interface{}, q string, args ...interface{}) error {
 	return nil
 }
 
+func (mt *mockTx) QueryRowsPartialCtx(ctx context.Context, v interface{}, query string, args ...interface{}) error {
+	return nil
+}
+
 func (mt *mockTx) Rollback() error {
 	mt.status |= mockRollback
 	return nil
@@ -59,18 +84,20 @@ func beginMock(mock *mockTx) beginnable {
 
 func TestTransactCommit(t *testing.T) {
 	mock := &mockTx{}
-	err := transactOnConn(nil, beginMock(mock), func(Session) error {
+	err := transactOnConn(context.Background(), nil, beginMock(mock),
-		return nil
+		func(context.Context, Session) error {
-	})
+			return nil
+		})
 	assert.Equal(t, mockCommit, mock.status)
 	assert.Nil(t, err)
 }
 
 func TestTransactRollback(t *testing.T) {
 	mock := &mockTx{}
-	err := transactOnConn(nil, beginMock(mock), func(Session) error {
+	err := transactOnConn(context.Background(), nil, beginMock(mock),
-		return errors.New("rollback")
+		func(context.Context, Session) error {
-	})
+			return errors.New("rollback")
+		})
 	assert.Equal(t, mockRollback, mock.status)
 	assert.NotNil(t, err)
 }

core/stores/sqlx/utils.go

@@ -1,6 +1,7 @@
 package sqlx
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"strings"
@@ -109,9 +110,9 @@ func logInstanceError(datasource string, err error) {
 	logx.Errorf("Error on getting sql instance of %s: %v", datasource, err)
 }
 
-func logSqlError(stmt string, err error) {
+func logSqlError(ctx context.Context, stmt string, err error) {
 	if err != nil && err != ErrNotFound {
-		logx.Errorf("stmt: %s, error: %s", stmt, err.Error())
+		logx.WithContext(ctx).Errorf("stmt: %s, error: %s", stmt, err.Error())
 	}
 }
 

go.mod

@@ -17,8 +17,8 @@ require (
 	github.com/prometheus/client_golang v1.11.0
 	github.com/spaolacci/murmur3 v1.1.0
 	github.com/stretchr/testify v1.7.0
-	go.etcd.io/etcd/api/v3 v3.5.1
+	go.etcd.io/etcd/api/v3 v3.5.2
-	go.etcd.io/etcd/client/v3 v3.5.1
+	go.etcd.io/etcd/client/v3 v3.5.2
 	go.opentelemetry.io/otel v1.3.0
 	go.opentelemetry.io/otel/exporters/jaeger v1.3.0
 	go.opentelemetry.io/otel/exporters/zipkin v1.3.0
@@ -44,8 +44,11 @@ require (
 	github.com/go-redis/redis/v8 v8.11.4
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/openzipkin/zipkin-go v0.4.0 // indirect
-	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect
+	go.uber.org/atomic v1.9.0 // indirect
-	golang.org/x/sys v0.0.0-20220209214540-3681064d5158 // indirect
+	go.uber.org/multierr v1.8.0 // indirect
-	google.golang.org/genproto v0.0.0-20220211171837-173942840c17 // indirect
+	go.uber.org/zap v1.21.0 // indirect
+	golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
+	golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9 // indirect
+	google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7 // indirect
 	k8s.io/klog/v2 v2.40.1 // indirect
 )

go.sum

@@ -52,6 +52,7 @@ github.com/alicebob/miniredis/v2 v2.17.0 h1:EwLdrIS50uczw71Jc7iVSxZluTKj5nfSP8n7
 github.com/alicebob/miniredis/v2 v2.17.0/go.mod h1:gquAfGbzn92jvtrSC69+6zZnwSODVXVpYDRaGhWaL6I=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
+github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
@@ -378,10 +379,16 @@ github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da h1:NimzV1aGyq29m5u
 github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da/go.mod h1:E1AXubJBdNmFERAOucpDIxNzeGfLzg0mYh+UfMWdChA=
 go.etcd.io/etcd/api/v3 v3.5.1 h1:v28cktvBq+7vGyJXF8G+rWJmj+1XUmMtqcLnH8hDocM=
 go.etcd.io/etcd/api/v3 v3.5.1/go.mod h1:cbVKeC6lCfl7j/8jBhAK6aIYO9XOjdptoxU/nLQcPvs=
+go.etcd.io/etcd/api/v3 v3.5.2 h1:tXok5yLlKyuQ/SXSjtqHc4uzNaMqZi2XsoSPr/LlJXI=
+go.etcd.io/etcd/api/v3 v3.5.2/go.mod h1:5GB2vv4A4AOn3yk7MftYGHkUfGtDHnEraIjym4dYz5A=
 go.etcd.io/etcd/client/pkg/v3 v3.5.1 h1:XIQcHCFSG53bJETYeRJtIxdLv2EWRGxcfzR8lSnTH4E=
 go.etcd.io/etcd/client/pkg/v3 v3.5.1/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
+go.etcd.io/etcd/client/pkg/v3 v3.5.2 h1:4hzqQ6hIb3blLyQ8usCU4h3NghkqcsohEQ3o3VetYxE=
+go.etcd.io/etcd/client/pkg/v3 v3.5.2/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
 go.etcd.io/etcd/client/v3 v3.5.1 h1:oImGuV5LGKjCqXdjkMHCyWa5OO1gYKCnC/1sgdfj1Uk=
 go.etcd.io/etcd/client/v3 v3.5.1/go.mod h1:OnjH4M8OnAotwaB2l9bVgZzRFKru7/ZMoS46OtKyd3Q=
+go.etcd.io/etcd/client/v3 v3.5.2 h1:WdnejrUtQC4nCxK0/dLTMqKOB+U5TP/2Ya0BJL+1otA=
+go.etcd.io/etcd/client/v3 v3.5.2/go.mod h1:kOOaWFFgHygyT0WlSmL8TJiXmMysO/nNUlEsSsN6W4o=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=
@@ -399,14 +406,21 @@ go.opentelemetry.io/otel/trace v1.3.0/go.mod h1:c/VDhno8888bvQYmbYLqe41/Ldmr/KKu
 go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
+go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
 go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
+go.uber.org/goleak v1.1.11/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
+go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8=
+go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/zap v1.17.0 h1:MTjgFu6ZLKvY6Pvaqk97GlxNBuMpV4Hy/3P6tRGlI2U=
 go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
+go.uber.org/zap v1.21.0 h1:WefMeulhovoZ2sYXz7st6K0sLj7bBhpiFaud4r4zST8=
+go.uber.org/zap v1.21.0/go.mod h1:wjWOCqI0f2ZZrJF/UufIOkiC8ii6tm1iqIsLo76RfJw=
 golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
@@ -485,6 +499,8 @@ golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT
 golang.org/x/net v0.0.0-20210917221730-978cfadd31cf/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd h1:O7DYs+zxREGLKzKoMQrtrEacpb0ZVXA5rIwylE2Xchk=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
+golang.org/x/net v0.0.0-20220225172249-27dd8689420f h1:oA4XRj0qtSt8Yo1Zms0CUlsT3KG69V2UGQWPBxujDmc=
+golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
@@ -548,6 +564,8 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220209214540-3681064d5158 h1:rm+CHSpPEEW2IsXUib1ThaHIjuBVZjxNgSKmBLFfD4c=
 golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9 h1:nhht2DYV/Sn3qOayu8lM+cU1ii9sTLUeBQwQQfUHtrs=
+golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
@@ -648,6 +666,8 @@ google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEY
 google.golang.org/genproto v0.0.0-20210602131652-f16073e35f0c/go.mod h1:UODoCrxHCcBojKKwX1terBiRUaqAsFqJiF615XL43r0=
 google.golang.org/genproto v0.0.0-20220211171837-173942840c17 h1:2X+CNIheCutWRyKRte8szGxrE5ggtV4U+NKAbh/oLhg=
 google.golang.org/genproto v0.0.0-20220211171837-173942840c17/go.mod h1:kGP+zUP2Ddo0ayMi4YuN7C3WZyJvGLZRh8Z5wnAqvEI=
+google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7 h1:ntPPoHzFW6Xp09ueznmahONZufyoSakK/piXnr2BU3I=
+google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7/go.mod h1:kGP+zUP2Ddo0ayMi4YuN7C3WZyJvGLZRh8Z5wnAqvEI=
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
 google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=

readme-cn.md

@@ -13,7 +13,15 @@
 
 > ***缩短从需求到上线的距离***
 
-**注意：为了满足开源基金会要求，go-zero 从好未来（tal-tech）组织下迁移至中立的 GitHub 组织（zeromicro）。**
+**为了满足开源基金会要求，go-zero 从好未来（tal-tech）组织下迁移至中立的 GitHub 组织（zeromicro）。**
+
+> ***注意：***
+>
+> 从 v1.3.0 之前版本升级请执行以下命令：
+> 
+> GOPROXY=https://goproxy.cn/,direct go install github.com/zeromicro/go-zero/tools/goctl@latest
+> 
+> goctl migrate —verbose —version v1.3.0
 
 ## 0. go-zero 介绍
 
@@ -242,6 +250,8 @@ go-zero 已被许多公司用于生产部署，接入场景如在线教育、电
 >54. 一犀科技成都有限公司
 >55. 北京术杰科技有限公司
 >56. 时代脉搏网络科技（云浮市）有限公司
+>57. 店有帮
+>58. 七牛云
 
 如果贵公司也已使用 go-zero，欢迎在 [登记地址](https://github.com/zeromicro/go-zero/issues/602) 登记，仅仅为了推广，不做其它用途。
 

readme.md

@@ -9,9 +9,19 @@ English | [简体中文](readme-cn.md)
 [![Go Report Card](https://goreportcard.com/badge/github.com/zeromicro/go-zero)](https://goreportcard.com/report/github.com/zeromicro/go-zero)
 [![Release](https://img.shields.io/github/v/release/zeromicro/go-zero.svg?style=flat-square)](https://github.com/zeromicro/go-zero)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Discord](https://img.shields.io/discord/794530774463414292?label=chat&logo=discord)](https://discord.gg/4JQvC5A4Fe)
+<a href="https://www.producthunt.com/posts/go-zero?utm_source=badge-featured&utm_medium=badge&utm_souce=badge-go&#0045;zero" target="_blank"><img src="https://api.producthunt.com/widgets/embed-image/v1/featured.svg?post_id=334030&theme=light" alt="go&#0045;zero - A&#0032;web&#0032;&#0038;&#0032;rpc&#0032;framework&#0032;written&#0032;in&#0032;Go&#0046; | Product Hunt" style="width: 250px; height: 54px;" width="250" height="54" /></a>
 
 **Note: To meet the requirements of Open Source Foundation, we moved go-zero from tal-tech to zeromicro (a neutral GitHub organization).**
 
+> ***Important!***
+> 
+> To upgrade from previous versions, run the following commands.
+> 
+> `go install github.com/zeromicro/go-zero/tools/goctl@latest`
+> 
+> `goctl migrate —verbose —version v1.3.0`
+
 ## 0. what is go-zero
 
 go-zero (listed in CNCF Landscape: [https://landscape.cncf.io/?selected=go-zero](https://landscape.cncf.io/?selected=go-zero)) is a web and rpc framework with lots of builtin engineering practices. It’s born to ensure the stability of the busy services with resilience design and has been serving sites with tens of millions of users for years.
@@ -221,7 +231,7 @@ go get -u github.com/zeromicro/go-zero
 
 ## 9. Chat group
 
-Join the chat via https://join.slack.com/t/go-zero/shared_invite/zt-10ruju779-BE4y6lQNB_R21samtyKTgA
+Join the chat via https://discord.gg/4JQvC5A4Fe
 
 ## 10. Cloud Native Landscape
 

rest/engine.go

@@ -14,6 +14,7 @@ import (
 	"github.com/zeromicro/go-zero/rest/handler"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zeromicro/go-zero/rest/internal"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 // use 1000m to represent 100%
@@ -154,6 +155,27 @@ func (ng *engine) getShedder(priority bool) load.Shedder {
 	return ng.shedder
 }
 
+// notFoundHandler returns a middleware that handles 404 not found requests.
+func (ng *engine) notFoundHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		chain := alice.New(
+			handler.TracingHandler(ng.conf.Name, ""),
+			ng.getLogHandler(),
+		)
+
+		var h http.Handler
+		if next != nil {
+			h = chain.Then(next)
+		} else {
+			h = chain.Then(http.NotFoundHandler())
+		}
+
+		cw := response.NewHeaderOnceResponseWriter(w)
+		h.ServeHTTP(cw, r)
+		cw.WriteHeader(http.StatusNotFound)
+	})
+}
+
 func (ng *engine) setTlsConfig(cfg *tls.Config) {
 	ng.tlsConfig = cfg
 }

rest/engine_test.go

@@ -1,13 +1,17 @@
 package rest
 
 import (
+	"context"
 	"errors"
 	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/conf"
+	"github.com/zeromicro/go-zero/core/logx"
 )
 
 func TestNewEngine(t *testing.T) {
@@ -190,6 +194,75 @@ func TestEngine_checkedTimeout(t *testing.T) {
 	}
 }
 
+func TestEngine_notFoundHandler(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	ts := httptest.NewServer(ng.notFoundHandler(nil))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+}
+
+func TestEngine_notFoundHandlerNotNil(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	var called int32
+	ts := httptest.NewServer(ng.notFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&called, 1)
+	})))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusNotFound, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&called))
+}
+
+func TestEngine_notFoundHandlerNotNilWriteHeader(t *testing.T) {
+	logx.Disable()
+
+	ng := newEngine(RestConf{})
+	var called int32
+	ts := httptest.NewServer(ng.notFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&called, 1)
+		w.WriteHeader(http.StatusExpectationFailed)
+	})))
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, err := http.NewRequest("GET", ts.URL+"/bad", nil)
+		assert.Nil(t, err)
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		assert.Equal(t, http.StatusExpectationFailed, res.StatusCode)
+		return res.Body.Close()
+	}(context.Background())
+
+	assert.Nil(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&called))
+}
+
 type mockedRouter struct{}
 
 func (m mockedRouter) ServeHTTP(writer http.ResponseWriter, request *http.Request) {

rest/handler/authhandler.go

@@ -1,15 +1,14 @@
 package handler
 
 import (
-	"bufio"
 	"context"
 	"errors"
-	"net"
 	"net/http"
 	"net/http/httputil"
 
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 	"github.com/zeromicro/go-zero/rest/token"
 )
 
@@ -105,7 +104,7 @@ func detailAuthLog(r *http.Request, reason string) {
 }
 
 func unauthorized(w http.ResponseWriter, r *http.Request, err error, callback UnauthorizedCallback) {
-	writer := newGuardedResponseWriter(w)
+	writer := response.NewHeaderOnceResponseWriter(w)
 
 	if err != nil {
 		detailAuthLog(r, err.Error())
@@ -121,47 +120,3 @@ func unauthorized(w http.ResponseWriter, r *http.Request, err error, callback Un
 	// if user not setting HTTP header, we set header with 401
 	writer.WriteHeader(http.StatusUnauthorized)
 }
-
-type guardedResponseWriter struct {
-	writer      http.ResponseWriter
-	wroteHeader bool
-}
-
-func newGuardedResponseWriter(w http.ResponseWriter) *guardedResponseWriter {
-	return &guardedResponseWriter{
-		writer: w,
-	}
-}
-
-func (grw *guardedResponseWriter) Flush() {
-	if flusher, ok := grw.writer.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
-func (grw *guardedResponseWriter) Header() http.Header {
-	return grw.writer.Header()
-}
-
-// Hijack implements the http.Hijacker interface.
-// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
-func (grw *guardedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := grw.writer.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (grw *guardedResponseWriter) Write(body []byte) (int, error) {
-	return grw.writer.Write(body)
-}
-
-func (grw *guardedResponseWriter) WriteHeader(statusCode int) {
-	if grw.wroteHeader {
-		return
-	}
-
-	grw.wroteHeader = true
-	grw.writer.WriteHeader(statusCode)
-}

rest/handler/authhandler_test.go

@@ -90,26 +90,6 @@ func TestAuthHandler_NilError(t *testing.T) {
 	})
 }
 
-func TestAuthHandler_Flush(t *testing.T) {
-	resp := httptest.NewRecorder()
-	handler := newGuardedResponseWriter(resp)
-	handler.Flush()
-	assert.True(t, resp.Flushed)
-}
-
-func TestAuthHandler_Hijack(t *testing.T) {
-	resp := httptest.NewRecorder()
-	writer := newGuardedResponseWriter(resp)
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = newGuardedResponseWriter(mockedHijackable{resp})
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-}
-
 func buildToken(secretKey string, payloads map[string]interface{}, seconds int64) (string, error) {
 	now := time.Now().Unix()
 	claims := make(jwt.MapClaims)

rest/handler/breakerhandler.go

@@ -9,7 +9,7 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/rest/httpx"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const breakerSeparator = "://"
@@ -28,7 +28,7 @@ func BreakerHandler(method, path string, metrics *stat.Metrics) func(http.Handle
 				return
 			}
 
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				if cw.Code < http.StatusInternalServerError {
 					promise.Accept()

rest/handler/prometheushandler.go

@@ -8,7 +8,7 @@ import (
 	"github.com/zeromicro/go-zero/core/metric"
 	"github.com/zeromicro/go-zero/core/prometheus"
 	"github.com/zeromicro/go-zero/core/timex"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const serverNamespace = "http_server"
@@ -41,7 +41,7 @@ func PrometheusHandler(path string) func(http.Handler) http.Handler {
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			startTime := timex.Now()
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				metricServerReqDur.Observe(int64(timex.Since(startTime)/time.Millisecond), path)
 				metricServerReqCodeTotal.Inc(path, strconv.Itoa(cw.Code))

rest/handler/sheddinghandler.go

@@ -8,7 +8,7 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/rest/httpx"
-	"github.com/zeromicro/go-zero/rest/internal/security"
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const serviceType = "api"
@@ -41,7 +41,7 @@ func SheddingHandler(shedder load.Shedder, metrics *stat.Metrics) func(http.Hand
 				return
 			}
 
-			cw := &security.WithCodeResponseWriter{Writer: w}
+			cw := &response.WithCodeResponseWriter{Writer: w}
 			defer func() {
 				if cw.Code == http.StatusServiceUnavailable {
 					promise.Fail()

rest/handler/tracinghandler.go

@@ -18,12 +18,16 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			ctx := propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+			spanName := path
+			if len(spanName) == 0 {
+				spanName = r.URL.Path
+			}
 			spanCtx, span := tracer.Start(
 				ctx,
-				path,
+				spanName,
 				oteltrace.WithSpanKind(oteltrace.SpanKindServer),
 				oteltrace.WithAttributes(semconv.HTTPServerAttributesFromHTTPRequest(
-					serviceName, path, r)...),
+					serviceName, spanName, r)...),
 			)
 			defer span.End()
 

rest/handler/tracinghandler_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/justinas/alice"
 	"github.com/stretchr/testify/assert"
 	ztrace "github.com/zeromicro/go-zero/core/trace"
 	"go.opentelemetry.io/otel"
@@ -21,28 +22,31 @@ func TestOtelHandler(t *testing.T) {
 		Sampler:  1.0,
 	})
 
-	ts := httptest.NewServer(
+	for _, test := range []string{"", "bar"} {
-		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Run(test, func(t *testing.T) {
-			ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+			h := alice.New(TracingHandler("foo", test)).Then(
-			spanCtx := trace.SpanContextFromContext(ctx)
+				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			assert.Equal(t, true, spanCtx.IsValid())
+					ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
-		}),
+					spanCtx := trace.SpanContextFromContext(ctx)
-	)
+					assert.True(t, spanCtx.IsValid())
-	defer ts.Close()
+				}))
+			ts := httptest.NewServer(h)
+			defer ts.Close()
 
-	client := ts.Client()
+			client := ts.Client()
-	err := func(ctx context.Context) error {
+			err := func(ctx context.Context) error {
-		ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
+				ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
-		defer span.End()
+				defer span.End()
 
-		req, _ := http.NewRequest("GET", ts.URL, nil)
+				req, _ := http.NewRequest("GET", ts.URL, nil)
-		otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
+				otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
 
-		res, err := client.Do(req)
+				res, err := client.Do(req)
-		assert.Equal(t, err, nil)
+				assert.Nil(t, err)
-		_ = res.Body.Close()
+				return res.Body.Close()
-		return nil
+			}(context.Background())
-	}(context.Background())
 
-	assert.Equal(t, err, nil)
+			assert.Nil(t, err)
+		})
+	}
 }

rest/internal/cors/handlers.go

@@ -1,10 +1,9 @@
 package cors
 
 import (
-	"bufio"
-	"errors"
-	"net"
 	"net/http"
+
+	"github.com/zeromicro/go-zero/rest/internal/response"
 )
 
 const (
@@ -30,7 +29,7 @@ const (
 // At most one origin can be specified, other origins are ignored if given, default to be *.
 func NotAllowedHandler(fn func(w http.ResponseWriter), origins ...string) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		gw := &guardedResponseWriter{w: w}
+		gw := response.NewHeaderOnceResponseWriter(w)
 		checkAndSetHeaders(gw, r, origins)
 		if fn != nil {
 			fn(gw)
@@ -62,44 +61,6 @@ func Middleware(fn func(w http.Header), origins ...string) func(http.HandlerFunc
 	}
 }
 
-type guardedResponseWriter struct {
-	w           http.ResponseWriter
-	wroteHeader bool
-}
-
-func (w *guardedResponseWriter) Flush() {
-	if flusher, ok := w.w.(http.Flusher); ok {
-		flusher.Flush()
-	}
-}
-
-func (w *guardedResponseWriter) Header() http.Header {
-	return w.w.Header()
-}
-
-// Hijack implements the http.Hijacker interface.
-// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
-func (w *guardedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	if hijacked, ok := w.w.(http.Hijacker); ok {
-		return hijacked.Hijack()
-	}
-
-	return nil, nil, errors.New("server doesn't support hijacking")
-}
-
-func (w *guardedResponseWriter) Write(bytes []byte) (int, error) {
-	return w.w.Write(bytes)
-}
-
-func (w *guardedResponseWriter) WriteHeader(code int) {
-	if w.wroteHeader {
-		return
-	}
-
-	w.w.WriteHeader(code)
-	w.wroteHeader = true
-}
-
 func checkAndSetHeaders(w http.ResponseWriter, r *http.Request, origins []string) {
 	setVaryHeaders(w, r)
 

rest/internal/cors/handlers_test.go

@@ -1,8 +1,6 @@
 package cors
 
 import (
-	"bufio"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -131,48 +129,3 @@ func TestCorsHandlerWithOrigins(t *testing.T) {
 		}
 	}
 }
-
-func TestGuardedResponseWriter_Flush(t *testing.T) {
-	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
-	handler := NotAllowedHandler(func(w http.ResponseWriter) {
-		w.Header().Set("X-Test", "test")
-		w.WriteHeader(http.StatusServiceUnavailable)
-		_, err := w.Write([]byte("content"))
-		assert.Nil(t, err)
-
-		flusher, ok := w.(http.Flusher)
-		assert.True(t, ok)
-		flusher.Flush()
-	}, "foo.com")
-
-	resp := httptest.NewRecorder()
-	handler.ServeHTTP(resp, req)
-	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
-	assert.Equal(t, "test", resp.Header().Get("X-Test"))
-	assert.Equal(t, "content", resp.Body.String())
-}
-
-func TestGuardedResponseWriter_Hijack(t *testing.T) {
-	resp := httptest.NewRecorder()
-	writer := &guardedResponseWriter{
-		w: resp,
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-
-	writer = &guardedResponseWriter{
-		w: mockedHijackable{resp},
-	}
-	assert.NotPanics(t, func() {
-		writer.Hijack()
-	})
-}
-
-type mockedHijackable struct {
-	*httptest.ResponseRecorder
-}
-
-func (m mockedHijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	return nil, nil, nil
-}

rest/internal/response/headeronceresponsewriter.go

@@ -0,0 +1,57 @@
+package response
+
+import (
+	"bufio"
+	"errors"
+	"net"
+	"net/http"
+)
+
+// HeaderOnceResponseWriter is a http.ResponseWriter implementation
+// that only the first WriterHeader takes effect.
+type HeaderOnceResponseWriter struct {
+	w           http.ResponseWriter
+	wroteHeader bool
+}
+
+// NewHeaderOnceResponseWriter returns a HeaderOnceResponseWriter.
+func NewHeaderOnceResponseWriter(w http.ResponseWriter) http.ResponseWriter {
+	return &HeaderOnceResponseWriter{w: w}
+}
+
+// Flush flushes the response writer.
+func (w *HeaderOnceResponseWriter) Flush() {
+	if flusher, ok := w.w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+// Header returns the http header.
+func (w *HeaderOnceResponseWriter) Header() http.Header {
+	return w.w.Header()
+}
+
+// Hijack implements the http.Hijacker interface.
+// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
+func (w *HeaderOnceResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hijacked, ok := w.w.(http.Hijacker); ok {
+		return hijacked.Hijack()
+	}
+
+	return nil, nil, errors.New("server doesn't support hijacking")
+}
+
+// Write writes bytes into w.
+func (w *HeaderOnceResponseWriter) Write(bytes []byte) (int, error) {
+	return w.w.Write(bytes)
+}
+
+// WriteHeader writes code into w, and not sealing the writer.
+func (w *HeaderOnceResponseWriter) WriteHeader(code int) {
+	if w.wroteHeader {
+		return
+	}
+
+	w.w.WriteHeader(code)
+	w.wroteHeader = true
+}

rest/internal/response/headeronceresponsewriter_test.go

@@ -0,0 +1,58 @@
+package response
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHeaderOnceResponseWriter_Flush(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cw := NewHeaderOnceResponseWriter(w)
+		cw.Header().Set("X-Test", "test")
+		cw.WriteHeader(http.StatusServiceUnavailable)
+		cw.WriteHeader(http.StatusExpectationFailed)
+		_, err := cw.Write([]byte("content"))
+		assert.Nil(t, err)
+
+		flusher, ok := cw.(http.Flusher)
+		assert.True(t, ok)
+		flusher.Flush()
+	})
+
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
+	assert.Equal(t, "test", resp.Header().Get("X-Test"))
+	assert.Equal(t, "content", resp.Body.String())
+}
+
+func TestHeaderOnceResponseWriter_Hijack(t *testing.T) {
+	resp := httptest.NewRecorder()
+	writer := &HeaderOnceResponseWriter{
+		w: resp,
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &HeaderOnceResponseWriter{
+		w: mockedHijackable{resp},
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+}
+
+type mockedHijackable struct {
+	*httptest.ResponseRecorder
+}
+
+func (m mockedHijackable) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, nil
+}

rest/internal/response/withcoderesponsewriter.go

@@ -1,7 +1,8 @@
-package security
+package response
 
 import (
 	"bufio"
+	"errors"
 	"net"
 	"net/http"
 )
@@ -27,7 +28,11 @@ func (w *WithCodeResponseWriter) Header() http.Header {
 // Hijack implements the http.Hijacker interface.
 // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
 func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	return w.Writer.(http.Hijacker).Hijack()
+	if hijacked, ok := w.Writer.(http.Hijacker); ok {
+		return hijacked.Hijack()
+	}
+
+	return nil, nil, errors.New("server doesn't support hijacking")
 }
 
 // Write writes bytes into w.

rest/internal/response/withcoderesponsewriter_test.go

@@ -1,4 +1,4 @@
-package security
+package response
 
 import (
 	"net/http"
@@ -31,3 +31,20 @@ func TestWithCodeResponseWriter(t *testing.T) {
 	assert.Equal(t, "test", resp.Header().Get("X-Test"))
 	assert.Equal(t, "content", resp.Body.String())
 }
+
+func TestWithCodeResponseWriter_Hijack(t *testing.T) {
+	resp := httptest.NewRecorder()
+	writer := &WithCodeResponseWriter{
+		Writer: resp,
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &WithCodeResponseWriter{
+		Writer: mockedHijackable{resp},
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+}

rest/server.go

@@ -49,6 +49,7 @@ func NewServer(c RestConf, opts ...RunOption) (*Server, error) {
 		router: router.NewRouter(),
 	}
 
+	opts = append([]RunOption{WithNotFoundHandler(nil)}, opts...)
 	for _, opt := range opts {
 		opt(server)
 	}
@@ -163,7 +164,8 @@ func WithMiddleware(middleware Middleware, rs ...Route) []Route {
 // WithNotFoundHandler returns a RunOption with not found handler set to given handler.
 func WithNotFoundHandler(handler http.Handler) RunOption {
 	return func(server *Server) {
-		server.router.SetNotFoundHandler(handler)
+		notFoundHandler := server.ngin.notFoundHandler(handler)
+		server.router.SetNotFoundHandler(notFoundHandler)
 	}
 }
 

tools/goctl/api/dartgen/gendata.go

@@ -19,12 +19,12 @@ class {{.Name}}{
 	});
 	factory {{.Name}}.fromJson(Map<String,dynamic> m) {
 		return {{.Name}}({{range .Members}}
-			{{lowCamelCase .Name}}: {{if isDirectType .Type.Name}}m['{{tagGet .Tag "json"}}']{{else if isClassListType .Type.Name}}(m['{{tagGet .Tag "json"}}'] as List<dynamic>).map((i) => {{getCoreType .Type.Name}}.fromJson(i)){{else}}{{.Type.Name}}.fromJson(m['{{tagGet .Tag "json"}}']){{end}},{{end}}
+			{{lowCamelCase .Name}}: {{if isDirectType .Type.Name}}m['{{getPropertyFromMember .}}']{{else if isClassListType .Type.Name}}(m['{{getPropertyFromMember .}}'] as List<dynamic>).map((i) => {{getCoreType .Type.Name}}.fromJson(i)){{else}}{{.Type.Name}}.fromJson(m['{{getPropertyFromMember .}}']){{end}},{{end}}
 		);
 	}
 	Map<String,dynamic> toJson() {
 		return { {{range .Members}}
-			'{{tagGet .Tag "json"}}': {{if isDirectType .Type.Name}}{{lowCamelCase .Name}}{{else if isClassListType .Type.Name}}{{lowCamelCase .Name}}.map((i) => i.toJson()){{else}}{{lowCamelCase .Name}}.toJson(){{end}},{{end}}
+			'{{getPropertyFromMember .}}': {{if isDirectType .Type.Name}}{{lowCamelCase .Name}}{{else if isClassListType .Type.Name}}{{lowCamelCase .Name}}.map((i) => i.toJson()){{else}}{{lowCamelCase .Name}}.toJson(){{end}},{{end}}
 		};
 	}
 }

tools/goctl/api/dartgen/util.go

@@ -34,18 +34,12 @@ func pathToFuncName(path string) string {
 	return util.ToLower(camel[:1]) + camel[1:]
 }
 
-func tagGet(tag, k string) string {
+func getPropertyFromMember(member spec.Member) string {
-	tags, err := spec.Parse(tag)
+	name, err := member.GetPropertyName()
 	if err != nil {
-		panic(k + " not exist")
+		panic(fmt.Sprintf("cannot get property name of %q", member.Name))
 	}
-
+	return name
-	v, err := tags.Get(k)
-	if err != nil {
-		panic(k + " value not exist")
-	}
-
-	return v.Name
 }
 
 func isDirectType(s string) bool {

tools/goctl/api/dartgen/util_test.go

@@ -0,0 +1,39 @@
+package dartgen
+
+import (
+	"testing"
+
+	"github.com/zeromicro/go-zero/tools/goctl/api/spec"
+)
+
+func Test_getPropertyFromMember(t *testing.T) {
+	tests := []struct {
+		name   string
+		member spec.Member
+		want   string
+	}{
+		{
+			name: "json tag should be ok",
+			member: spec.Member{
+				Tag:  "`json:\"foo\"`",
+				Name: "Foo",
+			},
+			want: "foo",
+		},
+		{
+			name: "form tag should be ok",
+			member: spec.Member{
+				Tag:  "`form:\"bar\"`",
+				Name: "Bar",
+			},
+			want: "bar",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := getPropertyFromMember(tt.member); got != tt.want {
+				t.Errorf("getPropertyFromMember() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

tools/goctl/api/dartgen/vars.go

@@ -3,12 +3,12 @@ package dartgen
 import "text/template"
 
 var funcMap = template.FuncMap{
-	"tagGet":          tagGet,
+	"getPropertyFromMember": getPropertyFromMember,
-	"isDirectType":    isDirectType,
+	"isDirectType":          isDirectType,
-	"isClassListType": isClassListType,
+	"isClassListType":       isClassListType,
-	"getCoreType":     getCoreType,
+	"getCoreType":           getCoreType,
-	"pathToFuncName":  pathToFuncName,
+	"pathToFuncName":        pathToFuncName,
-	"lowCamelCase":    lowCamelCase,
+	"lowCamelCase":          lowCamelCase,
 }
 
 const (

tools/goctl/api/docgen/gen.go

@@ -12,7 +12,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
 )
 
-// DocCommand generate markdown doc file
+// DocCommand generate Markdown doc file
 func DocCommand(c *cli.Context) error {
 	dir := c.String("dir")
 	if len(dir) == 0 {
@@ -45,7 +45,7 @@ func DocCommand(c *cli.Context) error {
 	for _, p := range files {
 		api, err := parser.Parse(p)
 		if err != nil {
-			return fmt.Errorf("parse file: %s, err: %s", p, err.Error())
+			return fmt.Errorf("parse file: %s, err: %w", p, err)
 		}
 
 		api.Service = api.Service.JoinPrefix()

tools/goctl/api/gogen/genhandlers.go

@@ -34,7 +34,7 @@ func {{.HandlerName}}(svcCtx *svc.ServiceContext) http.HandlerFunc {
 		}
 
 		{{end}}l := {{.LogicName}}.New{{.LogicType}}(r.Context(), svcCtx)
-		{{if .HasResp}}resp, {{end}}err := l.{{.Call}}({{if .HasRequest}}req{{end}})
+		{{if .HasResp}}resp, {{end}}err := l.{{.Call}}({{if .HasRequest}}&req{{end}})
 		if err != nil {
 			httpx.Error(w, err)
 		} else {

tools/goctl/api/gogen/genlogic.go

@@ -26,8 +26,8 @@ type {{.logic}} struct {
 	svcCtx *svc.ServiceContext
 }
 
-func New{{.logic}}(ctx context.Context, svcCtx *svc.ServiceContext) {{.logic}} {
+func New{{.logic}}(ctx context.Context, svcCtx *svc.ServiceContext) *{{.logic}} {
-	return {{.logic}}{
+	return &{{.logic}}{
 		Logger: logx.WithContext(ctx),
 		ctx:    ctx,
 		svcCtx: svcCtx,
@@ -73,7 +73,7 @@ func genLogicByRoute(dir, rootPkg string, cfg *config.Config, group spec.Group,
 		returnString = "return nil"
 	}
 	if len(route.RequestTypeName()) > 0 {
-		requestString = "req " + requestGoTypeName(route, typesPacket)
+		requestString = "req *" + requestGoTypeName(route, typesPacket)
 	}
 
 	subDir := getLogicFolderPath(group, route)

tools/goctl/api/parser/g4/ast/api.go

@@ -8,8 +8,10 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/api/parser/g4/gen/api"
 )
 
-const prefixKey = "prefix"
+const (
-const groupKey = "group"
+	prefixKey = "prefix"
+	groupKey  = "group"
+)
 
 // Api describes syntax for api
 type Api struct {

tools/goctl/api/parser/g4/ast/apiparser.go

@@ -114,13 +114,16 @@ func (p *Parser) parse(filename, content string) (*Api, error) {
 	apiAstList = append(apiAstList, root)
 	for _, imp := range root.Import {
 		dir := filepath.Dir(p.src)
-		imp := filepath.Join(dir, imp.Value.Text())
+		impPath := strings.ReplaceAll(imp.Value.Text(), "\"", "")
-		data, err := p.readContent(imp)
+		if !filepath.IsAbs(impPath) {
+			impPath = filepath.Join(dir, impPath)
+		}
+		data, err := p.readContent(impPath)
 		if err != nil {
 			return nil, err
 		}
 
-		nestedApi, err := p.invoke(imp, data)
+		nestedApi, err := p.invoke(impPath, data)
 		if err != nil {
 			return nil, err
 		}

tools/goctl/api/parser/g4/gen/api/baseparser.go

@@ -12,7 +12,7 @@ import (
 
 const (
 	versionRegex     = `(?m)"v[1-9][0-9]*"`
-	importValueRegex = `(?m)"(/?[a-zA-Z0-9_#-])+\.api"`
+	importValueRegex = `(?m)"\/?(([a-zA-Z0-9.]+)+(\/?){1})+([a-zA-Z0-9]+)+\.api"`
 	tagRegex         = `(?m)\x60[a-z]+:".+"\x60`
 )
 

tools/goctl/api/parser/g4/gen/api/baseparser_test.go

@@ -9,3 +9,24 @@ import (
 func TestMatch(t *testing.T) {
 	assert.False(t, matchRegex("v1ddd", versionRegex))
 }
+
+func TestImportRegex(t *testing.T) {
+	tests := []struct {
+		value   string
+		matched bool
+	}{
+		{`"bar.api"`, true},
+		{`"foo/bar.api"`, true},
+		{`"/foo/bar.api"`, true},
+		{`"../foo/bar.api"`, true},
+		{`"../../foo/bar.api"`, true},
+
+		{`"bar..api"`, false},
+		{`"//bar.api"`, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.value, func(t *testing.T) {
+			assert.Equal(t, tt.matched, matchRegex(tt.value, importValueRegex))
+		})
+	}
+}

tools/goctl/completion/const.go

@@ -1,7 +1,10 @@
 package completion
 
-const BashCompletionFlag = `generate-goctl-completion`
+const (
-const defaultCompletionFilename = "goctl_autocomplete"
+	BashCompletionFlag        = `generate-goctl-completion`
+	defaultCompletionFilename = "goctl_autocomplete"
+)
+
 const (
 	magic = 1 << iota
 	flagZsh

tools/goctl/env/check.go

@@ -0,0 +1,112 @@
+package env
+
+import (
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protoc"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengo"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengogrpc"
+	"github.com/zeromicro/go-zero/tools/goctl/util/console"
+)
+
+type bin struct {
+	name   string
+	exists bool
+	get    func(cacheDir string) (string, error)
+}
+
+var bins = []bin{
+	{
+		name:   "protoc",
+		exists: protoc.Exists(),
+		get:    protoc.Install,
+	},
+	{
+		name:   "protoc-gen-go",
+		exists: protocgengo.Exists(),
+		get:    protocgengo.Install,
+	},
+	{
+		name:   "protoc-gen-go-grpc",
+		exists: protocgengogrpc.Exists(),
+		get:    protocgengogrpc.Install,
+	},
+}
+
+func Check(ctx *cli.Context) error {
+	install := ctx.Bool("install")
+	force := ctx.Bool("force")
+	return check(install, force)
+}
+
+func check(install, force bool) error {
+	pending := true
+	console.Info("[goctl-env]: preparing to check env")
+	defer func() {
+		if p := recover(); p != nil {
+			console.Error("%+v", p)
+			return
+		}
+		if pending {
+			console.Success("\n[goctl-env]: congratulations! your goctl environment is ready!")
+		} else {
+			console.Error(`
+[goctl-env]: check env finish, some dependencies is not found in PATH, you can execute
+command 'goctl env check --install' or 'goctl env install' to install it, for details, 
+please see 'goctl env check --help' or 'goctl env install --help'`)
+		}
+	}()
+	for _, e := range bins {
+		time.Sleep(200 * time.Millisecond)
+		console.Info("")
+		console.Info("[goctl-env]: looking up %q", e.name)
+		if e.exists {
+			console.Success("[goctl-env]: %q is installed", e.name)
+			continue
+		}
+		console.Warning("[goctl-env]: %q is not found in PATH", e.name)
+		if install {
+			install := func() {
+				console.Info("[goctl-env]: preparing to install %q", e.name)
+				path, err := e.get(env.Get(env.GoctlCache))
+				if err != nil {
+					console.Error("[goctl-env]: an error interrupted the installation: %+v", err)
+					pending = false
+				} else {
+					console.Success("[goctl-env]: %q is already installed in %q", e.name, path)
+				}
+			}
+			if force {
+				install()
+				continue
+			}
+			console.Info("[goctl-env]: do you want to install %q [y: YES, n: No]", e.name)
+			for {
+				var in string
+				fmt.Scanln(&in)
+				var brk bool
+				switch {
+				case strings.EqualFold(in, "y"):
+					install()
+					brk = true
+				case strings.EqualFold(in, "n"):
+					pending = false
+					console.Info("[goctl-env]: %q installation is ignored", e.name)
+					brk = true
+				default:
+					console.Error("[goctl-env]: invalid input, input 'y' for yes, 'n' for no")
+				}
+				if brk {
+					break
+				}
+			}
+		} else {
+			pending = false
+		}
+	}
+	return nil
+}

tools/goctl/env/env.go

@@ -0,0 +1,17 @@
+package env
+
+import (
+	"fmt"
+
+	"github.com/urfave/cli"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
+)
+
+func Action(c *cli.Context) error {
+	write := c.StringSlice("write")
+	if len(write) > 0 {
+		return env.WriteEnv(write)
+	}
+	fmt.Println(env.Print())
+	return nil
+}

tools/goctl/go.mod

@@ -1,6 +1,6 @@
 module github.com/zeromicro/go-zero/tools/goctl
 
-go 1.15
+go 1.16
 
 require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
@@ -12,6 +12,6 @@ require (
 	github.com/stretchr/testify v1.7.0
 	github.com/urfave/cli v1.22.5
 	github.com/zeromicro/antlr v0.0.1
-	github.com/zeromicro/ddl-parser v0.0.0-20210712021150-63520aca7348
+	github.com/zeromicro/ddl-parser v1.0.3
 	github.com/zeromicro/go-zero v1.3.0
 )

tools/goctl/go.sum

@@ -173,7 +173,6 @@ github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QD
 github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
 github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
-github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
 github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
@@ -272,7 +271,6 @@ github.com/munnerz/goautoneg v0.0.0-20120707110453-a547fc61f48d/go.mod h1:+n7T8m
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw=
-github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32 h1:W6apQkHrMkS0Muv8G/TipAy/FJl/rCYT0+EuS8+Z0z4=
 github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uYLpLIr5fm8diHn0JbqRycJi6w0Ms=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
@@ -358,8 +356,8 @@ github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da h1:NimzV1aGyq29m5u
 github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da/go.mod h1:E1AXubJBdNmFERAOucpDIxNzeGfLzg0mYh+UfMWdChA=
 github.com/zeromicro/antlr v0.0.1 h1:CQpIn/dc0pUjgGQ81y98s/NGOm2Hfru2NNio2I9mQgk=
 github.com/zeromicro/antlr v0.0.1/go.mod h1:nfpjEwFR6Q4xGDJMcZnCL9tEfQRgszMwu3rDz2Z+p5M=
-github.com/zeromicro/ddl-parser v0.0.0-20210712021150-63520aca7348 h1:OhxL9tn28gDeJVzreIUiE5oVxZCjL3tBJ0XBNw8p5R8=
+github.com/zeromicro/ddl-parser v1.0.3 h1:hFecpbt0oPQMhHAbqG1tz78MUepHUnOkFJp1dvRBFyc=
-github.com/zeromicro/ddl-parser v0.0.0-20210712021150-63520aca7348/go.mod h1:ISU/8NuPyEpl9pa17Py9TBPetMjtsiHrb9f5XGiYbo8=
+github.com/zeromicro/ddl-parser v1.0.3/go.mod h1:ISU/8NuPyEpl9pa17Py9TBPetMjtsiHrb9f5XGiYbo8=
 github.com/zeromicro/go-zero v1.3.0 h1:Eyn36yBtR043sm4YKmxR6eS3UA/GtZDktQ+UqIJ3Lm0=
 github.com/zeromicro/go-zero v1.3.0/go.mod h1:Hy4o1VFAt32lXaQMbaBhoFeZjA/rJqJ4PTGNdGsURcc=
 go.etcd.io/etcd/api/v3 v3.5.1/go.mod h1:cbVKeC6lCfl7j/8jBhAK6aIYO9XOjdptoxU/nLQcPvs=
@@ -381,7 +379,6 @@ go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqe
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
 go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
-go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/zap v1.17.0/go.mod h1:MXVU+bhUf/A7Xi2HNOnopQOrmycQ5Ih87HtOu4q5SSo=
@@ -417,7 +414,6 @@ golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHl
 golang.org/x/lint v0.0.0-20191125180803-fdd1cda4f05f/go.mod h1:5qLYkcX4OjUUV8bRuDixDT3tpyyb+LUpUlRWLxfhWrs=
 golang.org/x/lint v0.0.0-20200130185559-910be7a94367/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
-golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 h1:VLliZ0d+/avPrXXH+OakdXhpJuEoBZuwh1m2j7U6Iug=
 golang.org/x/lint v0.0.0-20210508222113-6edffad5e616/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
 golang.org/x/mobile v0.0.0-20190312151609-d3739f865fa6/go.mod h1:z+o9i4GpDbdi3rU15maQ/Ox0txvL9dWGYEHz965HBQE=
 golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCcRqshq8CkpyQDoeVncDDYHnLhea+o=
@@ -427,7 +423,6 @@ golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzB
 golang.org/x/mod v0.1.1-0.20191107180719-034126e5016b/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -574,12 +569,10 @@ golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.2/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
-golang.org/x/tools v0.1.5 h1:ouewzE6p+/VEB31YYnTbEJdi8pFqKp4P4n85vwo3DHA=
 golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE=
 google.golang.org/api v0.7.0/go.mod h1:WtwebWUNSVBH/HAw79HIFXZNqEvBhG+Ra+ax0hx3E3M=
@@ -643,7 +636,6 @@ google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpAD
 google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
-google.golang.org/protobuf v1.27.1 h1:SnqbnDw1V7RiZcXPx5MEeqPv2s79L9i7BJUlG/+RurQ=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

tools/goctl/goctl.go

@@ -22,6 +22,7 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/bug"
 	"github.com/zeromicro/go-zero/tools/goctl/completion"
 	"github.com/zeromicro/go-zero/tools/goctl/docker"
+	"github.com/zeromicro/go-zero/tools/goctl/env"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/errorx"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
 	"github.com/zeromicro/go-zero/tools/goctl/kube"
@@ -47,6 +48,34 @@ var commands = []cli.Command{
 		Usage:  "upgrade goctl to latest version",
 		Action: upgrade.Upgrade,
 	},
+	{
+		Name:  "env",
+		Usage: "check or edit goctl environment",
+		Flags: []cli.Flag{
+			cli.StringSliceFlag{
+				Name:  "write, w",
+				Usage: "edit goctl environment",
+			},
+		},
+		Subcommands: []cli.Command{
+			{
+				Name:  "check",
+				Usage: "detect goctl env and dependency tools",
+				Flags: []cli.Flag{
+					cli.BoolFlag{
+						Name:  "install, i",
+						Usage: "install dependencies if not found",
+					},
+					cli.BoolFlag{
+						Name:  "force, f",
+						Usage: "silent installation of non-existent dependencies",
+					},
+				},
+				Action: env.Check,
+			},
+		},
+		Action: env.Action,
+	},
 	{
 		Name:        "migrate",
 		Usage:       "migrate from tal-tech to zeromicro",
@@ -408,6 +437,10 @@ var commands = []cli.Command{
 							"if they are, --remote has higher priority\n\tThe git repo directory must be consistent with the " +
 							"https://github.com/zeromicro/go-zero-template directory structure",
 					},
+					cli.StringFlag{
+						Name:  "serviceAccount",
+						Usage: "the ServiceAccount for the deployment",
+					},
 				},
 				Action: kube.DeploymentCommand,
 			},
@@ -658,7 +691,7 @@ var commands = []cli.Command{
 						Flags: []cli.Flag{
 							cli.StringFlag{
 								Name:  "url",
-								Usage: `the data source of database,like "postgres://root:password@127.0.0.1:54332/database?sslmode=disable"`,
+								Usage: `the data source of database,like "postgres://root:password@127.0.0.1:5432/database?sslmode=disable"`,
 							},
 							cli.StringFlag{
 								Name:  "table, t",
@@ -825,7 +858,7 @@ func main() {
 	app.Version = fmt.Sprintf("%s %s/%s", version.BuildVersion, runtime.GOOS, runtime.GOARCH)
 	app.Commands = commands
 
-	// cli already print error messages
+	// cli already print error messages.
 	if err := app.Run(os.Args); err != nil {
 		fmt.Println(aurora.Red(errorx.Wrap(err).Error()))
 		os.Exit(codeFailure)

tools/goctl/internal/errorx/errorx.go

@@ -2,14 +2,14 @@ package errorx
 
 import (
 	"fmt"
-	"runtime"
 	"strings"
 
-	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/env"
 )
 
-var errorFormat = `goctl: generation error: %+v
+var errorFormat = `goctl error: %+v
-goctl version: %s
+goctl env:
+%s
 %s`
 
 // GoctlError represents a goctl error.
@@ -20,8 +20,7 @@ type GoctlError struct {
 
 func (e *GoctlError) Error() string {
 	detail := wrapMessage(e.message...)
-	v := fmt.Sprintf("%s %s/%s", version.BuildVersion, runtime.GOOS, runtime.GOARCH)
+	return fmt.Sprintf(errorFormat, e.err, env.Print(), detail)
-	return fmt.Sprintf(errorFormat, e.err, v, detail)
 }
 
 // Wrap wraps an error with goctl version and message.

tools/goctl/kube/deployment.go

@@ -17,7 +17,8 @@ spec:
     metadata:
       labels:
         app: {{.Name}}
-    spec:
+    spec:{{if .ServiceAccount}}
+      serviceAccountName: {{.ServiceAccount}}{{end}}
       containers:
       - name: {{.Name}}
         image: {{.Image}}

tools/goctl/kube/job.go

@@ -11,8 +11,9 @@ spec:
   jobTemplate:
     spec:
       template:
-        spec:
+        spec:{{if .ServiceAccount}}
-          containers:
+          serviceAccountName: {{.ServiceAccount}}{{end}}
+	      {{end}}containers:
           - name: {{.Name}}
             image: # todo image url
             resources:

tools/goctl/kube/kube.go

@@ -21,21 +21,22 @@ const (
 
 // Deployment describes the k8s deployment yaml
 type Deployment struct {
-	Name        string
+	Name           string
-	Namespace   string
+	Namespace      string
-	Image       string
+	Image          string
-	Secret      string
+	Secret         string
-	Replicas    int
+	Replicas       int
-	Revisions   int
+	Revisions      int
-	Port        int
+	Port           int
-	NodePort    int
+	NodePort       int
-	UseNodePort bool
+	UseNodePort    bool
-	RequestCpu  int
+	RequestCpu     int
-	RequestMem  int
+	RequestMem     int
-	LimitCpu    int
+	LimitCpu       int
-	LimitMem    int
+	LimitMem       int
-	MinReplicas int
+	MinReplicas    int
-	MaxReplicas int
+	MaxReplicas    int
+	ServiceAccount string
 }
 
 // DeploymentCommand is used to generate the kubernetes deployment yaml files.
@@ -72,21 +73,22 @@ func DeploymentCommand(c *cli.Context) error {
 
 	t := template.Must(template.New("deploymentTemplate").Parse(text))
 	err = t.Execute(out, Deployment{
-		Name:        c.String("name"),
+		Name:           c.String("name"),
-		Namespace:   c.String("namespace"),
+		Namespace:      c.String("namespace"),
-		Image:       c.String("image"),
+		Image:          c.String("image"),
-		Secret:      c.String("secret"),
+		Secret:         c.String("secret"),
-		Replicas:    c.Int("replicas"),
+		Replicas:       c.Int("replicas"),
-		Revisions:   c.Int("revisions"),
+		Revisions:      c.Int("revisions"),
-		Port:        c.Int("port"),
+		Port:           c.Int("port"),
-		NodePort:    nodePort,
+		NodePort:       nodePort,
-		UseNodePort: nodePort > 0,
+		UseNodePort:    nodePort > 0,
-		RequestCpu:  c.Int("requestCpu"),
+		RequestCpu:     c.Int("requestCpu"),
-		RequestMem:  c.Int("requestMem"),
+		RequestMem:     c.Int("requestMem"),
-		LimitCpu:    c.Int("limitCpu"),
+		LimitCpu:       c.Int("limitCpu"),
-		LimitMem:    c.Int("limitMem"),
+		LimitMem:       c.Int("limitMem"),
-		MinReplicas: c.Int("minReplicas"),
+		MinReplicas:    c.Int("minReplicas"),
-		MaxReplicas: c.Int("maxReplicas"),
+		MaxReplicas:    c.Int("maxReplicas"),
+		ServiceAccount: c.String("serviceAccount"),
 	})
 	if err != nil {
 		return err

tools/goctl/migrate/migrate.go

@@ -164,12 +164,12 @@ func writeFile(pkgs []*ast.Package, verbose bool) error {
 			w := bytes.NewBuffer(nil)
 			err := format.Node(w, fset, file)
 			if err != nil {
-				return fmt.Errorf("[rewriteImport] format file %s error: %+v", filename, err)
+				return fmt.Errorf("[rewriteImport] format file %s error: %w", filename, err)
 			}
 
 			err = ioutil.WriteFile(filename, w.Bytes(), os.ModePerm)
 			if err != nil {
-				return fmt.Errorf("[rewriteImport] write file %s error: %+v", filename, err)
+				return fmt.Errorf("[rewriteImport] write file %s error: %w", filename, err)
 			}
 			if verbose {
 				console.Success("[OK] migrated %q successfully", filepath.Base(filename))

tools/goctl/migrate/proxy.go

@@ -9,8 +9,10 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
 )
 
-var defaultProxy = "https://goproxy.cn"
+var (
-var defaultProxies = []string{defaultProxy}
+	defaultProxy   = "https://goproxy.cn"
+	defaultProxies = []string{defaultProxy}
+)
 
 func goProxy() []string {
 	wd, err := os.Getwd()

tools/goctl/model/sql/model/postgresqlmodel.go

@@ -108,6 +108,7 @@ func (m *PostgreSqlModel) getColumns(schema, table string, in []*PostgreColumn)
 	if err != nil {
 		return nil, err
 	}
+
 	var list []*Column
 	for _, e := range in {
 		var dft interface{}
@@ -120,9 +121,15 @@ func (m *PostgreSqlModel) getColumns(schema, table string, in []*PostgreColumn)
 			isNullAble = "NO"
 		}
 
-		extra := "auto_increment"
+		var extra string
-		if e.IdentityIncrement.Int32 != 1 {
+		// when identity is true, the column is auto increment
-			extra = ""
+		if e.IdentityIncrement.Int32 == 1 {
+			extra = "auto_increment"
+		}
+
+		// when type is serial, it's auto_increment. and the default value is tablename_columnname_seq
+		if strings.Contains(e.ColumnDefault.String, table+"_"+e.Field.String+"_seq") {
+			extra = "auto_increment"
 		}
 
 		if len(index[e.Field.String]) > 0 {
@@ -172,6 +179,7 @@ func (m *PostgreSqlModel) getIndex(schema, table string) (map[string][]*DbIndex,
 	if err != nil {
 		return nil, err
 	}
+
 	index := make(map[string][]*DbIndex)
 	for _, e := range indexes {
 		if e.IsPrimary.Bool {
@@ -193,6 +201,7 @@ func (m *PostgreSqlModel) getIndex(schema, table string) (map[string][]*DbIndex,
 			SeqInIndex: int(e.IndexSort.Int32),
 		})
 	}
+
 	return index, nil
 }
 

tools/goctl/model/sql/template/delete.go

@@ -2,21 +2,21 @@ package template
 
 // Delete defines a delete template
 var Delete = `
-func (m *default{{.upperStartCamelObject}}Model) Delete({{.lowerStartCamelPrimaryKey}} {{.dataType}}) error {
+func (m *default{{.upperStartCamelObject}}Model) Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error {
-	{{if .withCache}}{{if .containsIndexCache}}data, err:=m.FindOne({{.lowerStartCamelPrimaryKey}})
+	{{if .withCache}}{{if .containsIndexCache}}data, err:=m.FindOneCtx(ctx, {{.lowerStartCamelPrimaryKey}})
 	if err!=nil{
 		return err
-	}{{end}}
+	}
 
-	{{.keys}}
+{{end}}	{{.keys}}
-    _, err {{if .containsIndexCache}}={{else}}:={{end}} m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) {
+    _, err {{if .containsIndexCache}}={{else}}:={{end}} m.ExecCtx(ctx, func(ctx context.Context, conn sqlx.SqlConn) (result sql.Result, err error) {
 		query := fmt.Sprintf("delete from %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}}", m.table)
-		return conn.Exec(query, {{.lowerStartCamelPrimaryKey}})
+		return conn.ExecCtx(ctx, query, {{.lowerStartCamelPrimaryKey}})
 	}, {{.keyValues}}){{else}}query := fmt.Sprintf("delete from %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}}", m.table)
-		_,err:=m.conn.Exec(query, {{.lowerStartCamelPrimaryKey}}){{end}}
+		_,err:=m.conn.ExecCtx(ctx, query, {{.lowerStartCamelPrimaryKey}}){{end}}
 	return err
 }
 `
 
 // DeleteMethod defines a delete template for interface method
-var DeleteMethod = `Delete({{.lowerStartCamelPrimaryKey}} {{.dataType}}) error`
+var DeleteMethod = `Delete(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) error`

tools/goctl/model/sql/template/find.go

@@ -2,12 +2,12 @@ package template
 
 // FindOne defines find row by id.
 var FindOne = `
-func (m *default{{.upperStartCamelObject}}Model) FindOne({{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error) {
+func (m *default{{.upperStartCamelObject}}Model) FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
-	err := m.QueryRow(&resp, {{.cacheKeyVariable}}, func(conn sqlx.SqlConn, v interface{}) error {
+	err := m.QueryRowCtx(ctx, &resp, {{.cacheKeyVariable}}, func(ctx context.Context, conn sqlx.SqlConn, v interface{}) error {
 		query :=  fmt.Sprintf("select %s from %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}} limit 1", {{.lowerStartCamelObject}}Rows, m.table)
-		return conn.QueryRow(v, query, {{.lowerStartCamelPrimaryKey}})
+		return conn.QueryRowCtx(ctx, v, query, {{.lowerStartCamelPrimaryKey}})
 	})
 	switch err {
 	case nil:
@@ -18,7 +18,7 @@ func (m *default{{.upperStartCamelObject}}Model) FindOne({{.lowerStartCamelPrima
 		return nil, err
 	}{{else}}query := fmt.Sprintf("select %s from %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}} limit 1", {{.lowerStartCamelObject}}Rows, m.table)
 	var resp {{.upperStartCamelObject}}
-	err := m.conn.QueryRow(&resp, query, {{.lowerStartCamelPrimaryKey}})
+	err := m.conn.QueryRowCtx(ctx, &resp, query, {{.lowerStartCamelPrimaryKey}})
 	switch err {
 	case nil:
 		return &resp, nil
@@ -35,9 +35,9 @@ var FindOneByField = `
 func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}({{.in}}) (*{{.upperStartCamelObject}}, error) {
 	{{if .withCache}}{{.cacheKey}}
 	var resp {{.upperStartCamelObject}}
-	err := m.QueryRowIndex(&resp, {{.cacheKeyVariable}}, m.formatPrimary, func(conn sqlx.SqlConn, v interface{}) (i interface{}, e error) {
+	err := m.QueryRowIndexCtx(ctx, &resp, {{.cacheKeyVariable}}, m.formatPrimary, func(ctx context.Context, conn sqlx.SqlConn, v interface{}) (i interface{}, e error) {
 		query := fmt.Sprintf("select %s from %s where {{.originalField}} limit 1", {{.lowerStartCamelObject}}Rows, m.table)
-		if err := conn.QueryRow(&resp, query, {{.lowerStartCamelField}}); err != nil {
+		if err := conn.QueryRowCtx(ctx, &resp, query, {{.lowerStartCamelField}}); err != nil {
 			return nil, err
 		}
 		return resp.{{.upperStartCamelPrimaryKey}}, nil
@@ -52,7 +52,7 @@ func (m *default{{.upperStartCamelObject}}Model) FindOneBy{{.upperField}}({{.in}
 	}
 }{{else}}var resp {{.upperStartCamelObject}}
 	query := fmt.Sprintf("select %s from %s where {{.originalField}} limit 1", {{.lowerStartCamelObject}}Rows, m.table )
-	err := m.conn.QueryRow(&resp, query, {{.lowerStartCamelField}})
+	err := m.conn.QueryRowCtx(ctx, &resp, query, {{.lowerStartCamelField}})
 	switch err {
 	case nil:
 		return &resp, nil
@@ -70,14 +70,14 @@ func (m *default{{.upperStartCamelObject}}Model) formatPrimary(primary interface
 	return fmt.Sprintf("%s%v", {{.primaryKeyLeft}}, primary)
 }
 
-func (m *default{{.upperStartCamelObject}}Model) queryPrimary(conn sqlx.SqlConn, v, primary interface{}) error {
+func (m *default{{.upperStartCamelObject}}Model) queryPrimary(ctx context.Context, conn sqlx.SqlConn, v, primary interface{}) error {
 	query := fmt.Sprintf("select %s from %s where {{.originalPrimaryField}} = {{if .postgreSql}}$1{{else}}?{{end}} limit 1", {{.lowerStartCamelObject}}Rows, m.table )
-	return conn.QueryRow(v, query, primary)
+	return conn.QueryRowCtx(ctx, v, query, primary)
 }
 `
 
 // FindOneMethod defines find row method.
-var FindOneMethod = `FindOne({{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error)`
+var FindOneMethod = `FindOne(ctx context.Context, {{.lowerStartCamelPrimaryKey}} {{.dataType}}) (*{{.upperStartCamelObject}}, error)`
 
 // FindOneByFieldMethod defines find row by field method.
-var FindOneByFieldMethod = `FindOneBy{{.upperField}}({{.in}}) (*{{.upperStartCamelObject}}, error) `
+var FindOneByFieldMethod = `FindOneBy{{.upperField}}(ctx context.Context, {{.in}}) (*{{.upperStartCamelObject}}, error) `

tools/goctl/model/sql/template/import.go

@@ -3,6 +3,7 @@ package template
 var (
 	// Imports defines a import template for model in cache case
 	Imports = `import (
+	"context"
 	"database/sql"
 	"fmt"
 	"strings"
@@ -17,6 +18,7 @@ var (
 `
 	// ImportsNoCache defines a import template for model in normal case
 	ImportsNoCache = `import (
+	"context"
 	"database/sql"
 	"fmt"
 	"strings"

tools/goctl/model/sql/template/insert.go

@@ -2,18 +2,18 @@ package template
 
 // Insert defines a template for insert code in model
 var Insert = `
-func (m *default{{.upperStartCamelObject}}Model) Insert(data *{{.upperStartCamelObject}}) (sql.Result,error) {
+func (m *default{{.upperStartCamelObject}}Model) Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error) {
 	{{if .withCache}}{{if .containsIndexCache}}{{.keys}}
-    ret, err := m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) {
+    ret, err := m.ExecCtx(ctx, func(ctx context.Context, conn sqlx.SqlConn) (result sql.Result, err error) {
 		query := fmt.Sprintf("insert into %s (%s) values ({{.expression}})", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet)
-		return conn.Exec(query, {{.expressionValues}})
+		return conn.ExecCtx(ctx, query, {{.expressionValues}})
 	}, {{.keyValues}}){{else}}query := fmt.Sprintf("insert into %s (%s) values ({{.expression}})", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet)
-    ret,err:=m.ExecNoCache(query, {{.expressionValues}})
+    ret,err:=m.ExecNoCacheCtx(ctx, query, {{.expressionValues}})
 	{{end}}{{else}}query := fmt.Sprintf("insert into %s (%s) values ({{.expression}})", m.table, {{.lowerStartCamelObject}}RowsExpectAutoSet)
-    ret,err:=m.conn.Exec(query, {{.expressionValues}}){{end}}
+    ret,err:=m.conn.ExecCtx(ctx, query, {{.expressionValues}}){{end}}
 	return ret,err
 }
 `
 
 // InsertMethod defines an interface method template for insert code in model
-var InsertMethod = `Insert(data *{{.upperStartCamelObject}}) (sql.Result,error)`
+var InsertMethod = `Insert(ctx context.Context, data *{{.upperStartCamelObject}}) (sql.Result,error)`

tools/goctl/model/sql/template/update.go

@@ -2,16 +2,16 @@ package template
 
 // Update defines a template for generating update codes
 var Update = `
-func (m *default{{.upperStartCamelObject}}Model) Update(data *{{.upperStartCamelObject}}) error {
+func (m *default{{.upperStartCamelObject}}Model) Update(ctx context.Context, data *{{.upperStartCamelObject}}) error {
 	{{if .withCache}}{{.keys}}
-    _, err := m.Exec(func(conn sqlx.SqlConn) (result sql.Result, err error) {
+    _, err := m.ExecCtx(ctx, func(ctx context.Context, conn sqlx.SqlConn) (result sql.Result, err error) {
 		query := fmt.Sprintf("update %s set %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}}", m.table, {{.lowerStartCamelObject}}RowsWithPlaceHolder)
-		return conn.Exec(query, {{.expressionValues}})
+		return conn.ExecCtx(ctx, query, {{.expressionValues}})
 	}, {{.keyValues}}){{else}}query := fmt.Sprintf("update %s set %s where {{.originalPrimaryKey}} = {{if .postgreSql}}$1{{else}}?{{end}}", m.table, {{.lowerStartCamelObject}}RowsWithPlaceHolder)
-    _,err:=m.conn.Exec(query, {{.expressionValues}}){{end}}
+    _,err:=m.conn.ExecCtx(ctx, query, {{.expressionValues}}){{end}}
 	return err
 }
 `
 
 // UpdateMethod defines an interface method template for generating update codes
-var UpdateMethod = `Update(data *{{.upperStartCamelObject}}) error`
+var UpdateMethod = `Update(ctx context.Context, data *{{.upperStartCamelObject}}) error`

tools/goctl/pkg/collection/sortedmap.go

@@ -0,0 +1,210 @@
+package sortedmap
+
+import (
+	"container/list"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/stringx"
+)
+
+var (
+	ErrInvalidKVExpression = errors.New(`invalid key-value expression`)
+	ErrInvalidKVS          = errors.New("the length of kv must be a even number")
+)
+
+type KV []interface{}
+
+type SortedMap struct {
+	kv   *list.List
+	keys map[interface{}]*list.Element
+}
+
+func New() *SortedMap {
+	return &SortedMap{
+		kv:   list.New(),
+		keys: make(map[interface{}]*list.Element),
+	}
+}
+
+func (m *SortedMap) SetExpression(expression string) (key interface{}, value interface{}, err error) {
+	idx := strings.Index(expression, "=")
+	if idx == -1 {
+		return "", "", ErrInvalidKVExpression
+	}
+	key = expression[:idx]
+	if len(expression) == idx {
+		value = ""
+	} else {
+		value = expression[idx+1:]
+	}
+	if keys, ok := key.(string); ok && stringx.ContainsWhiteSpace(keys) {
+		return "", "", ErrInvalidKVExpression
+	}
+	if values, ok := value.(string); ok && stringx.ContainsWhiteSpace(values) {
+		return "", "", ErrInvalidKVExpression
+	}
+	if len(key.(string)) == 0 {
+		return "", "", ErrInvalidKVExpression
+	}
+
+	m.SetKV(key, value)
+	return
+}
+
+func (m *SortedMap) SetKV(key, value interface{}) {
+	e, ok := m.keys[key]
+	if !ok {
+		e = m.kv.PushBack(KV{
+			key, value,
+		})
+	} else {
+		e.Value.(KV)[1] = value
+	}
+	m.keys[key] = e
+}
+
+func (m *SortedMap) Set(kv KV) error {
+	if len(kv) == 0 {
+		return nil
+	}
+	if len(kv)%2 != 0 {
+		return ErrInvalidKVS
+	}
+	for idx := 0; idx < len(kv); idx += 2 {
+		m.SetKV(kv[idx], kv[idx+1])
+	}
+	return nil
+}
+
+func (m *SortedMap) Get(key interface{}) (interface{}, bool) {
+	e, ok := m.keys[key]
+	if !ok {
+		return nil, false
+	}
+	return e.Value.(KV)[1], true
+}
+
+func (m *SortedMap) GetOr(key interface{}, dft interface{}) interface{} {
+	e, ok := m.keys[key]
+	if !ok {
+		return dft
+	}
+	return e.Value.(KV)[1]
+}
+
+func (m *SortedMap) GetString(key interface{}) (string, bool) {
+	value, ok := m.Get(key)
+	if !ok {
+		return "", false
+	}
+	vs, ok := value.(string)
+	return vs, ok
+}
+
+func (m *SortedMap) GetStringOr(key interface{}, dft string) string {
+	value, ok := m.GetString(key)
+	if !ok {
+		return dft
+	}
+	return value
+}
+
+func (m *SortedMap) HasKey(key interface{}) bool {
+	_, ok := m.keys[key]
+	return ok
+}
+
+func (m *SortedMap) HasValue(value interface{}) bool {
+	var contains bool
+	m.RangeIf(func(key, v interface{}) bool {
+		if value == v {
+			contains = true
+			return false
+		}
+		return true
+	})
+	return contains
+}
+
+func (m *SortedMap) Keys() []interface{} {
+	keys := make([]interface{}, 0)
+	next := m.kv.Front()
+	for next != nil {
+		keys = append(keys, next.Value.(KV)[0])
+		next = next.Next()
+	}
+	return keys
+}
+
+func (m *SortedMap) Values() []interface{} {
+	keys := m.Keys()
+	values := make([]interface{}, len(keys))
+	for idx, key := range keys {
+		values[idx] = m.keys[key].Value.(KV)[1]
+	}
+	return values
+}
+
+func (m *SortedMap) Range(iterator func(key, value interface{})) {
+	next := m.kv.Front()
+	for next != nil {
+		value := next.Value.(KV)
+		iterator(value[0], value[1])
+		next = next.Next()
+	}
+}
+
+func (m *SortedMap) RangeIf(iterator func(key, value interface{}) bool) {
+	next := m.kv.Front()
+	for next != nil {
+		value := next.Value.(KV)
+		loop := iterator(value[0], value[1])
+		if !loop {
+			return
+		}
+		next = next.Next()
+	}
+}
+
+func (m *SortedMap) Remove(key interface{}) (value interface{}, ok bool) {
+	v, ok := m.keys[key]
+	if !ok {
+		return nil, false
+	}
+	value = v.Value.(KV)[1]
+	ok = true
+	m.kv.Remove(v)
+	delete(m.keys, key)
+	return
+}
+
+func (m *SortedMap) Insert(sm *SortedMap) {
+	sm.Range(func(key, value interface{}) {
+		m.SetKV(key, value)
+	})
+}
+
+func (m *SortedMap) Copy() *SortedMap {
+	sm := New()
+	m.Range(func(key, value interface{}) {
+		sm.SetKV(key, value)
+	})
+	return sm
+}
+
+func (m *SortedMap) Format() []string {
+	format := make([]string, 0)
+	m.Range(func(key, value interface{}) {
+		format = append(format, fmt.Sprintf("%s=%s", key, value))
+	})
+	return format
+}
+
+func (m *SortedMap) Reset() {
+	m.kv.Init()
+	for key := range m.keys {
+		delete(m.keys, key)
+	}
+}

tools/goctl/pkg/collection/sortedmap_test.go

@@ -0,0 +1,235 @@
+package sortedmap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_SortedMap(t *testing.T) {
+	sm := New()
+	t.Run("SetExpression", func(t *testing.T) {
+		_, _, err := sm.SetExpression("")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo= ")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression(" foo=")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("foo =")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("=")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		_, _, err = sm.SetExpression("=bar")
+		assert.ErrorIs(t, err, ErrInvalidKVExpression)
+		key, value, err := sm.SetExpression("foo=bar")
+		assert.Nil(t, err)
+		assert.Equal(t, "foo", key)
+		assert.Equal(t, "bar", value)
+		key, value, err = sm.SetExpression("foo=")
+		assert.Nil(t, err)
+		assert.Equal(t, value, sm.GetOr(key, ""))
+		sm.Reset()
+	})
+
+	t.Run("SetKV", func(t *testing.T) {
+		sm.SetKV("foo", "bar")
+		assert.Equal(t, "bar", sm.GetOr("foo", ""))
+		sm.SetKV("foo", "bar-changed")
+		assert.Equal(t, "bar-changed", sm.GetOr("foo", ""))
+		sm.Reset()
+	})
+
+	t.Run("Set", func(t *testing.T) {
+		err := sm.Set(KV{})
+		assert.Nil(t, err)
+		err = sm.Set(KV{"foo"})
+		assert.ErrorIs(t, ErrInvalidKVS, err)
+		err = sm.Set(KV{"foo", "bar", "bar", "foo"})
+		assert.Nil(t, err)
+		assert.Equal(t, "bar", sm.GetOr("foo", ""))
+		assert.Equal(t, "foo", sm.GetOr("bar", ""))
+		sm.Reset()
+	})
+
+	t.Run("Get", func(t *testing.T) {
+		_, ok := sm.Get("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.Get("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		sm.Reset()
+	})
+
+	t.Run("GetString", func(t *testing.T) {
+		_, ok := sm.GetString("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.GetString("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		sm.Reset()
+	})
+
+	t.Run("GetStringOr", func(t *testing.T) {
+		value := sm.GetStringOr("foo", "bar")
+		assert.Equal(t, "bar", value)
+		sm.SetKV("foo", "foo")
+		value = sm.GetStringOr("foo", "bar")
+		assert.Equal(t, "foo", value)
+		sm.Reset()
+	})
+
+	t.Run("GetOr", func(t *testing.T) {
+		value := sm.GetOr("foo", "bar")
+		assert.Equal(t, "bar", value)
+		sm.SetKV("foo", "foo")
+		value = sm.GetOr("foo", "bar")
+		assert.Equal(t, "foo", value)
+		sm.Reset()
+	})
+
+	t.Run("HasKey", func(t *testing.T) {
+		ok := sm.HasKey("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "")
+		assert.True(t, sm.HasKey("foo"))
+		sm.Reset()
+	})
+
+	t.Run("HasValue", func(t *testing.T) {
+		assert.False(t, sm.HasValue("bar"))
+		sm.SetKV("foo", "bar")
+		assert.True(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Keys", func(t *testing.T) {
+		keys := sm.Keys()
+		assert.Equal(t, 0, len(keys))
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, "")
+		}
+		keys = sm.Keys()
+		var actual []string
+		for _, key := range keys {
+			actual = append(actual, key.(string))
+		}
+
+		assert.Equal(t, expected, actual)
+		sm.Reset()
+	})
+
+	t.Run("Values", func(t *testing.T) {
+		values := sm.Values()
+		assert.Equal(t, 0, len(values))
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		values = sm.Values()
+		var actual []string
+		for _, value := range values {
+			actual = append(actual, value.(string))
+		}
+
+		assert.Equal(t, expected, actual)
+		sm.Reset()
+	})
+
+	t.Run("Range", func(t *testing.T) {
+		var keys, values []string
+		sm.Range(func(key, value interface{}) {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+		})
+		assert.Len(t, keys, 0)
+		assert.Len(t, values, 0)
+
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		sm.Range(func(key, value interface{}) {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+		})
+		assert.Equal(t, expected, keys)
+		assert.Equal(t, expected, values)
+		sm.Reset()
+	})
+
+	t.Run("RangeIf", func(t *testing.T) {
+		var keys, values []string
+		sm.RangeIf(func(key, value interface{}) bool {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+			return true
+		})
+		assert.Len(t, keys, 0)
+		assert.Len(t, values, 0)
+
+		expected := []string{"foo1", "foo2", "foo3"}
+		for _, key := range expected {
+			sm.SetKV(key, key)
+		}
+		sm.RangeIf(func(key, value interface{}) bool {
+			keys = append(keys, key.(string))
+			values = append(values, value.(string))
+			if key.(string) == "foo1" {
+				return false
+			}
+			return true
+		})
+		assert.Equal(t, []string{"foo1"}, keys)
+		assert.Equal(t, []string{"foo1"}, values)
+		sm.Reset()
+	})
+
+	t.Run("Remove", func(t *testing.T) {
+		_, ok := sm.Remove("foo")
+		assert.False(t, ok)
+		sm.SetKV("foo", "bar")
+		value, ok := sm.Remove("foo")
+		assert.True(t, ok)
+		assert.Equal(t, "bar", value)
+		assert.False(t, sm.HasKey("foo"))
+		assert.False(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Insert", func(t *testing.T) {
+		data := New()
+		data.SetKV("foo", "bar")
+		sm.SetKV("foo1", "bar1")
+		sm.Insert(data)
+		assert.True(t, sm.HasKey("foo"))
+		assert.True(t, sm.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Copy", func(t *testing.T) {
+		sm.SetKV("foo", "bar")
+		data := sm.Copy()
+		assert.True(t, data.HasKey("foo"))
+		assert.True(t, data.HasValue("bar"))
+		sm.SetKV("foo", "bar1")
+		assert.True(t, data.HasKey("foo"))
+		assert.True(t, data.HasValue("bar"))
+		sm.Reset()
+	})
+
+	t.Run("Format", func(t *testing.T) {
+		format := sm.Format()
+		assert.Equal(t, []string{}, format)
+		sm.SetKV("foo1", "bar1")
+		sm.SetKV("foo2", "bar2")
+		sm.SetKV("foo3", "")
+		format = sm.Format()
+		assert.Equal(t, []string{"foo1=bar1", "foo2=bar2", "foo3="}, format)
+		sm.Reset()
+	})
+}

tools/goctl/pkg/downloader/downloader.go

@@ -0,0 +1,23 @@
+package downloader
+
+import (
+	"io"
+	"net/http"
+	"os"
+)
+
+func Download(url string, filename string) error {
+	resp, err := http.Get(url)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	f, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = io.Copy(f, resp.Body)
+	return err
+}

tools/goctl/pkg/env/env.go

@@ -0,0 +1,147 @@
+package env
+
+import (
+	"fmt"
+	"io/ioutil"
+	"log"
+	"path/filepath"
+	"runtime"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
+	sortedmap "github.com/zeromicro/go-zero/tools/goctl/pkg/collection"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protoc"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengo"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengogrpc"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+var goctlEnv *sortedmap.SortedMap
+
+const (
+	GoctlOS                = "GOCTL_OS"
+	GoctlArch              = "GOCTL_ARCH"
+	GoctlHome              = "GOCTL_HOME"
+	GoctlDebug             = "GOCTL_DEBUG"
+	GoctlCache             = "GOCTL_CACHE"
+	GoctlVersion           = "GOCTL_VERSION"
+	ProtocVersion          = "PROTOC_VERSION"
+	ProtocGenGoVersion     = "PROTOC_GEN_GO_VERSION"
+	ProtocGenGoGRPCVersion = "PROTO_GEN_GO_GRPC_VERSION"
+
+	envFileDir = "env"
+)
+
+// init initializes the goctl environment variables, the environment variables of the function are set in order,
+// please do not change the logic order of the code.
+func init() {
+	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
+	if err != nil {
+		log.Fatalln(err)
+	}
+	goctlEnv = sortedmap.New()
+	goctlEnv.SetKV(GoctlOS, runtime.GOOS)
+	goctlEnv.SetKV(GoctlArch, runtime.GOARCH)
+	existsEnv := readEnv(defaultGoctlHome)
+	if existsEnv != nil {
+		goctlHome, ok := existsEnv.GetString(GoctlHome)
+		if ok && len(goctlHome) > 0 {
+			goctlEnv.SetKV(GoctlHome, goctlHome)
+		}
+		if debug := existsEnv.GetOr(GoctlDebug, "").(string); debug != "" {
+			if strings.EqualFold(debug, "true") || strings.EqualFold(debug, "false") {
+				goctlEnv.SetKV(GoctlDebug, debug)
+			}
+		}
+		if value := existsEnv.GetStringOr(GoctlCache, ""); value != "" {
+			goctlEnv.SetKV(GoctlCache, value)
+		}
+	}
+	if !goctlEnv.HasKey(GoctlHome) {
+		goctlEnv.SetKV(GoctlHome, defaultGoctlHome)
+	}
+	if !goctlEnv.HasKey(GoctlDebug) {
+		goctlEnv.SetKV(GoctlDebug, "False")
+	}
+
+	if !goctlEnv.HasKey(GoctlCache) {
+		cacheDir, _ := pathx.GetCacheDir()
+		goctlEnv.SetKV(GoctlCache, cacheDir)
+	}
+
+	goctlEnv.SetKV(GoctlVersion, version.BuildVersion)
+	protocVer, _ := protoc.Version()
+	goctlEnv.SetKV(ProtocVersion, protocVer)
+
+	protocGenGoVer, _ := protocgengo.Version()
+	goctlEnv.SetKV(ProtocGenGoVersion, protocGenGoVer)
+
+	protocGenGoGrpcVer, _ := protocgengogrpc.Version()
+	goctlEnv.SetKV(ProtocGenGoGRPCVersion, protocGenGoGrpcVer)
+}
+
+func Print() string {
+	return strings.Join(goctlEnv.Format(), "\n")
+}
+
+func Get(key string) string {
+	return GetOr(key, "")
+}
+
+func GetOr(key string, def string) string {
+	return goctlEnv.GetStringOr(key, def)
+}
+
+func readEnv(goctlHome string) *sortedmap.SortedMap {
+	envFile := filepath.Join(goctlHome, envFileDir)
+	data, err := ioutil.ReadFile(envFile)
+	if err != nil {
+		return nil
+	}
+	dataStr := string(data)
+	lines := strings.Split(dataStr, "\n")
+	sm := sortedmap.New()
+	for _, line := range lines {
+		_, _, err = sm.SetExpression(line)
+		if err != nil {
+			continue
+		}
+	}
+	return sm
+}
+
+func WriteEnv(kv []string) error {
+	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
+	if err != nil {
+		log.Fatalln(err)
+	}
+	data := sortedmap.New()
+	for _, e := range kv {
+		_, _, err := data.SetExpression(e)
+		if err != nil {
+			return err
+		}
+	}
+	data.RangeIf(func(key, value interface{}) bool {
+		switch key.(string) {
+		case GoctlHome, GoctlCache:
+			path := value.(string)
+			if !pathx.FileExists(path) {
+				err = fmt.Errorf("[writeEnv]: path %q is not exists", path)
+				return false
+			}
+		}
+		if goctlEnv.HasKey(key) {
+			goctlEnv.SetKV(key, value)
+			return true
+		} else {
+			err = fmt.Errorf("[writeEnv]: invalid key: %v", key)
+			return false
+		}
+	})
+	if err != nil {
+		return err
+	}
+	envFile := filepath.Join(defaultGoctlHome, envFileDir)
+	return ioutil.WriteFile(envFile, []byte(strings.Join(goctlEnv.Format(), "\n")), 0o777)
+}

tools/goctl/pkg/goctl/goctl.go

@@ -0,0 +1,41 @@
+package goctl
+
+import (
+	"path/filepath"
+	"runtime"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/util/console"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+	"github.com/zeromicro/go-zero/tools/goctl/vars"
+)
+
+func Install(cacheDir, name string, installFn func(dest string) (string, error)) (string, error) {
+	goBin := golang.GoBin()
+	cacheFile := filepath.Join(cacheDir, name)
+	binFile := filepath.Join(goBin, name)
+
+	goos := runtime.GOOS
+	if goos == vars.OsWindows {
+		cacheFile = cacheFile + ".exe"
+		binFile = binFile + ".exe"
+	}
+	// read cache.
+	err := pathx.Copy(cacheFile, binFile)
+	if err == nil {
+		console.Info("%q installed from cache", name)
+		return binFile, nil
+	}
+
+	binFile, err = installFn(binFile)
+	if err != nil {
+		return "", err
+	}
+
+	// write cache.
+	err = pathx.Copy(binFile, cacheFile)
+	if err != nil {
+		console.Warning("write cache error: %+v", err)
+	}
+	return binFile, nil
+}

tools/goctl/pkg/golang/bin.go

@@ -0,0 +1,27 @@
+package golang
+
+import (
+	"go/build"
+	"os"
+	"path/filepath"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+// GoBin returns a path of GOBIN.
+func GoBin() string {
+	def := build.Default
+	goroot := os.Getenv("GOROOT")
+	bin := filepath.Join(goroot, "bin")
+	if !pathx.FileExists(bin) {
+		gopath := os.Getenv("GOPATH")
+		bin = filepath.Join(gopath, "bin")
+	}
+	if !pathx.FileExists(bin) {
+		bin = os.Getenv("GOBIN")
+	}
+	if !pathx.FileExists(bin) {
+		bin = filepath.Join(def.GOPATH, "bin")
+	}
+	return bin
+}

tools/goctl/pkg/golang/install.go

@@ -0,0 +1,17 @@
+package golang
+
+import (
+	"os"
+	"os/exec"
+)
+
+func Install(git string) error {
+	cmd := exec.Command("go", "install", git)
+	env := os.Environ()
+	env = append(env, "GO111MODULE=on", "GOPROXY=https://goproxy.cn,direct")
+	cmd.Env = env
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Run()
+	return err
+}

tools/goctl/pkg/protoc/protoc.go

@@ -0,0 +1,79 @@
+package protoc
+
+import (
+	"archive/zip"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/downloader"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+	"github.com/zeromicro/go-zero/tools/goctl/util/zipx"
+	"github.com/zeromicro/go-zero/tools/goctl/vars"
+)
+
+var url = map[string]string{
+	"linux_32":   "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-linux-x86_32.zip",
+	"linux_64":   "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-linux-x86_64.zip",
+	"darwin":     "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-osx-x86_64.zip",
+	"windows_32": "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-win32.zip",
+	"windows_64": "https://github.com/protocolbuffers/protobuf/releases/download/v3.19.4/protoc-3.19.4-win64.zip",
+}
+
+const (
+	Name        = "protoc"
+	ZipFileName = Name + ".zip"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		goos := runtime.GOOS
+		tempFile := filepath.Join(os.TempDir(), ZipFileName)
+		bit := 32 << (^uint(0) >> 63)
+		var downloadUrl string
+		switch goos {
+		case vars.OsMac:
+			downloadUrl = url[vars.OsMac]
+		case vars.OsWindows:
+			downloadUrl = url[fmt.Sprintf("%s_%d", vars.OsWindows, bit)]
+		case vars.OsLinux:
+			downloadUrl = url[fmt.Sprintf("%s_%d", vars.OsLinux, bit)]
+		default:
+			return "", fmt.Errorf("unsupport OS: %q", goos)
+		}
+
+		err := downloader.Download(downloadUrl, tempFile)
+		if err != nil {
+			return "", err
+		}
+
+		return dest, zipx.Unpacking(tempFile, filepath.Dir(dest), func(f *zip.File) bool {
+			return filepath.Base(f.Name) == filepath.Base(dest)
+		})
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtoc()
+	return err == nil
+}
+
+func Version() (string, error) {
+	path, err := env.LookUpProtoc()
+	if err != nil {
+		return "", err
+	}
+	version, err := execx.Run(path+" --version", "")
+	if err != nil {
+		return "", err
+	}
+	fields := strings.Fields(version)
+	if len(fields) > 1 {
+		return fields[1], nil
+	}
+	return "", nil
+}

tools/goctl/pkg/protocgengo/protocgengo.go

@@ -0,0 +1,53 @@
+package protocgengo
+
+import (
+	"strings"
+	"time"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+)
+
+const (
+	Name = "protoc-gen-go"
+	url  = "google.golang.org/protobuf/cmd/protoc-gen-go@latest"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		err := golang.Install(url)
+		return dest, err
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtocGenGo()
+	return err == nil
+}
+
+// Version is used to get the version of the protoc-gen-go plugin. For older versions, protoc-gen-go does not support
+// version fetching, so if protoc-gen-go --version is executed, it will cause the process to block, so it is controlled
+// by a timer to prevent the older version process from blocking.
+func Version() (string, error) {
+	path, err := env.LookUpProtocGenGo()
+	if err != nil {
+		return "", err
+	}
+	versionC := make(chan string)
+	go func(c chan string) {
+		version, _ := execx.Run(path+" --version", "")
+		fields := strings.Fields(version)
+		if len(fields) > 1 {
+			c <- fields[1]
+		}
+	}(versionC)
+	t := time.NewTimer(time.Second)
+	select {
+	case <-t.C:
+		return "", nil
+	case version := <-versionC:
+		return version, nil
+	}
+}

tools/goctl/pkg/protocgengogrpc/protocgengogrpc.go

@@ -0,0 +1,44 @@
+package protocgengogrpc
+
+import (
+	"strings"
+
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/goctl"
+	"github.com/zeromicro/go-zero/tools/goctl/pkg/golang"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/env"
+)
+
+const (
+	Name = "protoc-gen-go-grpc"
+	url  = "google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest"
+)
+
+func Install(cacheDir string) (string, error) {
+	return goctl.Install(cacheDir, Name, func(dest string) (string, error) {
+		err := golang.Install(url)
+		return dest, err
+	})
+}
+
+func Exists() bool {
+	_, err := env.LookUpProtocGenGoGrpc()
+	return err == nil
+}
+
+// Version is used to get the version of the protoc-gen-go-grpc plugin.
+func Version() (string, error) {
+	path, err := env.LookUpProtocGenGoGrpc()
+	if err != nil {
+		return "", err
+	}
+	version, err := execx.Run(path+" --version", "")
+	if err != nil {
+		return "", err
+	}
+	fields := strings.Fields(version)
+	if len(fields) > 1 {
+		return fields[1], nil
+	}
+	return "", nil
+}

tools/goctl/rpc/cli/zrpc.go

@@ -80,7 +80,7 @@ func ZRPC(c *cli.Context) error {
 		return err
 	}
 
-	var isGooglePlugin = len(grpcOut) > 0
+	isGooglePlugin := len(grpcOut) > 0
 	// If grpcOut is not empty means that user generates grpc code by
 	// https://google.golang.org/protobuf/cmd/protoc-gen-go and
 	// https://google.golang.org/grpc/cmd/protoc-gen-go-grpc,

tools/goctl/rpc/cli/zrpc_test.go

@@ -23,7 +23,7 @@ func Test_GetSourceProto(t *testing.T) {
 		return
 	}
 
-	var testData = []test{
+	testData := []test{
 		{
 			source:   []string{"a.proto"},
 			expected: filepath.Join(pwd, "a.proto"),
@@ -54,7 +54,7 @@ func Test_GetSourceProto(t *testing.T) {
 }
 
 func Test_RemoveGoctlFlag(t *testing.T) {
-	var testData = []test{
+	testData := []test{
 		{
 			source:   strings.Fields("protoc foo.proto --go_out=. --go_opt=bar --zrpc_out=. --style go-zero --home=foo"),
 			expected: "protoc foo.proto --go_out=. --go_opt=bar",
@@ -87,7 +87,7 @@ func Test_RemoveGoctlFlag(t *testing.T) {
 }
 
 func Test_RemovePluginFlag(t *testing.T) {
-	var testData = []test{
+	testData := []test{
 		{
 			source:   strings.Fields("plugins=grpc:."),
 			expected: ".",

tools/goctl/util/env/env.go

@@ -11,10 +11,11 @@ import (
 )
 
 const (
-	bin            = "bin"
+	bin                = "bin"
-	binGo          = "go"
+	binGo              = "go"
-	binProtoc      = "protoc"
+	binProtoc          = "protoc"
-	binProtocGenGo = "protoc-gen-go"
+	binProtocGenGo     = "protoc-gen-go"
+	binProtocGenGrpcGo = "protoc-gen-go-grpc"
 )
 
 // LookUpGo searches an executable go in the directories
@@ -46,6 +47,14 @@ func LookUpProtocGenGo() (string, error) {
 	return LookPath(xProtocGenGo)
 }
 
+// LookUpProtocGenGoGrpc searches an executable protoc-gen-go-grpc in the directories
+// named by the PATH environment variable.
+func LookUpProtocGenGoGrpc() (string, error) {
+	suffix := getExeSuffix()
+	xProtocGenGoGrpc := binProtocGenGrpcGo + suffix
+	return LookPath(xProtocGenGoGrpc)
+}
+
 // LookPath searches for an executable named file in the
 // directories named by the PATH environment variable,
 // for the os windows, the named file will be spliced with the

tools/goctl/util/pathx/file.go

@@ -3,6 +3,7 @@ package pathx
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
@@ -13,22 +14,23 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
 )
 
-// NL defines a new line
+// NL defines a new line.
 const (
 	NL              = "\n"
 	goctlDir        = ".goctl"
 	gitDir          = ".git"
 	autoCompleteDir = ".auto_complete"
+	cacheDir        = "cache"
 )
 
 var goctlHome string
 
-// RegisterGoctlHome register goctl home path
+// RegisterGoctlHome register goctl home path.
 func RegisterGoctlHome(home string) {
 	goctlHome = home
 }
 
-// CreateIfNotExist creates a file if it is not exists
+// CreateIfNotExist creates a file if it is not exists.
 func CreateIfNotExist(file string) (*os.File, error) {
 	_, err := os.Stat(file)
 	if !os.IsNotExist(err) {
@@ -38,7 +40,7 @@ func CreateIfNotExist(file string) (*os.File, error) {
 	return os.Create(file)
 }
 
-// RemoveIfExist deletes the specified file if it is exists
+// RemoveIfExist deletes the specified file if it is exists.
 func RemoveIfExist(filename string) error {
 	if !FileExists(filename) {
 		return nil
@@ -47,7 +49,7 @@ func RemoveIfExist(filename string) error {
 	return os.Remove(filename)
 }
 
-// RemoveOrQuit deletes the specified file if read a permit command from stdin
+// RemoveOrQuit deletes the specified file if read a permit command from stdin.
 func RemoveOrQuit(filename string) error {
 	if !FileExists(filename) {
 		return nil
@@ -60,23 +62,29 @@ func RemoveOrQuit(filename string) error {
 	return os.Remove(filename)
 }
 
-// FileExists returns true if the specified file is exists
+// FileExists returns true if the specified file is exists.
 func FileExists(file string) bool {
 	_, err := os.Stat(file)
 	return err == nil
 }
 
-// FileNameWithoutExt returns a file name without suffix
+// FileNameWithoutExt returns a file name without suffix.
 func FileNameWithoutExt(file string) string {
 	return strings.TrimSuffix(file, filepath.Ext(file))
 }
 
-// GetGoctlHome returns the path value of the goctl home where Join $HOME with .goctl
+// GetGoctlHome returns the path value of the goctl, the default path is ~/.goctl, if the path has
+// been set by calling the RegisterGoctlHome method, the user-defined path refers to.
 func GetGoctlHome() (string, error) {
 	if len(goctlHome) != 0 {
 		return goctlHome, nil
 	}
 
+	return GetDefaultGoctlHome()
+}
+
+// GetDefaultGoctlHome returns the path value of the goctl home where Join $HOME with .goctl.
+func GetDefaultGoctlHome() (string, error) {
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", err
@@ -104,7 +112,17 @@ func GetAutoCompleteHome() (string, error) {
 	return filepath.Join(goctlH, autoCompleteDir), nil
 }
 
-// GetTemplateDir returns the category path value in GoctlHome where could get it by GetGoctlHome
+// GetCacheDir returns the cache dit of goctl.
+func GetCacheDir() (string, error) {
+	goctlH, err := GetGoctlHome()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(goctlH, cacheDir), nil
+}
+
+// GetTemplateDir returns the category path value in GoctlHome where could get it by GetGoctlHome.
 func GetTemplateDir(category string) (string, error) {
 	home, err := GetGoctlHome()
 	if err != nil {
@@ -112,7 +130,7 @@ func GetTemplateDir(category string) (string, error) {
 	}
 	if home == goctlHome {
 		// backward compatible, it will be removed in the feature
-		// backward compatible start
+		// backward compatible start.
 		beforeTemplateDir := filepath.Join(home, version.GetGoctlVersion(), category)
 		fs, _ := ioutil.ReadDir(beforeTemplateDir)
 		var hasContent bool
@@ -124,7 +142,7 @@ func GetTemplateDir(category string) (string, error) {
 		if hasContent {
 			return beforeTemplateDir, nil
 		}
-		// backward compatible end
+		// backward compatible end.
 
 		return filepath.Join(home, category), nil
 	}
@@ -132,7 +150,7 @@ func GetTemplateDir(category string) (string, error) {
 	return filepath.Join(home, version.GetGoctlVersion(), category), nil
 }
 
-// InitTemplates creates template files GoctlHome where could get it by GetGoctlHome
+// InitTemplates creates template files GoctlHome where could get it by GetGoctlHome.
 func InitTemplates(category string, templates map[string]string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -152,7 +170,7 @@ func InitTemplates(category string, templates map[string]string) error {
 	return nil
 }
 
-// CreateTemplate writes template into file even it is exists
+// CreateTemplate writes template into file even it is exists.
 func CreateTemplate(category, name, content string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -161,7 +179,7 @@ func CreateTemplate(category, name, content string) error {
 	return createTemplate(filepath.Join(dir, name), content, true)
 }
 
-// Clean deletes all templates and removes the parent directory
+// Clean deletes all templates and removes the parent directory.
 func Clean(category string) error {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -170,7 +188,7 @@ func Clean(category string) error {
 	return os.RemoveAll(dir)
 }
 
-// LoadTemplate gets template content by the specified file
+// LoadTemplate gets template content by the specified file.
 func LoadTemplate(category, file, builtin string) (string, error) {
 	dir, err := GetTemplateDir(category)
 	if err != nil {
@@ -223,7 +241,7 @@ func createTemplate(file, content string, force bool) error {
 	return err
 }
 
-// MustTempDir creates a temporary directory
+// MustTempDir creates a temporary directory.
 func MustTempDir() string {
 	dir, err := ioutil.TempDir("", "")
 	if err != nil {
@@ -232,3 +250,25 @@ func MustTempDir() string {
 
 	return dir
 }
+
+func Copy(src, dest string) error {
+	f, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	dir := filepath.Dir(dest)
+	err = MkdirIfNotExist(dir)
+	if err != nil {
+		return err
+	}
+	w, err := os.Create(dest)
+	if err != nil {
+		return err
+	}
+	w.Chmod(os.ModePerm)
+	defer w.Close()
+	_, err = io.Copy(w, f)
+	return err
+}

tools/goctl/util/stringx/string.go

@@ -6,6 +6,8 @@ import (
 	"unicode"
 )
 
+var WhiteSpace = []rune{'\n', '\t', '\f', '\v', ' '}
+
 // String  provides for converting the source text into other spell case,like lower,snake,camel
 type String struct {
 	source string
@@ -114,3 +116,24 @@ func (s String) splitBy(fn func(r rune) bool, remove bool) []string {
 	}
 	return list
 }
+
+func ContainsAny(s string, runes ...rune) bool {
+	if len(runes) == 0 {
+		return true
+	}
+	tmp := make(map[rune]struct{}, len(runes))
+	for _, r := range runes {
+		tmp[r] = struct{}{}
+	}
+
+	for _, r := range s {
+		if _, ok := tmp[r]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+func ContainsWhiteSpace(s string) bool {
+	return ContainsAny(s, WhiteSpace...)
+}

tools/goctl/util/zipx/zipx.go

@@ -0,0 +1,51 @@
+package zipx
+
+import (
+	"archive/zip"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
+)
+
+func Unpacking(name, destPath string, mapper func(f *zip.File) bool) error {
+	r, err := zip.OpenReader(name)
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+
+	for _, file := range r.File {
+		ok := mapper(file)
+		if ok {
+			err = fileCopy(file, destPath)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func fileCopy(file *zip.File, destPath string) error {
+	rc, err := file.Open()
+	if err != nil {
+		return err
+	}
+	defer rc.Close()
+	filename := filepath.Join(destPath, filepath.Base(file.Name))
+	dir := filepath.Dir(filename)
+	err = pathx.MkdirIfNotExist(dir)
+	if err != nil {
+		return err
+	}
+
+	w, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer w.Close()
+	_, err = io.Copy(w, rc)
+	return err
+}