Move code to pkg and cmd

Author: aswinkarthik

Makefile

@@ -0,0 +1,6 @@
+-include .env
+
+VERSION := $(shell git describe --tags)
+
+compile:
+	go build -o pg-ping -ldflags "-s -w -X main.version=$(VERSION)" main.go

cmd/root.go

@@ -0,0 +1,102 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+
+	"github.com/thecasualcoder/pg-ping/pkg/pg"
+	"github.com/urfave/cli"
+)
+
+// Execute the app
+func Execute(app *cli.App) error {
+	app.Name = "pg-ping"
+	app.Usage = "Ping your postgres continously"
+	app.Action = run
+	app.Flags = []cli.Flag{
+		cli.BoolFlag{
+			Name:   "once",
+			Usage:  "Ping only once and quit",
+			EnvVar: "PGPING_ONCE",
+		},
+		cli.BoolFlag{
+			Name:   "debug",
+			Usage:  "Print debug logs",
+			EnvVar: "PGPING_DEBUG",
+		},
+		cli.StringFlag{
+			Name:   "username, U",
+			Usage:  "Username to connect to postgres",
+			EnvVar: "PGPING_USERNAME",
+		},
+		cli.StringFlag{
+			Name:   "password, p",
+			Usage:  "Password to connect to postgres",
+			EnvVar: "PGPING_PASSWORD",
+		},
+		cli.StringFlag{
+			Name:   "host, h",
+			Usage:  "Host of postgres server",
+			EnvVar: "PGPING_HOST",
+		},
+		cli.StringFlag{
+			Name:   "dbname, d",
+			Usage:  "DBName to connect to",
+			EnvVar: "PGPING_DBNAME",
+		},
+		cli.IntFlag{
+			Name:   "frequency, f",
+			Usage:  "Frequency to ping",
+			EnvVar: "PGPING_FREQUENCY_IN_MS",
+		},
+		cli.StringFlag{
+			Name:   "query, c",
+			Usage:  "Query to run",
+			EnvVar: "PGPING_QUERY",
+			Value:  "SELECT 1",
+		},
+	}
+	return app.Run(os.Args)
+}
+
+func run(c *cli.Context) error {
+	if len(c.Args()) > 0 {
+		cli.ShowAppHelp(c)
+		return fmt.Errorf("args are not allowed")
+	}
+	conf := pg.Config{
+		Username:      c.String("username"),
+		Password:      c.String("password"),
+		Host:          c.String("host"),
+		DBName:        c.String("dbname"),
+		FrequencyInMS: int32(c.Int("frequency")),
+		Query:         c.String("query"),
+	}
+
+	encoder := json.NewEncoder(os.Stdout)
+
+	if c.BoolT("debug") {
+		encoder.Encode(conf)
+	}
+
+	db, err := pg.NewDB(conf)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	var resultChan chan pg.SQLResult
+
+	if c.BoolT("once") {
+		resultChan = db.PingOnce()
+	} else {
+		resultChan = db.Ping()
+	}
+
+	for r := range resultChan {
+		encoder.Encode(r)
+	}
+
+	return nil
+}

go.mod

@@ -1,8 +1,8 @@
-module pg-ping
+module github.com/thecasualcoder/pg-ping
 
 go 1.12
 
 require (
-	github.com/kelseyhightower/envconfig v1.3.0
 	github.com/lib/pq v1.0.0
+	github.com/urfave/cli v1.20.0
 )

go.sum

@@ -1,4 +1,4 @@
-github.com/kelseyhightower/envconfig v1.3.0 h1:IvRS4f2VcIQy6j4ORGIf9145T/AsUB+oY8LyvN8BXNM=
-github.com/kelseyhightower/envconfig v1.3.0/go.mod h1:cccZRl6mQpaq41TPp5QxidR+Sa3axMbJDNb//FQX6Gg=
 github.com/lib/pq v1.0.0 h1:X5PMW56eZitiTeO7tKzZxFCSpbFZJtkMMooicw2us9A=
 github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
+github.com/urfave/cli v1.20.0 h1:fDqGv3UG/4jbVl/QkFwEdddtEDjh/5Ov6X+0B/3bPaw=
+github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA=

main.go

@@ -1,80 +1,28 @@
 package main
 
 import (
-	"context"
-	"database/sql"
-	"encoding/json"
-	"log"
+	"fmt"
 	"os"
-	"time"
 
-	"github.com/kelseyhightower/envconfig"
-	_ "github.com/lib/pq"
+	"github.com/thecasualcoder/pg-ping/cmd"
+	"github.com/urfave/cli"
 )
 
-func main() {
-	var conf PGConfig
-	err := envconfig.Process("PGPING", &conf)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	if conf.Debug {
-		json.NewEncoder(os.Stdout).Encode(conf)
-	}
+var version string
 
-	db, err := sql.Open("postgres", conf.ConnStr())
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer db.Close()
-
-	ticker := time.NewTicker(conf.GetFrequency())
-	encoder := json.NewEncoder(os.Stdout)
-	for range ticker.C {
-		encoder.Encode(executeQuery(db, conf.GetQuery()))
-	}
-}
+const defaultVersion = "dev"
 
-// SQLResult  tracks sql response and time taken
-type SQLResult struct {
-	Value          string
-	TimeTakenInMS  float64
-	Failed         bool
-	FailureMessage string
-}
-
-func executeQuery(db *sql.DB, query string) SQLResult {
-	res := SQLResult{}
-
-	ctx, cancelFunc := context.WithCancel(context.Background())
-	tenSecondTimer := time.NewTimer(10 * time.Second)
-	go func() {
-		<-tenSecondTimer.C
-		cancelFunc()
-	}()
-
-	start := time.Now()
-	rows, err := db.QueryContext(ctx, query)
-	res.TimeTakenInMS = time.Since(start).Seconds() * 1000
-
-	if err != nil {
-		res.Failed = true
-		res.FailureMessage = err.Error()
-		return res
+func main() {
+	var app = cli.NewApp()
+	if version == "" {
+		version = defaultVersion
 	}
-	defer rows.Close()
+	app.Version = version
 
-	for rows.Next() {
-		var name string
-		if err := rows.Scan(&name); err != nil {
-			res.Failed = true
-			res.FailureMessage = err.Error()
-			return res
-		}
+	cli.HelpFlag = cli.BoolFlag{Name: "help"}
 
-		res.Value = name
+	if err := cmd.Execute(app); err != nil {
+		fmt.Fprintf(os.Stderr, "pg-ping failed: %v\n", err)
+		os.Exit(1)
 	}
-
-	return res
 }

pg_config.go renamed to pkg/pg/config.go

@@ -1,12 +1,12 @@
-package main
+package pg
 
 import (
 	"fmt"
 	"time"
 )
 
-// PGConfig contains configurations to connect to a PG database
-type PGConfig struct {
+// Config contains configurations to connect to a PG database
+type Config struct {
 	Username      string
 	Password      string
 	Host          string
@@ -17,12 +17,12 @@ type PGConfig struct {
 }
 
 // ConnStr returns a connection string to connect to postgres
-func (c *PGConfig) ConnStr() string {
+func (c *Config) ConnStr() string {
 	return fmt.Sprintf("postgres://%s:%s@%s/%s?sslmode=disable", c.Username, c.Password, c.Host, c.DBName)
 }
 
 // GetQuery returns the query to use to ping
-func (c *PGConfig) GetQuery() string {
+func (c *Config) GetQuery() string {
 	if c.Query == "" {
 		return "select 1"
 	}
@@ -31,7 +31,7 @@ func (c *PGConfig) GetQuery() string {
 }
 
 // GetFrequency returns the frequence in MS in which the query should be run
-func (c *PGConfig) GetFrequency() time.Duration {
+func (c *Config) GetFrequency() time.Duration {
 	if c.FrequencyInMS == 0 {
 		return time.Second
 	}

pkg/pg/db.go

@@ -0,0 +1,95 @@
+package pg
+
+import (
+	"context"
+	"database/sql"
+	"time"
+
+	_ "github.com/lib/pq"
+)
+
+// DB represents a pingable DB
+type DB struct {
+	db   *sql.DB
+	conf Config
+}
+
+// NewDB creates a new DB connection
+func NewDB(conf Config) (*DB, error) {
+	db, err := sql.Open("postgres", conf.ConnStr())
+	if err != nil {
+		return nil, err
+	}
+	return &DB{db: db, conf: conf}, nil
+}
+
+// Close a pingable DB
+func (db *DB) Close() error {
+	return db.db.Close()
+}
+
+// PingOnce will execute query only once
+func (db *DB) PingOnce() chan SQLResult {
+	result := make(chan SQLResult, 1)
+	go func() {
+		result <- executeQuery(db.db, db.conf.GetQuery())
+		close(result)
+	}()
+	return result
+}
+
+// Ping will execute query indefinitely
+func (db *DB) Ping() chan SQLResult {
+	result := make(chan SQLResult, 10)
+	go func() {
+		ticker := time.NewTicker(db.conf.GetFrequency())
+		for range ticker.C {
+			result <- executeQuery(db.db, db.conf.GetQuery())
+		}
+		close(result)
+	}()
+	return result
+}
+
+// SQLResult tracks sql response and time taken
+type SQLResult struct {
+	Value          string
+	TimeTakenInMS  float64
+	Failed         bool
+	FailureMessage string
+}
+
+func executeQuery(db *sql.DB, query string) SQLResult {
+	res := SQLResult{}
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	tenSecondTimer := time.NewTimer(10 * time.Second)
+	go func() {
+		<-tenSecondTimer.C
+		cancelFunc()
+	}()
+
+	start := time.Now()
+	rows, err := db.QueryContext(ctx, query)
+	res.TimeTakenInMS = time.Since(start).Seconds() * 1000
+
+	if err != nil {
+		res.Failed = true
+		res.FailureMessage = err.Error()
+		return res
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var name string
+		if err := rows.Scan(&name); err != nil {
+			res.Failed = true
+			res.FailureMessage = err.Error()
+			return res
+		}
+
+		res.Value = name
+	}
+
+	return res
+}