add connection timeout for ENT (#2115)

Signed-off-by: pxp928 <[email protected]>

cmd/guacgql/cmd/ent.go

@@ -19,6 +19,9 @@ package cmd
 
 import (
  "context"
+ "fmt"
+ "os"
+ "time"
 
  "github.com/guacsec/guac/pkg/assembler/backends"
  entbackend "github.com/guacsec/guac/pkg/assembler/backends/ent/backend"
@@ -32,10 +35,21 @@ func init() {
 }
 
 func getEnt(_ context.Context) backends.BackendArgs {
+ var connTimeout *time.Duration
+ if flags.dbConnTime != "" {
+ if timeout, err := time.ParseDuration(flags.dbConnTime); err != nil {
+ fmt.Printf("failed to parser duration with error: %v", err)
+ os.Exit(1)
+ } else {
+ connTimeout = &timeout
+ }
+ }
+
  return &entbackend.BackendOptions{
- DriverName: flags.dbDriver,
- Address: flags.dbAddress,
- Debug: flags.dbDebug,
- AutoMigrate: flags.dbMigrate,
+ DriverName: flags.dbDriver,
+ Address: flags.dbAddress,
+ Debug: flags.dbDebug,
+ AutoMigrate: flags.dbMigrate,
+ ConnectionMaxLifeTime: connTimeout,
  }
 }

cmd/guacgql/cmd/root.go

@@ -42,10 +42,11 @@ var flags = struct {
  nRealm string
 
  // Needed only if using ent backend
- dbAddress string
- dbDriver string
- dbDebug bool
- dbMigrate bool
+ dbAddress string
+ dbDriver string
+ dbDebug bool
+ dbMigrate bool
+ dbConnTime string
 
  // Needed only if using arangodb backend
  arangoAddr string
@@ -86,6 +87,7 @@ var rootCmd = &cobra.Command{
  flags.dbDriver = viper.GetString("db-driver")
  flags.dbDebug = viper.GetBool("db-debug")
  flags.dbMigrate = viper.GetBool("db-migrate")
+ flags.dbConnTime = viper.GetString("db-conn-time")
 
  flags.arangoUser = viper.GetString("arango-user")
  flags.arangoPass = viper.GetString("arango-pass")
@@ -112,7 +114,7 @@ func init() {
  "neo4j-addr", "neo4j-user", "neo4j-pass", "neo4j-realm",
  "neptune-endpoint", "neptune-port", "neptune-region", "neptune-user", "neptune-realm",
  "gql-listen-port", "gql-tls-cert-file", "gql-tls-key-file", "gql-debug", "gql-backend", "gql-trace",
- "db-address", "db-driver", "db-debug", "db-migrate",
+ "db-address", "db-driver", "db-debug", "db-migrate", "db-conn-time",
  "kv-store", "kv-redis", "kv-tikv", "enable-prometheus",
  })
  if err != nil {

pkg/assembler/backends/ent/backend/migrations.go

@@ -19,6 +19,7 @@ import (
  "context"
  "database/sql"
  "fmt"
+ "time"
 
  "entgo.io/ent/dialect"
  "github.com/guacsec/guac/pkg/assembler/backends/ent"
@@ -34,6 +35,9 @@ type BackendOptions struct {
  Address string
  Debug bool
  AutoMigrate bool
+ // sets the maximum amount of time a connection may be reused.
+ // If nil, this is set to default of 0
+ ConnectionMaxLifeTime *time.Duration
 }
 
 // GetReadOnlyClient sets up the ent backend and returns a read-only client.
@@ -68,6 +72,11 @@ func SetupBackend(ctx context.Context, options *BackendOptions) (*ent.Client, er
  return nil, fmt.Errorf("error opening db: %w", err)
  }
 
+ if options.ConnectionMaxLifeTime != nil {
+ // set timeout limit for connections
+ db.SetConnMaxLifetime(*options.ConnectionMaxLifeTime)
+ }
+
  client := ent.NewClient(ent.Driver(dialectsql.OpenDB(driver, db)))
 
  if options.AutoMigrate {

pkg/cli/store.go

@@ -81,6 +81,7 @@ func init() {
  set.String("db-driver", "postgres", "database driver to use, one of [postgres | sqlite3 | mysql] or anything supported by sql.DB")
  set.Bool("db-debug", false, "enable debug logging for database queries")
  set.Bool("db-migrate", true, "automatically run database migrations on start")
+ set.String("db-conn-time", "", "sets the maximum amount of time a connection may be reused in m, h, s, etc.")
 
  set.String("arango-addr", "http://localhost:8529", "address to arango db")
  set.String("arango-user", "", "arango user to connect to graph db")