feat: Adding method for simple exec or create query (#31)

go.mod

@@ -3,6 +3,7 @@ module github.com/ViBiOh/httputils/v3
 go 1.14
 
 require (
+	github.com/DATA-DOG/go-sqlmock v1.4.1
 	github.com/lib/pq v1.4.0
 	github.com/prometheus/client_golang v1.5.1
 	github.com/tdewolff/minify/v2 v2.7.4

go.sum

@@ -1,3 +1,5 @@
+github.com/DATA-DOG/go-sqlmock v1.4.1 h1:ThlnYciV1iM/V0OSF/dtkqWb6xo5qITT1TJBG1MRDJM=
+github.com/DATA-DOG/go-sqlmock v1.4.1/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
@@ -37,8 +39,6 @@ github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFB
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/lib/pq v1.3.0 h1:/qkRGz8zljWiDcFvgpwUpwIAPu3r07TDvs3Rws+o/pU=
-github.com/lib/pq v1.3.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.4.0 h1:TmtCFbH+Aw0AixwyttznSMQDgbR5Yed/Gg6S8Funrhc=
 github.com/lib/pq v1.4.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2/go.mod h1:0KeJpeMD6o+O4hW7qJOT7vyQPKrWmj26uf5wMc/IiIs=
@@ -77,8 +77,6 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
 github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
-github.com/tdewolff/minify/v2 v2.7.3 h1:ngzhF7SaunCtbsBjgm7WJzl9HdiKlA1gYC/Qyx9CVMo=
-github.com/tdewolff/minify/v2 v2.7.3/go.mod h1:BkDSm8aMMT0ALGmpt7j3Ra7nLUgZL0qhyrAHXwxcy5w=
 github.com/tdewolff/minify/v2 v2.7.4 h1:r0OZQ3QzWeDS5cXq53Bk4IFIBDZ7fiXIkw1a4bHONsw=
 github.com/tdewolff/minify/v2 v2.7.4/go.mod h1:BkDSm8aMMT0ALGmpt7j3Ra7nLUgZL0qhyrAHXwxcy5w=
 github.com/tdewolff/parse/v2 v2.4.2 h1:Bu2Qv6wepkc+Ou7iB/qHjAhEImlAP5vedzlQRUdj3BI=

pkg/db/db.go

@@ -1,11 +1,13 @@
 package db
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"flag"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/ViBiOh/httputils/v3/pkg/flags"
 	_ "github.com/lib/pq" // Not referenced but needed for database/sql
@@ -101,3 +103,43 @@ func RowsClose(rows *sql.Rows, err error) error {
 
 	return err
 }
+
+// CreateWithTimeout execute query with a RETURNING id
+func CreateWithTimeout(db *sql.DB, tx *sql.Tx, sqlTimeout time.Duration, query string, args ...interface{}) (newID uint64, err error) {
+	var usedTx *sql.Tx
+	if usedTx, err = GetTx(db, tx); err != nil {
+		return
+	}
+
+	if usedTx != tx {
+		defer func() {
+			err = EndTx(usedTx, err)
+		}()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), sqlTimeout)
+	defer cancel()
+
+	err = usedTx.QueryRowContext(ctx, query, args...).Scan(&newID)
+	return
+}
+
+// ExecWithTimeout execute query with specified timeout, disregarding result
+func ExecWithTimeout(db *sql.DB, tx *sql.Tx, sqlTimeout time.Duration, query string, args ...interface{}) (err error) {
+	var usedTx *sql.Tx
+	if usedTx, err = GetTx(db, tx); err != nil {
+		return
+	}
+
+	if usedTx != tx {
+		defer func() {
+			err = EndTx(usedTx, err)
+		}()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), sqlTimeout)
+	defer cancel()
+
+	_, err = usedTx.ExecContext(ctx, query, args...)
+	return
+}

pkg/db/db_test.go

@@ -1,9 +1,14 @@
 package db
 
 import (
+	"database/sql"
+	"errors"
 	"flag"
 	"strings"
 	"testing"
+	"time"
+
+	"github.com/DATA-DOG/go-sqlmock"
 )
 
 func TestFlags(t *testing.T) {
@@ -34,3 +39,187 @@ func TestFlags(t *testing.T) {
 		})
 	}
 }
+
+func TestCreateWithTimeout(t *testing.T) {
+	type args struct {
+		tx bool
+	}
+
+	var cases = []struct {
+		intention string
+		args      args
+		want      uint64
+		wantErr   error
+	}{
+		{
+			"simple",
+			args{
+				tx: false,
+			},
+			1,
+			nil,
+		},
+		{
+			"timeout",
+			args{
+				tx: false,
+			},
+			0,
+			sqlmock.ErrCancelled,
+		},
+		{
+			"with tx",
+			args{
+				tx: true,
+			},
+			1,
+			nil,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.intention, func(t *testing.T) {
+			db, mock, err := sqlmock.New()
+			if err != nil {
+				t.Fatalf("unable to create mock database: %s", err)
+			}
+			defer db.Close()
+
+			var tx *sql.Tx
+			if tc.args.tx {
+				mock.ExpectBegin()
+				dbTx, err := db.Begin()
+
+				if err != nil {
+					t.Errorf("unable to create tx: %v", err)
+				}
+				tx = dbTx
+			}
+
+			if !tc.args.tx {
+				mock.ExpectBegin()
+			}
+
+			expectedQuery := mock.ExpectQuery("INSERT INTO item VALUES").WithArgs(1).WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+
+			if !tc.args.tx {
+				if tc.wantErr != nil {
+					mock.ExpectRollback()
+				} else {
+					mock.ExpectCommit()
+				}
+			}
+
+			if tc.intention == "timeout" {
+				expectedQuery.WillDelayFor(time.Second * 2)
+			}
+
+			got, gotErr := CreateWithTimeout(db, tx, time.Second, "INSERT INTO item VALUES ($1)", 1)
+
+			failed := false
+
+			if tc.wantErr != nil && !errors.Is(gotErr, tc.wantErr) {
+				failed = true
+			} else if got != tc.want {
+				failed = true
+			}
+
+			if failed {
+				t.Errorf("CreateWithTimeout() = (%d, `%s`), want (%d, `%s`)", got, gotErr, tc.want, tc.wantErr)
+			}
+
+			if err := mock.ExpectationsWereMet(); err != nil {
+				t.Errorf("sqlmock unfilled expectations: %s", err)
+			}
+		})
+	}
+}
+
+func TestExecWithTimeout(t *testing.T) {
+	type args struct {
+		tx bool
+	}
+
+	var cases = []struct {
+		intention string
+		args      args
+		wantErr   error
+	}{
+		{
+			"simple",
+			args{
+				tx: false,
+			},
+			nil,
+		},
+		{
+			"timeout",
+			args{
+				tx: false,
+			},
+			sqlmock.ErrCancelled,
+		},
+		{
+			"with tx",
+			args{
+				tx: true,
+			},
+			nil,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.intention, func(t *testing.T) {
+			db, mock, err := sqlmock.New()
+			if err != nil {
+				t.Fatalf("unable to create mock database: %s", err)
+			}
+			defer db.Close()
+
+			var tx *sql.Tx
+			if tc.args.tx {
+				mock.ExpectBegin()
+				dbTx, err := db.Begin()
+
+				if err != nil {
+					t.Errorf("unable to create tx: %v", err)
+				}
+				tx = dbTx
+			}
+
+			if !tc.args.tx {
+				mock.ExpectBegin()
+			}
+
+			expectedQuery := mock.ExpectExec("DELETE FROM item WHERE id = (.+)").WithArgs(1).WillReturnResult(sqlmock.NewResult(0, 1))
+
+			if !tc.args.tx {
+				if tc.wantErr != nil {
+					mock.ExpectRollback()
+				} else {
+					mock.ExpectCommit()
+				}
+			}
+
+			if tc.intention == "timeout" {
+				expectedQuery.WillDelayFor(time.Second * 2)
+			}
+
+			gotErr := ExecWithTimeout(db, tx, time.Second, "DELETE FROM item WHERE id = $1", 1)
+
+			failed := false
+
+			if tc.wantErr != nil && !errors.Is(gotErr, tc.wantErr) {
+				failed = true
+			}
+
+			if failed {
+				t.Errorf("ExecWithTimeout() = `%s`, want `%s`", gotErr, tc.wantErr)
+			}
+
+			if err := mock.ExpectationsWereMet(); err != nil {
+				t.Errorf("sqlmock unfilled expectations: %s", err)
+			}
+		})
+	}
+}