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
in-toto · Sep 27, 2023 · 5abcad0 · 5abcad0
1 parent eb5ae1d
commit 5abcad0
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 10 deletions.
diff --git a/cmd/archivista/main.go b/cmd/archivista/main.go
@@ -86,7 +86,12 @@ 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,
+		sqlstore.ClientWithMaxIdleConns(cfg.SQLStoreMaxIdleConnections),
+		sqlstore.ClientWithMaxOpenConns(cfg.SQLStoreMaxOpenConnections),
+		sqlstore.ClientWithConnMaxLifetime(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
@@ -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 number 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
@@ -28,9 +28,48 @@ import (
 	_ "github.com/lib/pq"
 )
 
+type ClientOption func(*clientOptions)
+
+type clientOptions struct {
+	maxIdleConns    int
+	maxOpenConns    int
+	connMaxLifetime time.Duration
+}
+
+// Configures a client with the specified max idle connections. Default is 10 connections
+func ClientWithMaxIdleConns(maxIdleConns int) ClientOption {
+	return func(co *clientOptions) {
+		co.maxIdleConns = maxIdleConns
+	}
+}
+
+// Configures a client with the specified max open connections. Default is 100 connections
+func ClientWithMaxOpenConns(maxOpenConns int) ClientOption {
+	return func(co *clientOptions) {
+		co.maxOpenConns = maxOpenConns
+	}
+}
+
+// Congiures a client with the specified max connection lifetime. Default is 3 minutes
+func ClientWithConnMaxLifetime(connMaxLifetime time.Duration) ClientOption {
+	return func(co *clientOptions) {
+		co.connMaxLifetime = connMaxLifetime
+	}
+}
+
 // 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, opts ...ClientOption) (*ent.Client, error) {
+	clientOpts := &clientOptions{
+		maxIdleConns:    10,
+		maxOpenConns:    100,
+		connMaxLifetime: 3 * time.Minute,
+	}
+
+	for _, opt := range opts {
+		opt(clientOpts)
+	}
+
 	var entDialect string
 	switch strings.ToUpper(sqlBackend) {
 	case "MYSQL":
@@ -56,9 +95,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(clientOpts.maxIdleConns)
+	db.SetMaxOpenConns(clientOpts.maxOpenConns)
+	db.SetConnMaxLifetime(clientOpts.connMaxLifetime)
 	sqlcommentDrv := sqlcomment.NewDriver(drv,
 		sqlcomment.WithDriverVerTag(),
 		sqlcomment.WithTags(sqlcomment.Tags{