From e5ca75b278b881d1914181234a08a7fbbde5b959 Mon Sep 17 00:00:00 2001
From: Mark Raynsford <code@io7m.com>
Date: Mon, 27 May 2024 18:25:53 +0000
Subject: [PATCH] Add put/get for transactions.

---
 .../api/DDatabaseTransactionAbstract.java     | 41 +++++++++++++++++++
 .../darco/api/DDatabaseTransactionType.java   | 25 +++++++++++
 2 files changed, 66 insertions(+)

diff --git a/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionAbstract.java b/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionAbstract.java
index b0c2821..00efe43 100644
--- a/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionAbstract.java
+++ b/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionAbstract.java
@@ -17,11 +17,14 @@
 
 package com.io7m.darco.api;
 
+import com.io7m.jmulticlose.core.CloseableCollection;
+import com.io7m.jmulticlose.core.CloseableCollectionType;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -48,6 +51,8 @@ public abstract class DDatabaseTransactionAbstract<
   private final N connection;
   private final Span transactionSpan;
   private final Map<Class<?>, Q> queries;
+  private final CloseableCollectionType<DDatabaseException> resources;
+  private final HashMap<Class<?>, Object> values;
 
   protected DDatabaseTransactionAbstract(
     final C inConfiguration,
@@ -63,6 +68,42 @@ protected DDatabaseTransactionAbstract(
       Objects.requireNonNull(inTransactionScope, "inMetricsScope");
     this.queries =
       Objects.requireNonNull(inQueries, "queries");
+    this.resources =
+      CloseableCollection.create(() -> {
+        return new DDatabaseException(
+          "One or more resources could not be closed.",
+          "error-resource-close",
+          Map.of(),
+          Optional.empty()
+        );
+      });
+
+    this.values =
+      new HashMap<>();
+  }
+
+  @Override
+  public final <V> void put(
+    final Class<? extends V> clazz,
+    final V value)
+  {
+    if (value instanceof final AutoCloseable closeable) {
+      this.resources.add(closeable);
+    }
+    this.values.put(clazz, value);
+  }
+
+  @Override
+  public final <V> V get(
+    final Class<V> clazz)
+  {
+    return Optional.ofNullable(this.values.get(clazz))
+      .map(clazz::cast)
+      .orElseThrow(() -> {
+        return new IllegalStateException(
+          "No object registered for class %s".formatted(clazz)
+        );
+      });
   }
 
   @Override
diff --git a/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionType.java b/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionType.java
index d622ef7..c46ffba 100644
--- a/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionType.java
+++ b/com.io7m.darco.api/src/main/java/com/io7m/darco/api/DDatabaseTransactionType.java
@@ -85,4 +85,29 @@ void commit()
 
   Span createSubSpan(
     String name);
+
+  /**
+   * Register an object on the transaction.
+   *
+   * @param clazz The class reference used to retrieve the object
+   * @param value The object
+   * @param <V>   The type of object
+   */
+
+  <V> void put(
+    Class<? extends V> clazz,
+    V value);
+
+  /**
+   * Retrieve an object from the transaction.
+   *
+   * @param clazz The class reference used to retrieve the object
+   * @param <V>   The type of object
+   *
+   * @return The object
+   *
+   * @see #put(Class, Object)
+   */
+
+  <V> V get(Class<V> clazz);
 }