From 1a999b7191016bbc21806884a5c1b185e6169e2f Mon Sep 17 00:00:00 2001
From: Jonathan Pearlin <jpearlin1@gmail.com>
Date: Wed, 11 May 2022 16:12:49 -0400
Subject: [PATCH] Close underlying connections during migration (#12710)

---
 .../db/instance/FlywayMigrationDatabase.java  | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/FlywayMigrationDatabase.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/FlywayMigrationDatabase.java
index 11d529eff252..94e23762baea 100644
--- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/FlywayMigrationDatabase.java
+++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/FlywayMigrationDatabase.java
@@ -8,6 +8,7 @@
 import io.airbyte.db.factory.DSLContextFactory;
 import io.airbyte.db.factory.DataSourceFactory;
 import io.airbyte.db.factory.FlywayFactory;
+import java.io.Closeable;
 import java.io.IOException;
 import java.sql.Connection;
 import javax.sql.DataSource;
@@ -18,6 +19,8 @@
 import org.jooq.meta.postgres.PostgresDatabase;
 import org.jooq.tools.StringUtils;
 import org.jooq.tools.jdbc.JDBCUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.testcontainers.containers.PostgreSQLContainer;
 
 /**
@@ -33,10 +36,16 @@
  */
 public abstract class FlywayMigrationDatabase extends PostgresDatabase {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(FlywayMigrationDatabase.class);
+
   private static final String DEFAULT_DOCKER_IMAGE = "postgres:13-alpine";
 
   private Connection connection;
 
+  private DataSource dataSource;
+
+  private DSLContext dslContext;
+
   protected abstract Database getAndInitializeDatabase(DSLContext dslContext) throws IOException;
 
   protected abstract DatabaseMigrator getDatabaseMigrator(Database database, Flyway flyway);
@@ -75,9 +84,9 @@ private void createInternalConnection() throws Exception {
         .withPassword("jooq_generator");
     container.start();
 
-    final DataSource dataSource =
+    dataSource =
         DataSourceFactory.create(container.getUsername(), container.getPassword(), container.getDriverClassName(), container.getJdbcUrl());
-    final DSLContext dslContext = DSLContextFactory.create(dataSource, SQLDialect.POSTGRES);
+    dslContext = DSLContextFactory.create(dataSource, SQLDialect.POSTGRES);
     final Flyway flyway = FlywayFactory.create(dataSource, getInstalledBy(), getDbIdentifier(), getMigrationFileLocations());
     final Database database = getAndInitializeDatabase(dslContext);
     final DatabaseMigrator migrator = getDatabaseMigrator(database, flyway);
@@ -91,6 +100,14 @@ private void createInternalConnection() throws Exception {
   public void close() {
     JDBCUtils.safeClose(connection);
     connection = null;
+    dslContext.close();
+    if (dataSource instanceof Closeable closeable) {
+      try {
+        closeable.close();
+      } catch (final IOException e) {
+        LOGGER.warn("Unable to close data source.", e);
+      }
+    }
     super.close();
   }