Allow retrying a connection on startup

pkg/cli/cli.go

@@ -77,35 +77,15 @@ func initClient() {
 	}
 
 	if command.Opts.Debug {
-		fmt.Println("Server connection string:", cl.ConnectionString)
+		fmt.Println("Opening database connection using string:", cl.ConnectionString)
 	}
 
 	fmt.Println("Connecting to server...")
-	if err := cl.Test(); err != nil {
-		msg := err.Error()
-
-		// Check if we're trying to connect to the default database.
-		if command.Opts.DbName == "" && command.Opts.URL == "" {
-			// If database does not exist, allow user to connect from the UI.
-			if strings.Contains(msg, "database") && strings.Contains(msg, "does not exist") {
-				fmt.Println("Error:", msg)
-				return
-			}
-
-			// Do not bail if local server is not running.
-			if regexErrConnectionRefused.MatchString(msg) {
-				fmt.Println("Error:", msg)
-				return
-			}
-
-			// Do not bail if local auth is invalid
-			if regexErrAuthFailed.MatchString(msg) {
-				fmt.Println("Error:", msg)
-				return
-			}
-		}
 
-		exitWithMessage(msg)
+	// TODO: move args to CLI options with a default value
+	if err := testClient(*cl, 5, 3*time.Second); err != nil {
+		exitWithMessage(err.Error())
+		return
 	}
 
 	if !command.Opts.Sessions {
@@ -280,6 +260,38 @@ func openPage() {
 	}
 }
 
+// testWithRetry attempts to establish a database connection until it succeeds or
+// give up after certain number of retries.
+func testClient(cl client.Client, retryCount int, retryDelay time.Duration) (err error) {
+	allowRetry := command.Opts.DbName != "" && command.Opts.URL != ""
+
+	for i := 0; i < retryCount-1; i++ {
+		err = cl.Test()
+		if err == nil {
+			return nil
+		}
+
+		// Retrying on a default database is not allowed.
+		if !allowRetry {
+			fmt.Printf("Connection error: %v\n", err)
+			if err == client.ErrConnectionRefused || err == client.ErrAuthFailed || err == client.ErrDatabaseNotExist {
+				return nil
+			}
+			return err
+		}
+
+		// Only retry connection errors.
+		if !errors.Is(err, client.ErrConnectionRefused) {
+			return err
+		}
+
+		fmt.Printf("Connection error: %v, retrying in %v\n", err, retryDelay)
+		<-time.After(retryDelay)
+	}
+
+	return err
+}
+
 func Run() {
 	initOptions()
 	initClient()

pkg/client/client.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	neturl "net/url"
 	"reflect"
+	"regexp"
 	"strings"
 	"time"
 
@@ -21,6 +22,18 @@ import (
 	"github.com/sosedoff/pgweb/pkg/statements"
 )
 
+var (
+	regexErrAuthFailed        = regexp.MustCompile(`authentication failed`)
+	regexErrConnectionRefused = regexp.MustCompile(`(connection|actively) refused`)
+	regexErrDatabaseNotExist  = regexp.MustCompile(`database "(.*)" does not exist`)
+)
+
+var (
+	ErrAuthFailed        = errors.New("authentication failed")
+	ErrConnectionRefused = errors.New("connection refused")
+	ErrDatabaseNotExist  = errors.New("database does not exist")
+)
+
 type Client struct {
 	db               *sqlx.DB
 	tunnel           *Tunnel
@@ -179,7 +192,27 @@ func (client *Client) setServerVersion() {
 }
 
 func (client *Client) Test() error {
-	return client.db.Ping()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	err := client.db.PingContext(ctx)
+	if err == nil {
+		return nil
+	}
+
+	errMsg := err.Error()
+
+	if regexErrConnectionRefused.MatchString(errMsg) {
+		return ErrConnectionRefused
+	}
+	if regexErrAuthFailed.MatchString(errMsg) {
+		return ErrAuthFailed
+	}
+	if regexErrDatabaseNotExist.MatchString(errMsg) {
+		return ErrDatabaseNotExist
+	}
+
+	return err
 }
 
 func (client *Client) TestWithTimeout(timeout time.Duration) (result error) {