From 621f8a9ff20eaa38de08e2bccd7cc4c3f30e1f1d Mon Sep 17 00:00:00 2001
From: Hugo Vieira <hugoamvieira@users.noreply.github.com>
Date: Tue, 6 Sep 2022 15:51:30 +0100
Subject: [PATCH] fix(schema): Database connection not cleanly closed

This commit introduces closing the Schema Manager's underlying driver
connection. We were only releasing the shared lock but not actually
cleaning up the db connection and relied on the fact that the process
would exit to do that for us.
---
 internal/cmd/commands/server/server.go        | 32 +++++++++----------
 .../db/schema/internal/postgres/postgres.go   |  5 +++
 internal/db/schema/manager.go                 | 13 ++++++++
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/internal/cmd/commands/server/server.go b/internal/cmd/commands/server/server.go
index 0605b04e9..ce3c08c48 100644
--- a/internal/cmd/commands/server/server.go
+++ b/internal/cmd/commands/server/server.go
@@ -410,24 +410,22 @@ func (c *Command) Run(args []string) int {
 			c.UI.Error(fmt.Errorf("Can't get schema manager: %w.", err).Error())
 			return base.CommandCliError
 		}
-		if !c.Config.Controller.Database.SkipSharedLockAcquisition {
-			// This is an advisory lock on the DB which is released when the db session ends.
-			if err := sMan.SharedLock(c.Context); err != nil {
-				c.UI.Error(fmt.Errorf("Unable to gain shared access to the database: %w", err).Error())
-				return base.CommandCliError
-			}
-			defer func() {
-				// The base context has already been canceled so we shouldn't use it here.
-				// 1 second is chosen so the shutdown is still responsive and this is a mostly
-				// non critical step since the lock should be released when the session with the
-				// database is closed.
-				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
-				defer cancel()
-				if err := sMan.SharedUnlock(ctx); err != nil {
-					c.UI.Error(fmt.Errorf("Unable to release shared lock to the database: %w", err).Error())
-				}
-			}()
+		// This is an advisory locks on the DB which is released when the db session ends.
+		if err := sMan.SharedLock(c.Context); err != nil {
+			c.UI.Error(fmt.Errorf("Unable to gain shared access to the database: %w", err).Error())
+			return base.CommandCliError
 		}
+		defer func() {
+			// The base context has already been canceled so we shouldn't use it here.
+			// 1 second is chosen so the shutdown is still responsive and this is a mostly
+			// non critical step since the lock should be released when the session with the
+			// database is closed.
+			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+			defer cancel()
+			if err := sMan.Close(ctx); err != nil {
+				c.UI.Error(fmt.Errorf("Unable to release shared lock to the database: %w", err).Error())
+			}
+		}()
 		ckState, err := sMan.CurrentState(c.Context)
 		if err != nil {
 			c.UI.Error(fmt.Errorf("Error checking schema state: %w", err).Error())
diff --git a/internal/db/schema/internal/postgres/postgres.go b/internal/db/schema/internal/postgres/postgres.go
index e0b1c49ad..e0ed1000d 100644
--- a/internal/db/schema/internal/postgres/postgres.go
+++ b/internal/db/schema/internal/postgres/postgres.go
@@ -251,6 +251,11 @@ func (p *Postgres) Run(ctx context.Context, migration io.Reader, version int, ed
 	return nil
 }
 
+// Close closes the underlying Postgres database connection.
+func (p *Postgres) Close() error {
+	return p.conn.Close()
+}
+
 var errOldMigrationTable = stderrors.New("old schema migration table")
 
 func (p *Postgres) schemaInitialized(ctx context.Context) (bool, error) {
diff --git a/internal/db/schema/manager.go b/internal/db/schema/manager.go
index a2d75e17e..95d042540 100644
--- a/internal/db/schema/manager.go
+++ b/internal/db/schema/manager.go
@@ -22,6 +22,7 @@ type driver interface {
 	Lock(context.Context) error
 	Unlock(context.Context) error
 	UnlockShared(context.Context) error
+	Close() error
 	// StartRun begins a transaction internal to the driver.
 	StartRun(context.Context) error
 	// CommitRun commits a transaction, if there is an error it should rollback the transaction.
@@ -123,6 +124,18 @@ func (b *Manager) CurrentState(ctx context.Context) (*State, error) {
 	return &dbS, nil
 }
 
+// Close unlocks the database shared lock and closes the underlying
+// database connection afterwards.
+func (b *Manager) Close(ctx context.Context) error {
+	const op = "schema.(Manager).Close"
+	err := b.SharedUnlock(ctx)
+	if err != nil {
+		return errors.Wrap(ctx, err, op)
+	}
+
+	return b.driver.Close()
+}
+
 // SharedLock attempts to obtain a shared lock on the database.  This can fail
 // if an exclusive lock is already held.  If the lock can't be obtained an
 // error is returned.