parse tso for snapshot (#124)

pingcap · Jul 18, 2020 · 7cb3b39 · 7cb3b39
1 parent 6d1ba10
commit 7cb3b39
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/v4/export/dump.go b/v4/export/dump.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"database/sql"
 	"errors"
-	"strconv"
 	"strings"
 	"time"
 
@@ -89,7 +88,7 @@ func Dump(pCtx context.Context, conf *Config) (err error) {
 	}
 
 	if doPdGC {
-		snapshotTS, err := strconv.ParseUint(conf.Snapshot, 10, 64)
+		snapshotTS, err := parseSnapshotToTSO(pool, conf.Snapshot)
 		if err != nil {
 			return err
 		}

diff --git a/v4/export/sql.go b/v4/export/sql.go
@@ -632,6 +632,28 @@ found:
 	return uint64(estRows)
 }
 
+func parseSnapshotToTSO(pool *sql.DB, snapshot string) (uint64, error) {
+	snapshotTS, err := strconv.ParseUint(snapshot, 10, 64)
+	if err == nil {
+		return snapshotTS, nil
+	}
+	var tso sql.NullInt64
+	err = simpleQueryWithArgs(pool, func(rows *sql.Rows) error {
+		err := rows.Scan(&tso)
+		if err != nil {
+			return err
+		}
+		if !tso.Valid {
+			return fmt.Errorf("snapshot %s format not supported. please use tso or '2006-01-02 15:04:05' format time", snapshot)
+		}
+		return nil
+	}, "SELECT unix_timestamp(?)", snapshot)
+	if err != nil {
+		return 0, withStack(err)
+	}
+	return (uint64(tso.Int64)<<18)*1000 + 1, nil
+}
+
 func buildWhereCondition(conf *Config, where string) string {
 	var query strings.Builder
 	separator := "WHERE"

diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go
@@ -190,6 +190,32 @@ func (s *testDumpSuite) TestBuildSelectField(c *C) {
 
 }
 
+func (s *testDumpSuite) TestParseSnapshotToTSO(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	snapshot := "2020/07/18 20:31:50"
+	var unixTimeStamp uint64 = 1595075510
+	// generate columns valid snapshot
+	mock.ExpectQuery(`SELECT unix_timestamp(?)`).
+		WithArgs(sqlmock.AnyArg()).
+		WillReturnRows(sqlmock.NewRows([]string{`unix_timestamp("2020/07/18 20:31:50")`}).AddRow(1595075510))
+	tso, err := parseSnapshotToTSO(db, snapshot)
+	c.Assert(err, IsNil)
+	c.Assert(tso, Equals, (unixTimeStamp<<18)*1000+1)
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+
+	// generate columns not valid snapshot
+	mock.ExpectQuery(`SELECT unix_timestamp(?)`).
+		WithArgs(sqlmock.AnyArg()).
+		WillReturnRows(sqlmock.NewRows([]string{`unix_timestamp("XXYYZZ")`}).AddRow(nil))
+	tso, err = parseSnapshotToTSO(db, "XXYYZZ")
+	c.Assert(err, ErrorMatches, "snapshot XXYYZZ format not supported. please use tso or '2006-01-02 15:04:05' format time")
+	c.Assert(tso, Equals, uint64(0))
+	c.Assert(mock.ExpectationsWereMet(), IsNil)
+}
+
 func makeVersion(major, minor, patch int64, preRelease string) *semver.Version {
 	return &semver.Version{
 		Major:      major,