From ac2fb0fb2a4b405f10cc1e0789cae66021a3021e Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Mon, 3 Feb 2025 15:07:36 +0200
Subject: [PATCH] cmd/atlas/internal/docker: set client conn as ephemeral and
 limit its lifetime

---
 cmd/atlas/internal/docker/docker.go | 32 +++++++++++++++++++++++++++++
 sql/sqlclient/client.go             |  6 ++++++
 2 files changed, 38 insertions(+)

diff --git a/cmd/atlas/internal/docker/docker.go b/cmd/atlas/internal/docker/docker.go
index 5118455c56d..aed96368e71 100644
--- a/cmd/atlas/internal/docker/docker.go
+++ b/cmd/atlas/internal/docker/docker.go
@@ -47,6 +47,8 @@ type (
 		Database string
 		// Out is a custom writer to send docker cli output to.
 		Out io.Writer
+		// ConnOptions allows configuring the underlying connection pool.
+		ConnOptions *ConnOptions
 	}
 	// A Container is an instance of a created container.
 	Container struct {
@@ -56,6 +58,13 @@ type (
 		// Port on the host this containers service is bound to.
 		Port string
 	}
+	// ConnOptions allows configuring the underlying connection pool.
+	ConnOptions struct {
+		MaxOpen     int
+		MaxIdle     int
+		MaxLifetime time.Duration
+		MaxIdleTime time.Duration
+	}
 	// ConfigOption allows configuring Config with functional arguments.
 	ConfigOption func(*Config) error
 )
@@ -369,6 +378,14 @@ func Setup(s ...string) ConfigOption {
 	}
 }
 
+// Conn sets the config connection configuration.
+func Conn(s *ConnOptions) ConfigOption {
+	return func(c *Config) error {
+		c.ConnOptions = s
+		return nil
+	}
+}
+
 // Run pulls and starts a new docker container from the Config.
 func (c *Config) Run(ctx context.Context) (*Container, error) {
 	// Make sure the configuration is not missing critical values.
@@ -586,6 +603,21 @@ func OpenWithOpts(ctx context.Context, u *url.URL, opts ...ConfigOption) (client
 	if client, err = sqlclient.Open(ctx, u1.String()); err != nil {
 		return nil, err
 	}
+	client.Ephemeral = true
+	if s := c.ConnOptions; s != nil {
+		if s.MaxOpen > 0 {
+			client.DB.SetMaxOpenConns(s.MaxOpen)
+		}
+		if s.MaxIdle > 0 {
+			client.DB.SetMaxIdleConns(s.MaxIdle)
+		}
+		if s.MaxLifetime > 0 {
+			client.DB.SetConnMaxLifetime(s.MaxLifetime)
+		}
+		if s.MaxIdleTime > 0 {
+			client.DB.SetConnMaxIdleTime(s.MaxIdleTime)
+		}
+	}
 	client.AddClosers(c)
 	return client, nil
 }
diff --git a/sql/sqlclient/client.go b/sql/sqlclient/client.go
index ad9b0269900..e4417c94b7f 100644
--- a/sql/sqlclient/client.go
+++ b/sql/sqlclient/client.go
@@ -42,6 +42,12 @@ type (
 		schemahcl.Marshaler
 		schemahcl.Evaluator
 
+		// Ephemeral indicates that the database we connect to is "ephemeral"
+		// (e.g., a temporary running container). This can be set by the driver
+		// that opens the client to signal to its consumers that there is no need
+		// to guard against race conditions with other Atlas clients.
+		Ephemeral bool
+
 		// Functions registered by the drivers and used for opening transactions and their clients.
 		openDriver func(schema.ExecQuerier) (migrate.Driver, error)
 		openTx     TxOpener