refactoring QueryAble

Author: ffffwh

driver/mysql/applier_gtid_executed.go

@@ -1,6 +1,7 @@
 package mysql
 
 import (
+	"context"
 	gosql "database/sql"
 	"fmt"
 	"strconv"
@@ -323,7 +324,7 @@ func SelectAllGtidExecuted(db sql.QueryAble, jid string, gtidSet *mysql.MysqlGTI
 	query := fmt.Sprintf(`SELECT source_uuid,gtid,gtid_set FROM %v.%v where job_name=?`,
 		g.DtleSchemaName, g.GtidExecutedTableV4)
 
-	rows, err := db.Query(query, jid)
+	rows, err := db.QueryContext(context.TODO(), query, jid)
 	if err != nil {
 		return nil, err
 	}

driver/mysql/base/utils.go

@@ -8,6 +8,7 @@ package base
 
 import (
 	"bytes"
+	"context"
 	gosql "database/sql"
 	"fmt"
 	"github.com/hashicorp/go-hclog"
@@ -120,7 +121,7 @@ func GetSomeSysVars(db usql.QueryAble, logger g.LoggerType) (r struct {
 	NetWriteTimeout     int
 }) {
 	query := `select @@version, @@time_zone, @@lower_case_table_names, @@net_write_timeout`
-	r.Err = db.QueryRow(query).Scan(&r.Version, &r.TimeZome, &r.LowerCaseTableNames, &r.NetWriteTimeout)
+	r.Err = db.QueryRowContext(context.TODO(), query).Scan(&r.Version, &r.TimeZome, &r.LowerCaseTableNames, &r.NetWriteTimeout)
 	if r.Err != nil {
 		return
 	}
@@ -141,7 +142,7 @@ func GetSomeSysVars(db usql.QueryAble, logger g.LoggerType) (r struct {
 func ShowCreateTable(db usql.QueryAble, databaseName, tableName string) (statement string, err error) {
 	var dummy, createTableStatement string
 	query := fmt.Sprintf(`show create table %s.%s`, umconf.EscapeName(databaseName), umconf.EscapeName(tableName))
-	err = db.QueryRow(query).Scan(&dummy, &createTableStatement)
+	err = db.QueryRowContext(context.TODO(), query).Scan(&dummy, &createTableStatement)
 	return createTableStatement, err
 }
 

driver/mysql/dumper.go

@@ -167,7 +167,7 @@ func (d *dumper) getChunkData() (nRows int64, err error) {
 
 	if d.doChecksum != 0 {
 		if d.doChecksum == 2 || (d.doChecksum == 1 && d.Iteration == 0) {
-			row := d.db.QueryRow(fmt.Sprintf("checksum table %v.%v", d.TableSchema, d.TableName))
+			row := d.db.QueryRowContext(d.Ctx, fmt.Sprintf("checksum table %v.%v", d.TableSchema, d.TableName))
 			var table string
 			var cs int64
 			err := row.Scan(&table, &cs)
@@ -181,7 +181,7 @@ func (d *dumper) getChunkData() (nRows int64, err error) {
 
 	// this must be increased after building query
 	d.Iteration += 1
-	rows, err := d.db.Query(query)
+	rows, err := d.db.QueryContext(context.TODO(), query)
 	if err != nil {
 		d.Logger.Error("error at select chunk", "query", query)
 		return 0, errors.Wrap(err, "select chunk")
@@ -249,7 +249,7 @@ func (d *dumper) getChunkData() (nRows int64, err error) {
 			case <-timer.C:
 				d.Logger.Debug("resultsChannel full. waiting and ping conn")
 				var dummy int
-				errPing := d.db.QueryRow("select 1").Scan(&dummy)
+				errPing := d.db.QueryRowContext(d.Ctx, "select 1").Scan(&dummy)
 				if errPing != nil {
 					d.Logger.Debug("ping query row got error.", "err", errPing)
 				}

driver/mysql/extractor.go

@@ -863,7 +863,7 @@ func (e *Extractor) CountTableRows(db sql.QueryAble, table *common.Table) (int64
 	// It only requires select privilege on target table to select its information_schema item.
 	query = fmt.Sprintf(`select table_rows from information_schema.tables where table_schema = ? and table_name = ?`)
 	var rowsEstimate int64
-	err := db.QueryRow(query, table.TableSchema, table.TableName).Scan(&rowsEstimate)
+	err := db.QueryRowContext(e.ctx, query, table.TableSchema, table.TableName).Scan(&rowsEstimate)
 	if err != nil {
 		e.logger.Error("error when getting estimated row number (using information_schema)", "err", err,
 			"schema", table.TableSchema, "table", table.TableName)

driver/mysql/sql/sqlutils.go

@@ -225,7 +225,7 @@ func QueryRowsMap(db QueryAble, query string, on_row func(RowMap) error, args ..
 		}
 	}()
 
-	rows, err := db.Query(query, args...)
+	rows, err := db.QueryContext(context.TODO(), query, args...)
 	defer rows.Close()
 	if err != nil && err != gosql.ErrNoRows {
 		return err
@@ -236,23 +236,21 @@ func QueryRowsMap(db QueryAble, query string, on_row func(RowMap) error, args ..
 
 // from https://github.com/golang/go/issues/14468
 type QueryAble interface {
-	Exec(query string, args ...interface{}) (gosql.Result, error)
-	Prepare(query string) (*gosql.Stmt, error)
-	Query(query string, args ...interface{}) (*gosql.Rows, error)
-	QueryRow(query string, args ...interface{}) *gosql.Row
+	ExecContext(ctx context.Context, query string, args ...interface{}) (gosql.Result, error)
+	QueryContext(ctx context.Context, query string, args ...interface{}) (*gosql.Rows, error)
 	QueryRowContext(ctx context.Context, query string, args ...interface{}) *gosql.Row
 }
 
 func GetServerUUID(db QueryAble) (result string, err error) {
-	err = db.QueryRow(`SELECT @@SERVER_UUID /*dtle*/`).Scan(&result)
+	err = db.QueryRowContext(context.TODO(), `SELECT @@SERVER_UUID /*dtle*/`).Scan(&result)
 	if err != nil {
 		return "", err
 	}
 	return result, nil
 }
 
 func ShowMasterStatus(db QueryAble) *gosql.Row {
-	return db.QueryRow("show master status /*dtle*/")
+	return db.QueryRowContext(context.TODO(), "show master status /*dtle*/")
 }
 
 // queryResultData returns a raw array of rows for a given query, optionally reading and returning column names
@@ -296,7 +294,7 @@ func ShowDatabases(db QueryAble) ([]string, error) {
 	dbs := make([]string, 0)
 
 	// Get table list
-	rows, err := db.Query("SHOW DATABASES")
+	rows, err := db.QueryContext(context.TODO(), "SHOW DATABASES")
 	if err != nil {
 		return dbs, err
 	}
@@ -341,7 +339,7 @@ func ShowTables(db QueryAble, dbName string, showType bool) (tables []*common.Ta
 		query = fmt.Sprintf("SHOW TABLES IN %s", escapedDbName)
 	}
 	g.Logger.Debug("ShowTables", "query", query)
-	rows, err := db.Query(query)
+	rows, err := db.QueryContext(context.TODO(), query)
 	if err != nil {
 		return tables, err
 	}