Merge pull request #91 from hellofresh/feature/allow-cancelation-long…

…-queries

Allow cancelation of long queries

cmd/steal.go

@@ -50,8 +50,8 @@ func NewStealCmd() *cobra.Command {
 	cmd.PersistentFlags().StringVarP(&opts.from, "from", "f", "root:root@tcp(localhost:3306)/klepto", "Database dsn to steal from")
 	cmd.PersistentFlags().StringVarP(&opts.to, "to", "t", "os://stdout/", "Database to output to (default writes to stdOut)")
 	cmd.PersistentFlags().IntVar(&opts.concurrency, "concurrency", 4, "Sets the amount of dumps to be performed concurrently")
-	cmd.PersistentFlags().StringVar(&opts.readOpts.timeout, "read-timeout", "30s", "Sets the timeout for all read operations")
-	cmd.PersistentFlags().StringVar(&opts.writeOpts.timeout, "write-timeout", "30s", "Sets the timeout for all write operations")
+	cmd.PersistentFlags().StringVar(&opts.readOpts.timeout, "read-timeout", "5m", "Sets the timeout for read operations")
+	cmd.PersistentFlags().StringVar(&opts.writeOpts.timeout, "write-timeout", "30s", "Sets the timeout for write operations")
 	cmd.PersistentFlags().StringVar(&opts.readOpts.maxConnLifetime, "read-conn-lifetime", "0", "Sets the maximum amount of time a connection may be reused on the read database")
 	cmd.PersistentFlags().IntVar(&opts.readOpts.maxConns, "read-max-conns", 10, "Sets the maximum number of open connections to the read database")
 	cmd.PersistentFlags().IntVar(&opts.readOpts.maxIdleConns, "read-max-idle-conns", 0, "Sets the maximum number of connections in the idle connection pool for the read database")

features/mysql_test.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path"
 	"testing"
+	"time"
 
 	"github.com/go-sql-driver/mysql"
 	"github.com/hellofresh/klepto/pkg/config"
@@ -19,15 +20,15 @@ import (
 
 type MysqlTestSuite struct {
 	suite.Suite
-
 	rootDSN        string
 	rootConnection *sql.DB
-
-	databases []string
+	databases      []string
+	timeout        time.Duration
 }
 
 func TestMysqlTestSuite(t *testing.T) {
-	suite.Run(t, new(MysqlTestSuite))
+	s := &MysqlTestSuite{timeout: time.Second * 3}
+	suite.Run(t, s)
 }
 
 func (s *MysqlTestSuite) TestExample() {
@@ -36,7 +37,7 @@ func (s *MysqlTestSuite) TestExample() {
 
 	s.loadFixture(readDSN, "mysql_simple.sql")
 
-	rdr, err := reader.Connect(reader.ConnOpts{DSN: readDSN})
+	rdr, err := reader.Connect(reader.ConnOpts{DSN: readDSN, Timeout: s.timeout})
 	s.Require().NoError(err, "Unable to create reader")
 	defer rdr.Close()
 

features/postgres_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"path"
 	"testing"
+	"time"
 
 	"strconv"
 
@@ -21,11 +22,10 @@ import (
 
 type PostgresTestSuite struct {
 	suite.Suite
-
 	rootDSN        string
 	rootConnection *sql.DB
-
-	databases []string
+	databases      []string
+	timeout        time.Duration
 }
 
 type tableInfo struct {
@@ -35,7 +35,8 @@ type tableInfo struct {
 }
 
 func TestPostgresTestSuite(t *testing.T) {
-	suite.Run(t, new(PostgresTestSuite))
+	s := &PostgresTestSuite{timeout: time.Second * 3}
+	suite.Run(t, s)
 }
 
 func (s *PostgresTestSuite) TestExample() {
@@ -44,7 +45,7 @@ func (s *PostgresTestSuite) TestExample() {
 
 	s.loadFixture(readDSN, "pg_simple.sql")
 
-	rdr, err := reader.Connect(reader.ConnOpts{DSN: readDSN})
+	rdr, err := reader.Connect(reader.ConnOpts{DSN: readDSN, Timeout: s.timeout})
 	s.Require().NoError(err, "Unable to create reader")
 	defer rdr.Close()
 

pkg/reader/generic/sql.go

@@ -1,9 +1,11 @@
 package generic
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"sync"
+	"time"
 
 	sq "github.com/Masterminds/squirrel"
 	"github.com/hellofresh/klepto/pkg/database"
@@ -20,6 +22,8 @@ type (
 		tables []string
 		// columns is a cache variable for tables and there columns in the db
 		columns sync.Map
+		// timeout is the sql read operation timeout
+		timeout time.Duration
 	}
 
 	SqlEngine interface {
@@ -44,8 +48,8 @@ type (
 )
 
 // NewSqlReader creates a new sql reader
-func NewSqlReader(se SqlEngine) *SqlReader {
-	return &SqlReader{SqlEngine: se}
+func NewSqlReader(se SqlEngine, t time.Duration) *SqlReader {
+	return &SqlReader{SqlEngine: se, timeout: t}
 }
 
 // GetTables gets a list of all tables in the database
@@ -93,20 +97,40 @@ func (s *SqlReader) ReadTable(tableName string, rowChan chan<- database.Row, opt
 		opts.Columns = s.formatColumns(tableName, columns)
 	}
 
-	query, err := s.buildQuery(tableName, opts)
+	var (
+		query sq.SelectBuilder
+		err   error
+	)
+	query, err = s.buildQuery(tableName, opts)
 	if err != nil {
 		return errors.Wrapf(err, "failed to build query for %s", tableName)
 	}
 
-	rows, err := query.RunWith(s.GetConnection()).Query()
-	if err != nil {
-		querySQL, queryParams, _ := query.ToSql()
-		logger.WithFields(log.Fields{
-			"query":  querySQL,
-			"params": queryParams,
-		}).Warn("failed to query rows")
-
-		return errors.Wrap(err, "failed to query rows")
+	var rows *sql.Rows
+	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
+	defer cancel()
+
+	errchan := make(chan error)
+	go func() {
+		defer close(errchan)
+		rows, err = query.RunWith(s.GetConnection()).QueryContext(ctx)
+		errchan <- err
+	}()
+
+	select {
+	case <-ctx.Done():
+		return errors.Wrapf(ctx.Err(), fmt.Sprintf("timeout during read %s table", tableName))
+	case err := <-errchan:
+		if err != nil {
+			querySQL, queryParams, _ := query.ToSql()
+			logger.WithError(err).
+				WithFields(log.Fields{
+					"query":  querySQL,
+					"params": queryParams,
+				}).Warn("failed to query rows")
+			return errors.Wrap(err, "failed to query rows")
+		}
+		break
 	}
 
 	return s.publishRows(rows, rowChan, tableName)
@@ -179,7 +203,7 @@ func (s *SqlReader) publishRows(rows *sql.Rows, rowChan chan<- database.Row, tab
 		}
 
 		if err := rows.Scan(fieldPointers...); err != nil {
-			log.WithError(err).Warn("failed to fetch row")
+			log.WithError(err).WithField("table", tableName).Warn("failed to fetch row")
 			continue
 		}
 

pkg/reader/mysql/mysql.go

@@ -28,7 +28,7 @@ func (m *driver) NewConnection(opts reader.ConnOpts) (reader.Reader, error) {
 	conn.SetMaxIdleConns(opts.MaxIdleConns)
 	conn.SetConnMaxLifetime(opts.MaxConnLifetime)
 
-	return NewStorage(conn), nil
+	return NewStorage(conn, opts.Timeout), nil
 }
 
 func init() {

pkg/reader/mysql/reader.go

@@ -17,8 +17,8 @@ type storage struct {
 }
 
 // NewStorage ...
-func NewStorage(conn *sql.DB) reader.Reader {
-	return generic.NewSqlReader(&storage{conn})
+func NewStorage(conn *sql.DB, timeout time.Duration) reader.Reader {
+	return generic.NewSqlReader(&storage{conn}, timeout)
 }
 
 // GetConnection return the connection

pkg/reader/postgres/postgres.go

@@ -29,7 +29,7 @@ func (m *driver) NewConnection(opts reader.ConnOpts) (reader.Reader, error) {
 		return nil, err
 	}
 
-	return NewStorage(conn, dump), nil
+	return NewStorage(conn, dump, opts.Timeout), nil
 }
 
 func init() {

pkg/reader/postgres/reader.go

@@ -3,6 +3,7 @@ package postgres
 import (
 	"database/sql"
 	"strconv"
+	"time"
 
 	"github.com/hellofresh/klepto/pkg/reader"
 	"github.com/hellofresh/klepto/pkg/reader/generic"
@@ -17,13 +18,12 @@ type storage struct {
 }
 
 // NewStorage ...
-func NewStorage(conn *sql.DB, dumper PgDump) reader.Reader {
-	return generic.NewSqlReader(
-		&storage{
-			PgDump:     dumper,
-			connection: conn,
-		},
-	)
+func NewStorage(conn *sql.DB, dumper PgDump, timeout time.Duration) reader.Reader {
+	s := &storage{
+		PgDump:     dumper,
+		connection: conn,
+	}
+	return generic.NewSqlReader(s, timeout)
 }
 
 func (s *storage) GetConnection() *sql.DB {