add max open/idle db connection

infobloxopen · Nov 6, 2024 · db62795 · db62795
1 parent 9b3ab52
commit db62795
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/atlas/templates/cmd/server/config.go.gotmpl b/atlas/templates/cmd/server/config.go.gotmpl
@@ -29,6 +29,9 @@ const (
 	defaultDatabaseSSL      = "disable"
 	defaultDatabaseOption   = ""
 
+	defaultDatabaseOpenConn    = 10
+	defaultDatabaseMaxIdleConn = 10
+
 	// PubSub
 	defaultPubsubEnable    =  {{ if .WithPubsub }}true{{ else }}false{{ end }}
 	defaultPubsubAddress   = "pubsub.atlas"
@@ -113,6 +116,9 @@ var (
 	flagDatabasePassword = pflag.String("database.password", defaultDatabasePassword, "database password")
 	flagDatabaseSSL      = pflag.String("database.ssl", defaultDatabaseSSL, "database ssl mode")
 	flagDatabaseOption   = pflag.String("database.option", defaultDatabaseOption, "define custom option to db driver")
+	flagDatabaseMaxOpenConns   = pflag.String("database.max.open.conns", defaultDatabaseOpenConn, "Database max open connections")
+	flagDatabaseMaxIdleConns   = pflag.String("database.max.idle.conns", defaultDatabaseMaxIdleConn, "Database max idle connections")
+
 
 	flagPubsubEnable       = pflag.Bool("atlas.pubsub.enable", defaultPubsubEnable, "enable application with pubsub")
 	flagPubsubAddress      = pflag.String("atlas.pubsub.address", defaultPubsubAddress, "address or FQDN of the pubsub service")

diff --git a/atlas/templates/cmd/server/grpc.go.gotmpl b/atlas/templates/cmd/server/grpc.go.gotmpl
@@ -50,7 +50,9 @@ func NewGRPCServer(logger *logrus.Logger{{ if or .WithPublishTopic .WithSubscrib
 	db, err := gorm.Open("postgres", dbConnectionString)
 	if err != nil {
 		return nil, err
-	}{{ end }}
+	}
+	setMaxOpenIdleConnection(db)
+	{{ end }}
 	// register service implementation with the grpcServer
 	s, err := svc.NewBasicServer({{ if or .WithPublishTopic .WithSubscribeTopic }}pubsub{{ if .WithDatabase }}, {{ end }}{{ end }}{{ if .WithDatabase }}db{{ end }})
 	if err != nil {
@@ -62,3 +64,10 @@ func NewGRPCServer(logger *logrus.Logger{{ if or .WithPublishTopic .WithSubscrib
 
 	return grpcServer, nil
 }
+
+{{ if .WithDatabase }}
+// setMaxOpenIdleConnection sets the db connection string
+func setMaxOpenIdleConnection(db *gorm.DB) {
+	db.DB().SetMaxOpenConns(viper.GetInt("database.max.open.conns"))
+	db.DB().SetMaxIdleConns(viper.GetInt("database.max.idle.conns"))
+}{{ end }}