From 1fee7356c7c9bd7bb7c5b73cf64adaf065416ddb Mon Sep 17 00:00:00 2001
From: chinmay <cbharadwaj@infoblox.com>
Date: Wed, 6 Nov 2024 22:21:41 +0530
Subject: [PATCH] add max open/idle db connection

---
 atlas/templates/cmd/server/config.go.gotmpl |  6 ++++++
 atlas/templates/cmd/server/grpc.go.gotmpl   | 11 ++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/atlas/templates/cmd/server/config.go.gotmpl b/atlas/templates/cmd/server/config.go.gotmpl
index ca4be28..ee8b7cf 100644
--- 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")
+	flagDatabaseOption   = pflag.String("database.max.open.conns", defaultDatabaseOpenConn, "Database max open connections")
+	flagDatabaseOption   = 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
index a14f024..399f105 100644
--- 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 }}