Fix connection pool exhaustion when inserting (#823)

The getHandle method was obtaining a connection from the pool and not freeing it

Signed-off-by: henneberger <[email protected]>

Author: henneberger

src/main/java/marquez/db/DatasetDao.java

@@ -35,30 +35,32 @@
 public interface DatasetDao extends SqlObject {
   @Transaction
   default void insert(DatasetRow row) {
-    getHandle()
-        .createUpdate(
-            "INSERT INTO datasets ("
-                + "uuid, "
-                + "type, "
-                + "created_at, "
-                + "updated_at, "
-                + "namespace_uuid, "
-                + "source_uuid, "
-                + "name, "
-                + "physical_name, "
-                + "description"
-                + ") VALUES ("
-                + ":uuid, "
-                + ":type, "
-                + ":createdAt, "
-                + ":updatedAt, "
-                + ":namespaceUuid, "
-                + ":sourceUuid, "
-                + ":name, "
-                + ":physicalName, "
-                + ":description)")
-        .bindBean(row)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "INSERT INTO datasets ("
+                        + "uuid, "
+                        + "type, "
+                        + "created_at, "
+                        + "updated_at, "
+                        + "namespace_uuid, "
+                        + "source_uuid, "
+                        + "name, "
+                        + "physical_name, "
+                        + "description"
+                        + ") VALUES ("
+                        + ":uuid, "
+                        + ":type, "
+                        + ":createdAt, "
+                        + ":updatedAt, "
+                        + ":namespaceUuid, "
+                        + ":sourceUuid, "
+                        + ":name, "
+                        + ":physicalName, "
+                        + ":description)")
+                .bindBean(row)
+                .execute());
     // Tags
     final Instant taggedAt = row.getCreatedAt();
     row.getTagUuids().forEach(tagUuid -> updateTags(row.getUuid(), tagUuid, taggedAt));

src/main/java/marquez/db/DatasetFieldDao.java

@@ -38,26 +38,28 @@ default void insertAll(List<DatasetFieldRow> rows) {
 
   @Transaction
   default void insert(DatasetFieldRow row) {
-    getHandle()
-        .createUpdate(
-            "INSERT INTO dataset_fields ("
-                + "uuid, "
-                + "type, "
-                + "created_at, "
-                + "updated_at, "
-                + "dataset_uuid, "
-                + "name, "
-                + "description"
-                + ") VALUES ("
-                + ":uuid, "
-                + ":type, "
-                + ":createdAt, "
-                + ":updatedAt, "
-                + ":datasetUuid, "
-                + ":name, "
-                + ":description)")
-        .bindBean(row)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "INSERT INTO dataset_fields ("
+                        + "uuid, "
+                        + "type, "
+                        + "created_at, "
+                        + "updated_at, "
+                        + "dataset_uuid, "
+                        + "name, "
+                        + "description"
+                        + ") VALUES ("
+                        + ":uuid, "
+                        + ":type, "
+                        + ":createdAt, "
+                        + ":updatedAt, "
+                        + ":datasetUuid, "
+                        + ":name, "
+                        + ":description)")
+                .bindBean(row)
+                .execute());
     // Tags
     final Instant taggedAt = row.getCreatedAt();
     row.getTagUuids().forEach(tagUuid -> updateTags(row.getUuid(), tagUuid, taggedAt));

src/main/java/marquez/db/JobVersionDao.java

@@ -41,28 +41,30 @@ enum IoType {
 
   @Transaction
   default void insert(@NonNull JobVersionRow row) {
-    getHandle()
-        .createUpdate(
-            "INSERT INTO job_versions ("
-                + "uuid, "
-                + "created_at, "
-                + "updated_at, "
-                + "job_uuid, "
-                + "version, "
-                + "location, "
-                + "latest_run_uuid, "
-                + "job_context_uuid"
-                + ") VALUES ("
-                + ":uuid, "
-                + ":createdAt, "
-                + ":updateAt, "
-                + ":jobUuid, "
-                + ":version, "
-                + ":location, "
-                + ":latestRunUuid, "
-                + ":jobContextUuid)")
-        .bindBean(row)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "INSERT INTO job_versions ("
+                        + "uuid, "
+                        + "created_at, "
+                        + "updated_at, "
+                        + "job_uuid, "
+                        + "version, "
+                        + "location, "
+                        + "latest_run_uuid, "
+                        + "job_context_uuid"
+                        + ") VALUES ("
+                        + ":uuid, "
+                        + ":createdAt, "
+                        + ":updateAt, "
+                        + ":jobUuid, "
+                        + ":version, "
+                        + ":location, "
+                        + ":latestRunUuid, "
+                        + ":jobContextUuid)")
+                .bindBean(row)
+                .execute());
     // I/O
     row.getInputUuids().forEach(inputUuid -> updateInputs(row.getUuid(), inputUuid));
     row.getOutputUuids().forEach(outputUuid -> updateOutputs(row.getUuid(), outputUuid));

src/main/java/marquez/db/RunDao.java

@@ -42,26 +42,28 @@ default void insertWith(RunRow row, List<UUID> inputVersionUuids) {
 
   @Transaction
   default void insert(RunRow row) {
-    getHandle()
-        .createUpdate(
-            "INSERT INTO runs ("
-                + "uuid, "
-                + "created_at, "
-                + "updated_at, "
-                + "job_version_uuid, "
-                + "run_args_uuid, "
-                + "nominal_start_time, "
-                + "nominal_end_time"
-                + ") VALUES ("
-                + ":uuid, "
-                + ":createdAt, "
-                + ":updatedAt, "
-                + ":jobVersionUuid, "
-                + ":runArgsUuid, "
-                + ":nominalStartTime, "
-                + ":nominalEndTime)")
-        .bindBean(row)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "INSERT INTO runs ("
+                        + "uuid, "
+                        + "created_at, "
+                        + "updated_at, "
+                        + "job_version_uuid, "
+                        + "run_args_uuid, "
+                        + "nominal_start_time, "
+                        + "nominal_end_time"
+                        + ") VALUES ("
+                        + ":uuid, "
+                        + ":createdAt, "
+                        + ":updatedAt, "
+                        + ":jobVersionUuid, "
+                        + ":runArgsUuid, "
+                        + ":nominalStartTime, "
+                        + ":nominalEndTime)")
+                .bindBean(row)
+                .execute());
     // Input versions
     row.getInputVersionUuids()
         .forEach(inputVersionUuid -> updateInputVersions(row.getUuid(), inputVersionUuid));
@@ -72,16 +74,18 @@ default void insert(RunRow row) {
 
   @Transaction
   default void updateJobVersionUuid(UUID rowUuid, Instant updatedAt, UUID jobVersionUuid) {
-    getHandle()
-        .createUpdate(
-            "UPDATE runs "
-                + "SET updated_at = :updatedAt, "
-                + "    job_version_uuid = :jobVersionUuid "
-                + "WHERE uuid = :rowUuid")
-        .bind("updatedAt", updatedAt)
-        .bind("jobVersionUuid", jobVersionUuid)
-        .bind("rowUuid", rowUuid)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "UPDATE runs "
+                        + "SET updated_at = :updatedAt, "
+                        + "    job_version_uuid = :jobVersionUuid "
+                        + "WHERE uuid = :rowUuid")
+                .bind("updatedAt", updatedAt)
+                .bind("jobVersionUuid", jobVersionUuid)
+                .bind("rowUuid", rowUuid)
+                .execute());
     createJobVersionDao().updateLatestRun(jobVersionUuid, updatedAt, rowUuid);
   }
 

src/main/java/marquez/db/RunStateDao.java

@@ -41,12 +41,14 @@ default void insert(
       boolean starting,
       boolean done,
       boolean complete) {
-    getHandle()
-        .createUpdate(
-            "INSERT INTO run_states (uuid, transitioned_at, run_uuid, state)"
-                + "VALUES (:uuid, :transitionedAt, :runUuid, :state)")
-        .bindBean(row)
-        .execute();
+    withHandle(
+        handle ->
+            handle
+                .createUpdate(
+                    "INSERT INTO run_states (uuid, transitioned_at, run_uuid, state)"
+                        + "VALUES (:uuid, :transitionedAt, :runUuid, :state)")
+                .bindBean(row)
+                .execute());
     // State transition
     final Instant updateAt = row.getTransitionedAt();
     createRunDao().updateRunState(row.getRunUuid(), updateAt, row.getState());