From 37a99771eca6c670f0d3428e978f763f45be9bd2 Mon Sep 17 00:00:00 2001
From: Mikhail Swift <mikhail@testifysec.com>
Date: Tue, 26 Sep 2023 17:18:10 -0500
Subject: [PATCH] feat: add configurable sql connection parameters

Allows the user to configure the SQL connection pool's max idle and open
connections as well as a connections maximum lifetime
---
 cmd/archivista/main.go                      |  2 +-
 internal/config/config.go                   | 15 ++++++++++-----
 internal/metadatastorage/sqlstore/client.go |  8 ++++----
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/cmd/archivista/main.go b/cmd/archivista/main.go
index e6aed61b..4aec745e 100644
--- a/cmd/archivista/main.go
+++ b/cmd/archivista/main.go
@@ -86,7 +86,7 @@ func main() {
 		logrus.Fatalf("error initializing storage clients: %+v", err)
 	}
 
-	entClient, err := sqlstore.NewEntClient(cfg.SQLStoreBackend, cfg.SQLStoreConnectionString)
+	entClient, err := sqlstore.NewEntClient(cfg.SQLStoreBackend, cfg.SQLStoreConnectionString, cfg.SQLStoreMaxIdleConnections, cfg.SQLStoreMaxOpenConnections, cfg.SQLStoreConnectionMaxLifetime)
 	if err != nil {
 		logrus.Fatalf("could not create ent client: %+v", err)
 	}
diff --git a/internal/config/config.go b/internal/config/config.go
index 92d29892..26559b6e 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -18,6 +18,7 @@ import (
 	"errors"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/kelseyhightower/envconfig"
 	"github.com/sirupsen/logrus"
@@ -28,11 +29,15 @@ type Config struct {
 	LogLevel         string   `default:"INFO" desc:"Log level" split_words:"true"`
 	CORSAllowOrigins []string `default:"" desc:"Comma separated list of origins to allow CORS requests from" split_words:"true"`
 
-	EnableSPIFFE             bool   `default:"TRUE" desc:"*** Enable SPIFFE support" split_words:"true"`
-	SPIFFEAddress            string `default:"unix:///tmp/spire-agent/public/api.sock" desc:"SPIFFE server address" split_words:"true"`
-	SPIFFETrustedServerId    string `default:"" desc:"Trusted SPIFFE server ID; defaults to any" split_words:"true"`
-	SQLStoreConnectionString string `default:"root:example@tcp(db)/testify" desc:"SQL store connection string" split_words:"true"`
-	SQLStoreBackend          string `default:"MYSQL" desc:"SQL backend to use. Options are MYSQL, PSQL" split_words:"true"`
+	EnableSPIFFE          bool   `default:"TRUE" desc:"*** Enable SPIFFE support" split_words:"true"`
+	SPIFFEAddress         string `default:"unix:///tmp/spire-agent/public/api.sock" desc:"SPIFFE server address" split_words:"true"`
+	SPIFFETrustedServerId string `default:"" desc:"Trusted SPIFFE server ID; defaults to any" split_words:"true"`
+
+	SQLStoreConnectionString      string        `default:"root:example@tcp(db)/testify" desc:"SQL store connection string" split_words:"true"`
+	SQLStoreBackend               string        `default:"MYSQL" desc:"SQL backend to use. Options are MYSQL, PSQL" split_words:"true"`
+	SQLStoreMaxIdleConnections    int           `default:"10" desc:"Maximum numver of connections in the idle connection pool" split_words:"true"`
+	SQLStoreMaxOpenConnections    int           `default:"100" desc:"Maximum number of open connections to the database" split_words:"true"`
+	SQLStoreConnectionMaxLifetime time.Duration `default:"3m" desc:"Maximum amount of time a connection may be reused" split_words:"true"`
 
 	StorageBackend             string `default:"" desc:"Backend to use for attestation storage. Options are FILE, BLOB, or empty string for disabled." split_words:"true"`
 	FileServeOn                string `default:"" desc:"What address to serve files on. Only valid when using FILE storage backend." split_words:"true"`
diff --git a/internal/metadatastorage/sqlstore/client.go b/internal/metadatastorage/sqlstore/client.go
index 45a81d91..fe212814 100644
--- a/internal/metadatastorage/sqlstore/client.go
+++ b/internal/metadatastorage/sqlstore/client.go
@@ -30,7 +30,7 @@ import (
 
 // NewEntClient creates an ent client for use in the sqlmetadata store.
 // Valid backends are MYSQL and PSQL.
-func NewEntClient(sqlBackend string, connectionString string) (*ent.Client, error) {
+func NewEntClient(sqlBackend string, connectionString string, maxxIdleConns, maxOpenConns int, connMaxxLifetime time.Duration) (*ent.Client, error) {
 	var entDialect string
 	switch strings.ToUpper(sqlBackend) {
 	case "MYSQL":
@@ -56,9 +56,9 @@ func NewEntClient(sqlBackend string, connectionString string) (*ent.Client, erro
 	}
 
 	db := drv.DB()
-	db.SetMaxIdleConns(10)
-	db.SetMaxOpenConns(100)
-	db.SetConnMaxLifetime(3 * time.Minute)
+	db.SetMaxIdleConns(maxxIdleConns)
+	db.SetMaxOpenConns(maxOpenConns)
+	db.SetConnMaxLifetime(connMaxxLifetime)
 	sqlcommentDrv := sqlcomment.NewDriver(drv,
 		sqlcomment.WithDriverVerTag(),
 		sqlcomment.WithTags(sqlcomment.Tags{