snippets

commit e8cd095d422f2902e2fc9c71c8dd6123ac721916
parent d9e6b2563a36ca59fdb41ee22ed2c1e50cdcc73d
Author: Vetle Haflan <vetle@haflan.dev>
Date:   Wed, 31 Aug 2022 15:13:04 +0200

Add loop mode for pgtester

Diffstat:
M
pgtester/pgtester.go
 | 
92
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----------------

1 file changed, 74 insertions(+), 18 deletions(-)
diff --git a/pgtester/pgtester.go b/pgtester/pgtester.go
@@ -4,7 +4,9 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
+	"time"
 
 	_ "github.com/lib/pq"
 )
@@ -23,28 +25,82 @@ func must(err error, msg string) {
 	os.Exit(1)
 }
 
-func main() {
-	if len(os.Args) < 2 {
-		fmt.Printf("Use:\n\t%v [postgresql://]<database_url>\nURL format:\n\t%v", os.Args[0], urlFmt)
-		return
-	}
-	dbURL := pgPrefix + strings.TrimPrefix(os.Args[1], pgPrefix)
-	db, err := sql.Open("postgres", dbURL)
-	must(err, "Failed to open database:")
+var db *sql.DB
+
+func test() error {
+	var (
+		err     error
+		readNum int
+		row     *sql.Row
+	)
 	err = db.Ping()
-	must(err, "Failed to ping database")
+	if err != nil {
+		return fmt.Errorf("Failed to ping database: %w", err)
+	}
 	_, err = db.Exec(`CREATE TABLE pgtester_temp_table (num int);`)
-	must(err, "Failed to create temporary table")
+	if err != nil {
+		return fmt.Errorf("Failed to create temporary table")
+	}
 	_, err = db.Exec(`INSERT INTO pgtester_temp_table VALUES ($1);`, testNum)
-	must(err, "Failed to insert into temporary table")
-	row := db.QueryRow(`SELECT num FROM pgtester_temp_table;`)
-	var readNum int
+	if err != nil {
+		return fmt.Errorf("Failed to insert into temporary table")
+	}
+	row = db.QueryRow(`SELECT num FROM pgtester_temp_table;`)
 	err = row.Scan(&readNum)
-	must(err, "Failed to read data from temporary table")
+	if err != nil {
+		return fmt.Errorf("Failed to read data from temporary table")
+	}
 	if readNum != testNum {
-		fmt.Printf("Inserted data does not match fetched data (%v != %v)\n", readNum, testNum)
-		os.Exit(1)
+		err = fmt.Errorf("Inserted data does not match fetched data (%v != %v)\n", readNum, testNum)
+	}
+	_, nerr := db.Exec(`DROP TABLE pgtester_temp_table;`)
+	if nerr != nil {
+		if err == nil {
+			err = fmt.Errorf("Failed to delete temporary database")
+		} else {
+			err = fmt.Errorf("Failed to to delete temporary database, also: %w", err)
+		}
+	}
+	return err
+}
+
+func testPrint() {
+	err := test()
+	if err != nil {
+		fmt.Println(err)
+	} else {
+		fmt.Println("All tests passed")
+	}
+}
+
+func loop(periodSeconds int) {
+	for {
+		testPrint()
+		time.Sleep(time.Duration(periodSeconds) * time.Second)
+	}
+
+}
+
+func main() {
+	if len(os.Args) < 2 {
+		fmt.Printf("Use:\n\t%v [postgresql://]<database_url> [loop_seconds]\nURL format:\n\t%v", os.Args[0], urlFmt)
+		return
+	}
+	dbURL := pgPrefix + strings.TrimPrefix(os.Args[1], pgPrefix)
+	var err error
+	db, err = sql.Open("postgres", dbURL)
+	if err != nil {
+		fmt.Println("Failed to open database:")
+		return
+	}
+	if len(os.Args) > 2 {
+		sec, err := strconv.Atoi(os.Args[2])
+		if err != nil {
+			fmt.Println("Invalid loop period, please give a whole number of seconds")
+			return
+		}
+		loop(sec)
+	} else {
+		testPrint()
 	}
-	_, err = db.Exec(`DROP TABLE pgtester_temp_table;`)
-	must(err, "Failed to delete temporary database")
 }