[datastore] Move flags and ConnectParameters to a dedicated package, merge datastoreutils into datastore

Author: the-glu

cmds/core-service/main.go

@@ -23,7 +23,7 @@ import (
 	aux "github.com/interuss/dss/pkg/aux_"
 	auxc "github.com/interuss/dss/pkg/aux_/store/datastore"
 	"github.com/interuss/dss/pkg/build"
-	"github.com/interuss/dss/pkg/datastoreutils"
+	"github.com/interuss/dss/pkg/datastore"
 	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/dss/pkg/rid/application"
 	rid_v1 "github.com/interuss/dss/pkg/rid/server/v1"
@@ -350,7 +350,7 @@ func main() {
 	backoff := 0
 	for {
 		if err := RunHTTPServer(ctx, cancel, *address, *locality); err != nil {
-			if stacktrace.GetCode(err) == datastoreutils.CodeRetryable {
+			if stacktrace.GetCode(err) == datastore.CodeRetryable {
 				logger.Info(fmt.Sprintf("Prerequisites not yet satisfied; waiting %.fs to retry...", backoffs[backoff].Seconds()), zap.Error(err))
 				time.Sleep(backoffs[backoff])
 				if backoff < len(backoffs)-1 {

cmds/db-manager/migration/migrate.go

@@ -12,7 +12,7 @@ import (
 
 	"github.com/coreos/go-semver/semver"
 	"github.com/interuss/dss/pkg/datastore"
-	datastoreflags "github.com/interuss/dss/pkg/datastore/flags"
+	"github.com/interuss/dss/pkg/datastore/params"
 
 	"github.com/interuss/stacktrace"
 	"github.com/spf13/cobra"
@@ -227,7 +227,7 @@ func migrate(cmd *cobra.Command, _ []string) error {
 
 func connectTo(ctx context.Context, dbName string) (*datastore.Datastore, error) {
 	// Connect to database server
-	connectParameters := datastoreflags.ConnectParameters()
+	connectParameters := params.GetConnectParameters()
 	connectParameters.DBName = dbName
 	return datastore.Dial(ctx, connectParameters)
 }

pkg/aux_/accepted_ca_certs.go

@@ -7,12 +7,12 @@ import (
 	"strings"
 
 	restapi "github.com/interuss/dss/pkg/api/auxv1"
-	"github.com/interuss/dss/pkg/datastore/flags"
+	"github.com/interuss/dss/pkg/datastore/params"
 )
 
 func (a *Server) GetAcceptedCAs(ctx context.Context, req *restapi.GetAcceptedCAsRequest) restapi.GetAcceptedCAsResponseSet {
 
-	connectParameters := flags.ConnectParameters()
+	connectParameters := params.GetConnectParameters()
 
 	CAFile := connectParameters.GetCAFile()
 

pkg/aux_/instance_ca_certs.go

@@ -7,12 +7,12 @@ import (
 	"strings"
 
 	restapi "github.com/interuss/dss/pkg/api/auxv1"
-	"github.com/interuss/dss/pkg/datastore/flags"
+	"github.com/interuss/dss/pkg/datastore/params"
 )
 
 func (a *Server) GetInstanceCAs(ctx context.Context, req *restapi.GetInstanceCAsRequest) restapi.GetInstanceCAsResponseSet {
 
-	connectParameters := flags.ConnectParameters()
+	connectParameters := params.GetConnectParameters()
 
 	CAFile := connectParameters.GetInstanceCAFile()
 

pkg/aux_/store/datastore/store.go

@@ -5,8 +5,7 @@ import (
 
 	"github.com/interuss/dss/pkg/aux_/repos"
 	"github.com/interuss/dss/pkg/datastore"
-	"github.com/interuss/dss/pkg/datastore/flags"
-	"github.com/interuss/dss/pkg/datastoreutils"
+	"github.com/interuss/dss/pkg/datastore/params"
 	"github.com/interuss/dss/pkg/logging"
 	dssql "github.com/interuss/dss/pkg/sql"
 	"github.com/jonboulle/clockwork"
@@ -34,7 +33,7 @@ func newStore(ctx context.Context, db *datastore.Datastore, logger *zap.Logger)
 
 	s := &Store{}
 
-	base, err := datastore.NewStore(ctx, db, flags.ConnectParameters().MaxRetries, currentCrdbMajorSchemaVersion, currentYugabyteMajorSchemaVersion, func(q dssql.Queryable) repos.Repository {
+	base, err := datastore.NewStore(ctx, db, params.GetConnectParameters().MaxRetries, currentCrdbMajorSchemaVersion, currentYugabyteMajorSchemaVersion, func(q dssql.Queryable) repos.Repository {
 		return &repo{
 			Queryable: q,
 			clock:     s.Clock,
@@ -51,7 +50,7 @@ func newStore(ctx context.Context, db *datastore.Datastore, logger *zap.Logger)
 
 func Dial(ctx context.Context, logger *zap.Logger, withCheckCron bool) (*Store, error) {
 
-	store, err := datastoreutils.DialStore(ctx, "aux", withCheckCron, func(db *datastore.Datastore) (*Store, error) {
+	store, err := datastore.DialStore(ctx, "aux", withCheckCron, func(db *datastore.Datastore) (*Store, error) {
 		return newStore(ctx, db, logger)
 	})
 

pkg/aux_/store/datastore/store_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/interuss/dss/pkg/datastore"
-	"github.com/interuss/dss/pkg/datastore/flags"
+	"github.com/interuss/dss/pkg/datastore/params"
 	"github.com/interuss/dss/pkg/logging"
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
@@ -16,7 +16,7 @@ var (
 )
 
 func setUpStore(ctx context.Context, t *testing.T) (*Store, func()) {
-	connectParameters := flags.ConnectParameters()
+	connectParameters := params.GetConnectParameters()
 	if connectParameters.Host == "" || connectParameters.Port == 0 {
 		t.Skip()
 	}
@@ -32,7 +32,7 @@ func setUpStore(ctx context.Context, t *testing.T) (*Store, func()) {
 	}
 }
 
-func newTestStore(ctx context.Context, t *testing.T, connectParameters datastore.ConnectParameters) (*Store, error) {
+func newTestStore(ctx context.Context, t *testing.T, connectParameters params.ConnectParameters) (*Store, error) {
 	db, err := datastore.Dial(ctx, connectParameters)
 	require.NoError(t, err)
 

pkg/datastore/connectParameters.go

@@ -3,22 +3,82 @@ package datastore
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/coreos/go-semver/semver"
 	"github.com/exaring/otelpgx"
+	"github.com/interuss/dss/pkg/datastore/params"
+	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/stacktrace"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/robfig/cron/v3"
+	"go.uber.org/zap"
 )
 
 type Datastore struct {
 	Version *Version
 	Pool    *pgxpool.Pool
 }
 
+const (
+	CodeRetryable = stacktrace.ErrorCode(1)
+)
+
 var UnknownVersion = &semver.Version{}
 
-func Dial(ctx context.Context, connParams ConnectParameters) (*Datastore, error) {
+func checkDatabase(ctx context.Context, db *Datastore, databaseName string) {
+	logger := logging.WithValuesFromContext(ctx, logging.Logger)
+	statsPtr := db.Pool.Stat()
+	if int(statsPtr.TotalConns()) == 0 {
+		logger.Warn("Failed periodic DB Ping (TotalConns=0)", zap.String("Database", databaseName))
+	} else {
+		logger.Info("Successful periodic DB Ping", zap.String("Database", databaseName))
+	}
+}
+
+func DialStore[S any](ctx context.Context, dbName string, withCheckCron bool, newStore func(*Datastore) (S, error)) (S, error) {
+
+	var zero S
+
+	cp := params.GetConnectParameters()
+	cp.DBName = dbName
+
+	db, err := Dial(ctx, cp)
+
+	if err != nil {
+		if strings.Contains(err.Error(), "connect: connection refused") {
+			return zero, stacktrace.PropagateWithCode(err, CodeRetryable, "Failed to connect to datastore server for %s", dbName)
+		}
+		return zero, stacktrace.Propagate(err, "Failed to connect to %s database", dbName)
+	}
+	s, err := newStore(db)
+	if err != nil {
+		db.Pool.Close()
+		if strings.Contains(err.Error(), "connect: connection refused") || strings.Contains(err.Error(), fmt.Sprintf("database \"%s\" does not exist", dbName)) || strings.Contains(err.Error(), "database has not been bootstrapped with Schema Manager") {
+			return zero, stacktrace.PropagateWithCode(err, CodeRetryable, "Failed to create %s store", dbName)
+		}
+		return zero, stacktrace.Propagate(err, "Failed to create %s store", dbName)
+	}
+
+	if withCheckCron {
+		c := cron.New()
+		if _, err := c.AddFunc("@every 1m", func() { checkDatabase(ctx, db, dbName) }); err != nil {
+			db.Pool.Close()
+			return zero, stacktrace.Propagate(err, "Failed to schedule db check for %s", dbName)
+		}
+		c.Start()
+
+		go func() {
+			<-ctx.Done()
+			c.Stop()
+		}()
+	}
+
+	return s, nil
+}
+
+func Dial(ctx context.Context, connParams params.ConnectParameters) (*Datastore, error) {
 	dsn, err := connParams.BuildDSN()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "Failed to create connection config for pgx")