From a83ea351f4ace30d8408f6decd9a802277b51a08 Mon Sep 17 00:00:00 2001
From: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>
Date: Wed, 5 Feb 2025 17:44:30 -0800
Subject: [PATCH] internal: ensure schema manager is closed

This releases the connection it holds indefinitely
---
 internal/cmd/commands/database/funcs.go      | 2 ++
 internal/db/schema/migrations/oss/testing.go | 1 +
 internal/db/schema/schema.go                 | 1 +
 internal/oplog/testing.go                    | 1 +
 4 files changed, 5 insertions(+)

diff --git a/internal/cmd/commands/database/funcs.go b/internal/cmd/commands/database/funcs.go
index 2e667982fb..0ab80316a7 100644
--- a/internal/cmd/commands/database/funcs.go
+++ b/internal/cmd/commands/database/funcs.go
@@ -46,11 +46,13 @@ func migrateDatabase(ctx context.Context, ui cli.Ui, dialect, u string, initiali
 	// This is an advisory lock on the DB which is released when the DB session ends.
 	if err := man.ExclusiveLock(ctx); err != nil {
 		ui.Error("Unable to capture a lock on the database.")
+		_ = man.Close(ctx)
 		return noop, 2
 	}
 	unlock := func() {
 		// We don't report anything since this should resolve itself anyways.
 		_ = man.ExclusiveUnlock(ctx)
+		_ = man.Close(ctx)
 	}
 
 	st, err := man.CurrentState(ctx)
diff --git a/internal/db/schema/migrations/oss/testing.go b/internal/db/schema/migrations/oss/testing.go
index 3c17567ab2..c7b5fc3645 100644
--- a/internal/db/schema/migrations/oss/testing.go
+++ b/internal/db/schema/migrations/oss/testing.go
@@ -19,6 +19,7 @@ func ApplyMigration(t *testing.T, ctx context.Context, d *sql.DB, migrationId in
 		schema.TestCreatePartialEditions(schema.Dialect(dialect), schema.PartialEditions{"oss": migrationId}),
 	))
 	require.NoError(t, err)
+	t.Cleanup(func() { m.Close(context.Background()) })
 	_, err = m.ApplyMigrations(ctx)
 	require.NoError(t, err)
 	state, err := m.CurrentState(ctx)
diff --git a/internal/db/schema/schema.go b/internal/db/schema/schema.go
index bf12da1e54..039991762c 100644
--- a/internal/db/schema/schema.go
+++ b/internal/db/schema/schema.go
@@ -25,6 +25,7 @@ func MigrateStore(ctx context.Context, dialect Dialect, url string, opt ...Optio
 	if err != nil {
 		return false, errors.Wrap(ctx, err, op)
 	}
+	defer sMan.Close(ctx)
 
 	st, err := sMan.CurrentState(ctx)
 	if err != nil {
diff --git a/internal/oplog/testing.go b/internal/oplog/testing.go
index 18509cafbc..f5f5ba5845 100644
--- a/internal/oplog/testing.go
+++ b/internal/oplog/testing.go
@@ -125,6 +125,7 @@ func testInitStore(t testing.TB, cleanup func() error, url string) {
 	require.NoError(t, err)
 	sm, err := schema.NewManager(ctx, schema.Dialect(dialect), d)
 	require.NoError(t, err)
+	t.Cleanup(func() { sm.Close(context.Background()) })
 	_, err = sm.ApplyMigrations(ctx)
 	require.NoError(t, err)
 }