From 999f5c0c593105cf1ceb3b171c7adf160610bc93 Mon Sep 17 00:00:00 2001
From: Henrique Chehad <hchehad@gmail.com>
Date: Fri, 17 Mar 2017 21:25:13 -0300
Subject: [PATCH] connection timeout configuration (#127)

* created connection timeout configuration

* created postgres timeout test

* moved timeout test to conn_test

* handle rows.Scan error

* back timeout tests to postgres_test
---
 adapters/postgres/connection/conn.go |  2 +-
 adapters/postgres/postgres.go        |  6 ++++--
 adapters/postgres/postgres_test.go   | 21 +++++++++++++++++++++
 config/config.go                     |  3 +++
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/adapters/postgres/connection/conn.go b/adapters/postgres/connection/conn.go
index 4ac37bf08..7d0f3a062 100644
--- a/adapters/postgres/connection/conn.go
+++ b/adapters/postgres/connection/conn.go
@@ -22,7 +22,7 @@ func MustGet() *sqlx.DB {
 	if db == nil {
 		cfg = config.Prest{}
 		config.Parse(&cfg)
-		dbURI := fmt.Sprintf("user=%s dbname=%s host=%s port=%v sslmode=disable", cfg.PGUser, cfg.PGDatabase, cfg.PGHost, cfg.PGPort)
+		dbURI := fmt.Sprintf("user=%s dbname=%s host=%s port=%v sslmode=disable connect_timeout=%d", cfg.PGUser, cfg.PGDatabase, cfg.PGHost, cfg.PGPort, cfg.PGConnTimeout)
 		if cfg.PGPass != "" {
 			dbURI += " password=" + cfg.PGPass
 		}
diff --git a/adapters/postgres/postgres.go b/adapters/postgres/postgres.go
index 365901f5d..22afd3eb9 100644
--- a/adapters/postgres/postgres.go
+++ b/adapters/postgres/postgres.go
@@ -258,7 +258,6 @@ func CountByRequest(req *http.Request) (countQuery string, err error) {
 func Query(SQL string, params ...interface{}) (jsonData []byte, err error) {
 	db := connection.MustGet()
 	prepare, err := db.Prepare(SQL)
-
 	if err != nil {
 		return
 	}
@@ -282,7 +281,10 @@ func Query(SQL string, params ...interface{}) (jsonData []byte, err error) {
 		for i := 0; i < count; i++ {
 			valuePtrs[i] = &values[i]
 		}
-		rows.Scan(valuePtrs...)
+		err = rows.Scan(valuePtrs...)
+		if err != nil {
+			return
+		}
 		entry := make(map[string]interface{})
 		for i, col := range columns {
 			var v interface{}
diff --git a/adapters/postgres/postgres_test.go b/adapters/postgres/postgres_test.go
index 847deab6a..7492b114d 100644
--- a/adapters/postgres/postgres_test.go
+++ b/adapters/postgres/postgres_test.go
@@ -791,3 +791,24 @@ func TestColumnsByRequest(t *testing.T) {
 		}
 	}
 }
+
+func TestTimeout(t *testing.T) {
+	_, err := Query("SET statement_timeout TO 10;")
+	if err != nil {
+		t.Errorf("Error setting statement_timeout: %s", err)
+	}
+
+	_, err = Query("SELECT pg_sleep(1000);")
+	if err == nil {
+		t.Errorf("Error should not be nil")
+	}
+
+	if !strings.Contains(err.Error(), "statement timeout") {
+		t.Errorf("Returned different error: %s", err)
+	}
+
+	_, err = Query("SET statement_timeout TO 0;")
+	if err != nil {
+		t.Errorf("Error disabling statement_timeout")
+	}
+}
diff --git a/config/config.go b/config/config.go
index e6083b894..0628ace95 100644
--- a/config/config.go
+++ b/config/config.go
@@ -36,6 +36,7 @@ type Prest struct {
 	PGDatabase      string
 	PGMaxIdleConn   int
 	PGMAxOpenConn   int
+	PGConnTimeout   int
 	JWTKey          string
 	MigrationsPath  string
 	QueriesPath     string
@@ -66,6 +67,7 @@ func viperCfg() {
 	viper.SetDefault("pg.port", 5432)
 	viper.SetDefault("pg.maxidleconn", 10)
 	viper.SetDefault("pg.maxopenconn", 10)
+	viper.SetDefault("pg.conntimeout", 10)
 
 	user, err := user.Current()
 	if err != nil {
@@ -89,6 +91,7 @@ func Parse(cfg *Prest) (err error) {
 	cfg.PGDatabase = viper.GetString("pg.database")
 	cfg.PGMaxIdleConn = viper.GetInt("pg.maxidleconn")
 	cfg.PGMAxOpenConn = viper.GetInt("pg.maxopenconn")
+	cfg.PGConnTimeout = viper.GetInt("pg.conntimeout")
 	cfg.JWTKey = viper.GetString("jwt.key")
 	cfg.MigrationsPath = viper.GetString("migrations")
 	cfg.AccessConf.Restrict = viper.GetBool("access.restrict")