planner, executor: supports select statement with AS OF

executor/stale_txn_test.go

@@ -15,6 +15,7 @@ package executor_test
 
 import (
 	"fmt"
+	"strings"
 	"time"
 
 	. "github.com/pingcap/check"
@@ -116,6 +117,104 @@ func (s *testStaleTxnSerialSuite) TestExactStalenessTransaction(c *C) {
 	}
 }
 
+func (s *testStaleTxnSerialSuite) TestSelectAsOf(c *C) {
+	c.Assert(failpoint.Enable("github.com/pingcap/tidb/executor/mockStalenessTxnSchemaVer", "return(false)"), IsNil)
+	defer func() {
+		err := failpoint.Disable("github.com/pingcap/tidb/executor/mockStalenessTxnSchemaVer")
+		c.Assert(err, IsNil)
+	}()
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec(`drop table if exists b`)
+	tk.MustExec("create table t (id int primary key);")
+	tk.MustExec("create table b (id int primary key);")
+	defer func() {
+		tk.MustExec(`drop table if exists b`)
+		tk.MustExec(`drop table if exists t`)
+	}()
+
+	testcases := []struct {
+		name             string
+		sql              string
+		expectPhysicalTS int64
+		preSec           int64
+		// IsStaleness is auto cleanup in select stmt.
+		errorStr string
+	}{
+		{
+			name:             "TimestampExactRead",
+			sql:              `select * from t as of timestamp '2020-09-06 00:00:00';`,
+			expectPhysicalTS: 1599321600000,
+		},
+		{
+			name:             "TimestampExactRead",
+			sql:              `select * from t as of timestamp TIMESTAMP('2020-09-06 00:00:00');`,
+			expectPhysicalTS: 1599321600000,
+		},
+		{
+			name:   "TimestampExactRead",
+			sql:    `select * from t as of timestamp NOW() - INTERVAL 20 SECOND;`,
+			preSec: 20,
+		},
+		{
+			name:   "TimestampExactRead",
+			sql:    `select * from t as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND);`,
+			preSec: 20,
+		},
+		{
+			name:   "TimestampExactRead",
+			sql:    `select * from t as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND), b as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND);`,
+			preSec: 20,
+		},
+		{
+			name:     "TimestampExactRead",
+			sql:      `select * from t as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND), b as of timestamp TIMESTAMP('2020-09-06 00:00:00');`,
+			preSec:   20,
+			errorStr: "not set different timestamp",
+		},
+		{
+			name:     "TimestampExactRead",
+			sql:      `select * from t as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND), b;`,
+			preSec:   20,
+			errorStr: "not set different timestamp",
+		},
+		{
+			name:     "TimestampExactRead",
+			sql:      `select * from t, b as of timestamp TIMESTAMP(NOW() - INTERVAL 20 SECOND);`,
+			preSec:   20,
+			errorStr: "not set different timestamp",
+		},
+		{
+			name:   "NomalRead",
+			sql:    `select * from t, b;`,
+			preSec: 0,
+		},
+	}
+
+	tk.MustExec("use test")
+	for _, testcase := range testcases {
+		c.Log(testcase.name)
+		_, err := tk.Exec(testcase.sql)
+		if len(testcase.errorStr) == 0 {
+			c.Assert(err, IsNil, Commentf("sql:%s, error stack %v", testcase.sql, errors.ErrorStack(err)))
+		} else {
+			c.Assert(err, NotNil)
+			c.Assert(strings.Contains(err.Error(), testcase.errorStr), IsTrue)
+			continue
+		}
+		if testcase.expectPhysicalTS > 0 {
+			c.Assert(oracle.ExtractPhysical(tk.Se.GetSessionVars().TxnCtx.StartTS), Equals, testcase.expectPhysicalTS)
+		} else if testcase.preSec > 0 {
+			curSec := time.Now().Unix()
+			startTS := oracle.ExtractPhysical(tk.Se.GetSessionVars().TxnCtx.StartTS)
+			// exact stale txn tolerate 2 seconds deviation for startTS
+			c.Assert(startTS, Greater, (curSec-testcase.preSec-2)*1000)
+			c.Assert(startTS, Less, (curSec-testcase.preSec+2)*1000)
+		}
+	}
+}
+
 func (s *testStaleTxnSerialSuite) TestStaleReadKVRequest(c *C) {
 	c.Assert(failpoint.Enable("github.com/pingcap/tidb/executor/mockStalenessTxnSchemaVer", "return(false)"), IsNil)
 	defer failpoint.Disable("github.com/pingcap/tidb/executor/mockStalenessTxnSchemaVer")

go.mod

@@ -47,7 +47,7 @@ require (
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
 	github.com/pingcap/kvproto v0.0.0-20210507054410-a8152f8a876c
 	github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4
-	github.com/pingcap/parser v0.0.0-20210513020953-ae2c4497c07b
+	github.com/pingcap/parser v0.0.0-20210517060338-92d468f210ff
 	github.com/pingcap/sysutil v0.0.0-20210315073920-cc0985d983a3
 	github.com/pingcap/tidb-tools v4.0.9-0.20201127090955-2707c97b3853+incompatible
 	github.com/pingcap/tipb v0.0.0-20210422074242-57dd881b81b1
@@ -68,7 +68,7 @@ require (
 	go.etcd.io/etcd v0.5.0-alpha.5.0.20200824191128-ae9734ed278b
 	go.uber.org/atomic v1.7.0
 	go.uber.org/automaxprocs v1.2.0
-	go.uber.org/multierr v1.6.0 // indirect
+	go.uber.org/multierr v1.7.0 // indirect
 	go.uber.org/zap v1.16.0
 	golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a // indirect
 	golang.org/x/mod v0.4.2 // indirect
@@ -81,7 +81,6 @@ require (
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 // indirect
 	honnef.co/go/tools v0.1.4 // indirect
 	modernc.org/mathutil v1.2.2 // indirect
 	sourcegraph.com/sourcegraph/appdash v0.0.0-20190731080439-ebfcffb1b5c0

go.sum

@@ -443,8 +443,8 @@ github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIf
 github.com/pingcap/log v0.0.0-20201112100606-8f1e84a3abc8/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4 h1:ERrF0fTuIOnwfGbt71Ji3DKbOEaP189tjym50u8gpC8=
 github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
-github.com/pingcap/parser v0.0.0-20210513020953-ae2c4497c07b h1:eLuDQ6eJCEKCbGwhGrkjzagwev1GJGU2Y2kFkAsBzV0=
-github.com/pingcap/parser v0.0.0-20210513020953-ae2c4497c07b/go.mod h1:xZC8I7bug4GJ5KtHhgAikjTfU4kBv1Sbo3Pf1MZ6lVw=
+github.com/pingcap/parser v0.0.0-20210517060338-92d468f210ff h1:wopnX7X3Z1iq/3OR5yInILPATjy6kvorLj0uXwtGI8M=
+github.com/pingcap/parser v0.0.0-20210517060338-92d468f210ff/go.mod h1:xZC8I7bug4GJ5KtHhgAikjTfU4kBv1Sbo3Pf1MZ6lVw=
 github.com/pingcap/sysutil v0.0.0-20200206130906-2bfa6dc40bcd/go.mod h1:EB/852NMQ+aRKioCpToQ94Wl7fktV+FNnxf3CX/TTXI=
 github.com/pingcap/sysutil v0.0.0-20210315073920-cc0985d983a3 h1:A9KL9R+lWSVPH8IqUuH1QSTRJ5FGoY1bT2IcfPKsWD8=
 github.com/pingcap/sysutil v0.0.0-20210315073920-cc0985d983a3/go.mod h1:tckvA041UWP+NqYzrJ3fMgC/Hw9wnmQ/tUkp/JaHly8=
@@ -541,8 +541,9 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.6.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/swaggo/files v0.0.0-20190704085106-630677cd5c14/go.mod h1:gxQT6pBGRuIGunNf/+tSOB5OHvguWi8Tbt82WOkf35E=
 github.com/swaggo/gin-swagger v1.2.0/go.mod h1:qlH2+W7zXGZkczuL+r2nEBR2JTT+/lX05Nn6vPhc7OI=
 github.com/swaggo/http-swagger v0.0.0-20200308142732-58ac5e232fba/go.mod h1:O1lAbCgAAX/KZ80LM/OXwtWFI/5TvZlwxSg8Cq08PV0=
@@ -628,8 +629,8 @@ go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/
 go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=
 go.uber.org/multierr v1.4.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=
 go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
-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.7.0 h1:zaiO/rmgFjbmCXdSYJWQcdvOCsthmdaHfr3Gm2Kx4Ec=
+go.uber.org/multierr v1.7.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
 go.uber.org/zap v1.8.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
@@ -924,8 +925,8 @@ gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 h1:tQIYjPdBoyREyB9XMu+nnTclpTYkz2zFM+lzLJFO4gQ=
-gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
+gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

planner/core/logical_plan_builder.go

@@ -2421,6 +2421,95 @@ func (b *PlanBuilder) resolveCorrelatedAggregates(ctx context.Context, sel *ast.
 	return correlatedAggMap, nil
 }
 
+type asofResolver struct {
+	ctx      sessionctx.Context
+	tsValues []*types.Datum
+	err      error
+}
+
+func (a *asofResolver) Enter(inNode ast.Node) (ast.Node, bool) {
+	switch n := inNode.(type) {
+	case *ast.AsOfClause:
+		if n.TsExpr != nil {
+			tsRes, err := evalAstExpr(a.ctx, n.TsExpr)
+			if err != nil {
+				a.err = err
+				return n, true
+			}
+			a.tsValues = append(a.tsValues, &tsRes)
+
+		}
+		return n, false
+	// for the cases: select * from a as of timestamp 'xxx', b;
+	case *ast.TableName:
+		if n.AsOf == nil {
+			a.tsValues = append(a.tsValues, nil)
+		}
+		return n, false
+	}
+
+	return inNode, false
+}
+
+func (a *asofResolver) Leave(inNode ast.Node) (ast.Node, bool) {
+	return inNode, true
+}
+
+// TryExtractTSFromAsOf trys to extract the specified timestamp.
+func TryExtractTSFromAsOf(ctx sessionctx.Context, node ast.Node) (*types.Datum, error) {
+	switch x := node.(type) {
+	case *ast.SelectStmt:
+		tblRefs := x.From
+		if tblRefs != nil {
+			resolve := &asofResolver{
+				ctx: ctx,
+			}
+			tblRefs.Accept(resolve)
+			if resolve.err != nil {
+				return nil, resolve.err
+			}
+			tsValues := resolve.tsValues
+			if len(tsValues) == 0 {
+				return nil, nil
+			}
+
+			var res *types.Datum
+			first := tsValues[0]
+			if first != nil {
+				ts, err := first.ConvertTo(ctx.GetSessionVars().StmtCtx, types.NewFieldType(mysql.TypeTimestamp))
+				if err != nil {
+					return nil, err
+				}
+				res = &ts
+			}
+
+			for i := 1; i < len(tsValues); i++ {
+				val := tsValues[i]
+				if (val == nil && first != nil) || (val != nil && first == nil) {
+					return nil, errors.New("can not set different timestamp in one statement")
+				}
+				if first == nil && val == nil {
+					continue
+				}
+
+				ts, err := val.ConvertTo(ctx.GetSessionVars().StmtCtx, types.NewFieldType(mysql.TypeTimestamp))
+				if err != nil {
+					return nil, err
+				}
+				cmp, err := res.CompareDatum(ctx.GetSessionVars().StmtCtx, &ts)
+				if err != nil {
+					return nil, err
+				}
+				if cmp != 0 {
+					return nil, errors.New("can not set different timestamp in one statement")
+				}
+			}
+			return res, nil
+		}
+	}
+	return nil, nil
+}
+
 // gbyResolver resolves group by items from select fields.
 type gbyResolver struct {
 	ctx        sessionctx.Context

planner/optimize.go

@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/store/tikv/oracle"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/hint"
 	"github.com/pingcap/tidb/util/logutil"
@@ -102,7 +103,10 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 			sctx.GetSessionVars().StmtCtx.AppendWarning(err)
 		}
 	}
-
+	err := tryPrepareStaledTS(ctx, sctx, node, is)
+	if err != nil {
+		return nil, nil, err
+	}
 	if _, isolationReadContainTiKV := sessVars.IsolationReadEngines[kv.TiKV]; isolationReadContainTiKV {
 		var fp plannercore.Plan
 		if fpv, ok := sctx.Value(plannercore.PointPlanKey).(plannercore.PointPlanVal); ok {
@@ -340,6 +344,30 @@ func extractSelectAndNormalizeDigest(stmtNode ast.StmtNode, specifiledDB string)
 	return nil, "", "", nil
 }
 
+// tryPrepareStaledTS try to prepare the staled timestamp for stale read.
+func tryPrepareStaledTS(ctx context.Context, sctx sessionctx.Context, node ast.Node, is infoschema.InfoSchema) error {
+	ts, err := plannercore.TryExtractTSFromAsOf(sctx, node)
+	if err != nil {
+		return err
+	}
+	if ts != nil {
+		tsTime, err := ts.GetMysqlTime().GoTime(sctx.GetSessionVars().TimeZone)
+		if err != nil {
+			return err
+		}
+		tso := oracle.ComposeTS(tsTime.Unix()*1000, 0)
+		opt := sessionctx.StalenessTxnOption{}
+		// TODO: remove TimestampBoundReadTimestamp
+		opt.Mode = ast.TimestampBoundReadTimestamp
+		opt.StartTS = tso
+		err = sctx.NewTxnWithStalenessOption(ctx, opt)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func getBindRecord(ctx sessionctx.Context, stmt ast.StmtNode) (*bindinfo.BindRecord, string, error) {
 	// When the domain is initializing, the bind will be nil.
 	if ctx.Value(bindinfo.SessionBindInfoKeyType) == nil {