allow cancelation of long queries

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)
+	defer close(errchan)
+	go func() {
+		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)

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, t time.Duration) reader.Reader {
+	return generic.NewSqlReader(&storage{conn}, t)
 }
 
 // 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, t time.Duration) reader.Reader {
+	s := &storage{
+		PgDump:     dumper,
+		connection: conn,
+	}
+	return generic.NewSqlReader(s, t)
 }
 
 func (s *storage) GetConnection() *sql.DB {