Attempt to connect database a few times before timing out (#241)

Author: Chris Vermeulen

cmd/migrate.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"github.com/compliance-framework/configuration-service/internal/config"
 	"github.com/compliance-framework/configuration-service/internal/service"
 	"github.com/spf13/cobra"
@@ -50,8 +51,8 @@ func migrateUp(cmd *cobra.Command, args []string) {
 	defer zapLogger.Sync() // flushes buffer, if any
 
 	cfg := config.NewConfig(sugar)
-
-	db, err := service.ConnectSQLDb(cfg, sugar)
+	ctx := context.Background()
+	db, err := service.ConnectSQLDb(ctx, cfg, sugar)
 	if err != nil {
 		panic("failed to connect database")
 	}
@@ -71,8 +72,8 @@ func migrateDown(cmd *cobra.Command, args []string) {
 	defer zapLogger.Sync() // flushes buffer, if any
 
 	cfg := config.NewConfig(sugar)
-
-	db, err := service.ConnectSQLDb(cfg, sugar)
+	ctx := context.Background()
+	db, err := service.ConnectSQLDb(ctx, cfg, sugar)
 	if err != nil {
 		panic("failed to connect database")
 	}

cmd/oscal/import.go

@@ -1,6 +1,7 @@
 package oscal
 
 import (
+	"context"
 	"encoding/json"
 	"github.com/compliance-framework/configuration-service/internal/service/relational"
 	oscalTypes_1_1_3 "github.com/defenseunicorns/go-oscal/src/types/oscal-1-1-3"
@@ -45,7 +46,7 @@ func importOscal(cmd *cobra.Command, args []string) {
 		panic(err)
 	}
 
-	db, err := service.ConnectSQLDb(config, sugar)
+	db, err := service.ConnectSQLDb(context.Background(), config, sugar)
 	if err != nil {
 		panic("failed to connect database")
 	}

cmd/run.go

@@ -34,7 +34,7 @@ func RunServer(cmd *cobra.Command, args []string) {
 
 	config := config.NewConfig(sugar)
 
-	db, err := service.ConnectSQLDb(config, sugar)
+	db, err := service.ConnectSQLDb(ctx, config, sugar)
 	if err != nil {
 		sugar.Fatal("Failed to connect to SQL database", "err", err)
 	}

cmd/seed/evidence.go

@@ -1,6 +1,7 @@
 package seed
 
 import (
+	"context"
 	"fmt"
 	"github.com/compliance-framework/configuration-service/internal"
 	"github.com/compliance-framework/configuration-service/internal/service/relational"
@@ -59,7 +60,7 @@ func generateEvidence(cmd *cobra.Command, args []string) {
 	defer zapLogger.Sync() // flushes buffer, if any
 
 	cmdConfig := config.NewConfig(sugar)
-	db, err := service.ConnectSQLDb(cmdConfig, sugar)
+	db, err := service.ConnectSQLDb(context.Background(), cmdConfig, sugar)
 	if err != nil {
 		panic("failed to connect database")
 	}

cmd/seed/heartbeat.go

@@ -1,6 +1,7 @@
 package seed
 
 import (
+	"context"
 	"github.com/google/uuid"
 	"github.com/schollz/progressbar/v3"
 	"log"
@@ -53,7 +54,7 @@ func generateHeartbeats(cmd *cobra.Command, args []string) {
 	defer zapLogger.Sync() // flushes buffer, if any
 
 	cmdConfig := config.NewConfig(sugar)
-	db, err := service.ConnectSQLDb(cmdConfig, sugar)
+	db, err := service.ConnectSQLDb(context.Background(), cmdConfig, sugar)
 	if err != nil {
 		panic("failed to connect database")
 	}

cmd/users/create.go

@@ -1,6 +1,7 @@
 package users
 
 import (
+	"context"
 	"crypto/rand"
 	"errors"
 	"github.com/compliance-framework/configuration-service/internal/config"
@@ -34,13 +35,14 @@ func newUserAddCmd() *cobra.Command {
 }
 
 func addUser(cmd *cobra.Command, args []string) {
+	ctx := context.Background()
 	logger, err := zap.NewProduction()
 	cobra.CheckErr(err)
 	defer logger.Sync() // flushes buffer, if any
 	sugar := logger.Sugar()
 
 	config := config.NewConfig(sugar)
-	db, err := service.ConnectSQLDb(config, sugar)
+	db, err := service.ConnectSQLDb(ctx, config, sugar)
 
 	if err != nil {
 		sugar.Errorw("Failed to connect to database", "error", err)

cmd/users/update.go

@@ -1,6 +1,7 @@
 package users
 
 import (
+	"context"
 	"github.com/compliance-framework/configuration-service/internal/config"
 	"github.com/compliance-framework/configuration-service/internal/service"
 	"github.com/compliance-framework/configuration-service/internal/service/relational"
@@ -33,13 +34,14 @@ func updateUserCmd() *cobra.Command {
 }
 
 func updateUser(cmd *cobra.Command, args []string) {
+	ctx := context.Background()
 	logger, err := zap.NewProduction()
 	cobra.CheckErr(err)
 	defer logger.Sync() // flushes buffer, if any
 	sugar := logger.Sugar()
 
 	config := config.NewConfig(sugar)
-	db, err := service.ConnectSQLDb(config, sugar)
+	db, err := service.ConnectSQLDb(ctx, config, sugar)
 
 	if err != nil {
 		sugar.Errorw("Failed to connect to database", "error", err)

internal/service/connector.go

@@ -1,7 +1,10 @@
 package service
 
 import (
+	"context"
 	"errors"
+	"strings"
+	"time"
 
 	"github.com/compliance-framework/configuration-service/internal/config"
 	"github.com/compliance-framework/configuration-service/internal/logging"
@@ -11,7 +14,7 @@ import (
 	"gorm.io/gorm/logger"
 )
 
-func ConnectSQLDb(config *config.Config, sugar *zap.SugaredLogger) (*gorm.DB, error) {
+func ConnectSQLDb(ctx context.Context, config *config.Config, sugar *zap.SugaredLogger) (*gorm.DB, error) {
 	gormLogLevel := logger.Warn
 	if config.DBDebug {
 		gormLogLevel = logger.Info
@@ -25,10 +28,42 @@ func ConnectSQLDb(config *config.Config, sugar *zap.SugaredLogger) (*gorm.DB, er
 
 	switch config.DBDriver {
 	case "postgres":
-		db, err = gorm.Open(postgres.Open(config.DBConnectionString), &gorm.Config{
+		dialect := postgres.New(postgres.Config{
+			DSN: config.DBConnectionString,
+		})
+		db, err = gorm.Open(dialect, &gorm.Config{
+			DisableAutomaticPing:                     true,
 			DisableForeignKeyConstraintWhenMigrating: true,
 			Logger:                                   logging.NewZapGormLogger(sugar, gormLogLevel),
 		})
+
+		pdb, err := db.DB()
+		if err != nil {
+			return nil, err
+		}
+
+		timeoutCtx, cancel := context.WithTimeout(ctx, time.Second*10)
+		defer cancel()
+
+		for true {
+			err = pdb.Ping()
+			if err == nil {
+				sugar.Warn("Connected to database")
+				break
+			}
+
+			if strings.Contains(err.Error(), "failed to connect") {
+				// The connection failed, we should see if we have timed out and return
+				if errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) {
+					sugar.Warn("Timed out trying to connect to database. Returning")
+					return nil, err
+				}
+			} else {
+				return nil, err
+			}
+			sugar.Warn("Failed to connect to database. Retrying in 0.5 seconds")
+			time.Sleep(time.Millisecond * 500)
+		}
 	default:
 		return nil, errors.New("unsupported DB driver: " + config.DBDriver)
 	}