Remove database connections leaks (1 prod, many in tests)

cloudfoundry · Dec 11, 2024 · d5854bd · d5854bd
1 parent 4519778
commit d5854bd
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 17 deletions.
diff --git a/...src/main/java/org/cloudfoundry/identity/uaa/db/mysql/V4_9_2__AddPrimaryKeysIfMissing.java b/...src/main/java/org/cloudfoundry/identity/uaa/db/mysql/V4_9_2__AddPrimaryKeysIfMissing.java
@@ -4,6 +4,10 @@
 import org.flywaydb.core.api.migration.Context;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.datasource.SingleConnectionDataSource;
+import org.springframework.jdbc.datasource.lookup.DataSourceLookupFailureException;
+
+import java.sql.Connection;
+import java.sql.SQLException;
 
 
 public class V4_9_2__AddPrimaryKeysIfMissing extends BaseJavaMigration {
@@ -16,11 +20,19 @@ public void migrate(Context context) throws Exception {
         JdbcTemplate jdbcTemplate = new JdbcTemplate(new SingleConnectionDataSource(
                 context.getConnection(), true));
         for (String table : tables) {
-            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), table);
+            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(jdbcTemplate), table);
             if (count == 0) {
                 String sql = "ALTER TABLE " + table + " ADD COLUMN `id` int(11) unsigned PRIMARY KEY AUTO_INCREMENT";
                 jdbcTemplate.execute(sql);
             }
         }
     }
+
+    private String getDatabaseCatalog(JdbcTemplate jdbcTemplate) {
+        try (Connection connection = jdbcTemplate.getDataSource().getConnection()) {
+            return  connection.getCatalog();
+        } catch (SQLException e) {
+            throw new DataSourceLookupFailureException("Unable to look up database schema.", e);
+        }
+    }
 }
diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/db/DbMigrationIntegrationTestParent.java b/server/src/test/java/org/cloudfoundry/identity/uaa/db/DbMigrationIntegrationTestParent.java
@@ -45,4 +45,8 @@ public void cleanup() throws SQLException {
             TestUtils.cleanAndSeedDb(jdbcTemplate);
         }
     }
+
+    protected String getDatabaseCatalog() {
+        return MigrationTest.getDatabaseCatalog(jdbcTemplate);
+    }
 }
diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/db/HsqlDbMigrationIntegrationTest.java b/server/src/test/java/org/cloudfoundry/identity/uaa/db/HsqlDbMigrationIntegrationTest.java
@@ -29,16 +29,16 @@ public String getTargetMigration() {
 
             @Override
             public void runAssertions() throws Exception {
-                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "OAUTH_CODE");
+                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "OAUTH_CODE");
                 assertThat("OAUTH_CODE is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "GROUP_MEMBERSHIP");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "GROUP_MEMBERSHIP");
                 assertThat("GROUP_MEMBERSHIP is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "SEC_AUDIT");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "SEC_AUDIT");
                 assertThat("SEC_AUDIT is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "EXTERNAL_GROUP_MAPPING");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "EXTERNAL_GROUP_MAPPING");
                 assertThat("EXTERNAL_GROUP_MAPPING is missing primary key", count, is(1));
 
                 try {
@@ -56,10 +56,10 @@ public void runAssertions() throws Exception {
     public void everyTableShouldHaveAPrimaryKeyColumn() throws Exception {
         flyway.migrate();
 
-        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, jdbcTemplate.getDataSource().getConnection().getCatalog());
+        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, getDatabaseCatalog());
         assertThat(tableNames, hasSize(greaterThan(0)));
         for (String tableName : tableNames) {
-            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), tableName);
+            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), tableName);
             assertThat("%s is missing primary key".formatted(tableName), count, greaterThanOrEqualTo(1));
         }
     }

diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/db/MigrationTest.java b/server/src/test/java/org/cloudfoundry/identity/uaa/db/MigrationTest.java
@@ -1,7 +1,21 @@
 package org.cloudfoundry.identity.uaa.db;
 
+import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.jdbc.datasource.lookup.DataSourceLookupFailureException;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+
 public interface MigrationTest {
     String getTargetMigration();
 
     void runAssertions() throws Exception;
+
+    static String getDatabaseCatalog(JdbcTemplate jdbcTemplate) {
+        try (Connection connection = jdbcTemplate.getDataSource().getConnection()) {
+            return  connection.getCatalog();
+        } catch (SQLException e) {
+            throw new DataSourceLookupFailureException("Unable to look up database schema.", e);
+        }
+    }
 }
diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/db/MySqlDbMigrationIntegrationTest.java b/server/src/test/java/org/cloudfoundry/identity/uaa/db/MySqlDbMigrationIntegrationTest.java
@@ -30,16 +30,16 @@ public String getTargetMigration() {
 
             @Override
             public void runAssertions() throws Exception {
-                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "oauth_code");
+                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "oauth_code");
                 assertThat("oauth_code is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "group_membership");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "group_membership");
                 assertThat("group_membership is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "sec_audit");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "sec_audit");
                 assertThat("sec_audit is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "external_group_mapping");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "external_group_mapping");
                 assertThat("external_group_membership is missing primary key", count, is(1));
 
                 try {
@@ -79,10 +79,10 @@ public String getTargetMigration() {
 
             @Override
             public void runAssertions() throws Exception {
-                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "group_membership");
+                int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "group_membership");
                 assertThat("group_membership is missing primary key", count, is(1));
 
-                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), "external_group_mapping");
+                count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), "external_group_mapping");
                 assertThat("external_group_mapping is missing primary key", count, is(1));
             }
         });
@@ -94,10 +94,10 @@ public void runAssertions() throws Exception {
     public void everyTableShouldHaveAPrimaryKeyColumn() throws Exception {
         flyway.migrate();
 
-        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, jdbcTemplate.getDataSource().getConnection().getCatalog());
+        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, getDatabaseCatalog());
         assertThat(tableNames, hasSize(greaterThan(0)));
         for (String tableName : tableNames) {
-            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), tableName);
+            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), tableName);
             assertThat("%s is missing primary key".formatted(tableName), count, greaterThanOrEqualTo(1));
         }
     }

diff --git a/...er/src/test/java/org/cloudfoundry/identity/uaa/db/PostgresDbMigrationIntegrationTest.java b/...er/src/test/java/org/cloudfoundry/identity/uaa/db/PostgresDbMigrationIntegrationTest.java
@@ -22,10 +22,10 @@ public class PostgresDbMigrationIntegrationTest extends DbMigrationIntegrationTe
     public void everyTableShouldHaveAPrimaryKeyColumn() throws Exception {
         flyway.migrate();
 
-        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, jdbcTemplate.getDataSource().getConnection().getCatalog());
+        List<String> tableNames = jdbcTemplate.queryForList(getAllTableNames, String.class, getDatabaseCatalog());
         assertThat(tableNames, hasSize(greaterThan(0)));
         for (String tableName : tableNames) {
-            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, jdbcTemplate.getDataSource().getConnection().getCatalog(), tableName, "%" + tableName + "_pk%");
+            int count = jdbcTemplate.queryForObject(checkPrimaryKeyExists, Integer.class, getDatabaseCatalog(), tableName, "%" + tableName + "_pk%");
             assertThat("%s is missing primary key".formatted(tableName), count, greaterThanOrEqualTo(1));
         }