From efca86783cb1cb01f5bed74aa375929fb33eae30 Mon Sep 17 00:00:00 2001
From: Greg Burd <gburd@onshape.com>
Date: Wed, 13 Sep 2017 12:28:12 -0400
Subject: [PATCH] Allow configuration of the exception class used when
 conflicts arise at UnitOfWork commit time to enable clients to provide
 custom/advanced/specific exceptions without having to do a bunch of
 encapsulation.

---
 pom.xml                                                | 2 +-
 src/main/java/net/helenus/core/HelenusSession.java     | 2 ++
 src/main/java/net/helenus/core/SessionInitializer.java | 7 +++++++
 src/main/java/net/helenus/core/UnitOfWork.java         | 7 ++++++-
 4 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/pom.xml b/pom.xml
index ad11ab3..7c21e2b 100644
--- a/pom.xml
+++ b/pom.xml
@@ -5,7 +5,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>net.helenus</groupId>
     <artifactId>helenus-core</artifactId>
-    <version>2.0.30-SNAPSHOT</version>
+    <version>2.0.31-SNAPSHOT</version>
     <packaging>jar</packaging>
 
     <name>helenus</name>
diff --git a/src/main/java/net/helenus/core/HelenusSession.java b/src/main/java/net/helenus/core/HelenusSession.java
index 67101d1..60c8c77 100644
--- a/src/main/java/net/helenus/core/HelenusSession.java
+++ b/src/main/java/net/helenus/core/HelenusSession.java
@@ -75,6 +75,7 @@ public final class HelenusSession extends AbstractSessionOperations implements C
           Executor executor,
           boolean dropSchemaOnClose,
           ConsistencyLevel consistencyLevel,
+          Class<? extends Exception> conflictExceptionClass,
           MetricRegistry metricRegistry,
           Tracer tracer) {
     this.session = session;
@@ -88,6 +89,7 @@ public final class HelenusSession extends AbstractSessionOperations implements C
     this.executor = executor;
     this.dropSchemaOnClose = dropSchemaOnClose;
     this.defaultConsistencyLevel = consistencyLevel;
+    UnitOfWork.conflictExceptionClass = conflictExceptionClass;
     this.metricRegistry = metricRegistry;
     this.zipkinTracer = tracer;
 
diff --git a/src/main/java/net/helenus/core/SessionInitializer.java b/src/main/java/net/helenus/core/SessionInitializer.java
index a46253a..0a58e7a 100644
--- a/src/main/java/net/helenus/core/SessionInitializer.java
+++ b/src/main/java/net/helenus/core/SessionInitializer.java
@@ -43,6 +43,7 @@ public final class SessionInitializer extends AbstractSessionOperations {
   private Tracer zipkinTracer;
   private PrintStream printStream = System.out;
   private Executor executor = MoreExecutors.directExecutor();
+  private Class<? extends Exception> conflictingUnitOfWorkClass = ConflictingUnitOfWorkException.class;
 
   private SessionRepositoryBuilder sessionRepository;
 
@@ -110,6 +111,11 @@ public final class SessionInitializer extends AbstractSessionOperations {
     return this;
   }
 
+  public SessionInitializer setConflictingUnitOfWorkException(Class<? extends Exception> e) {
+    this.conflictingUnitOfWorkClass = e;
+    return this;
+  }
+
   public SessionInitializer consistencyLevel(ConsistencyLevel consistencyLevel) {
     this.consistencyLevel = consistencyLevel;
     return this;
@@ -235,6 +241,7 @@ public final class SessionInitializer extends AbstractSessionOperations {
         executor,
         autoDdl == AutoDdl.CREATE_DROP,
         consistencyLevel,
+        conflictingUnitOfWorkClass,
         metricRegistry,
         zipkinTracer);
   }
diff --git a/src/main/java/net/helenus/core/UnitOfWork.java b/src/main/java/net/helenus/core/UnitOfWork.java
index ab017f0..bba8b85 100644
--- a/src/main/java/net/helenus/core/UnitOfWork.java
+++ b/src/main/java/net/helenus/core/UnitOfWork.java
@@ -9,6 +9,7 @@ import java.util.*;
 
 /** Encapsulates the concept of a "transaction" as a unit-of-work. */
 public class UnitOfWork implements AutoCloseable {
+  protected static Class<? extends Exception>conflictExceptionClass = ConflictingUnitOfWorkException.class;
   private final List<UnitOfWork> nested = new ArrayList<>();
   private final HelenusSession session;
   private final UnitOfWork parent;
@@ -74,7 +75,7 @@ public class UnitOfWork implements AutoCloseable {
    * @return a function from which to chain work that only happens when commit is successful
    * @throws ConflictingUnitOfWorkException when the work overlaps with other concurrent writers.
    */
-  public PostCommitFunction<Void, Void> commit() throws ConflictingUnitOfWorkException {
+  public PostCommitFunction<Void, Void> commit() throws Exception {
     // All nested UnitOfWork should be committed (not aborted) before calls to commit, check.
     boolean canCommit = true;
     TreeTraverser<UnitOfWork> traverser = TreeTraverser.using(node -> node::getChildNodes);
@@ -120,6 +121,10 @@ public class UnitOfWork implements AutoCloseable {
         return new PostCommitFunction(this, null);
       }
     }
+    // else {
+    // Constructor<?> ctor = clazz.getConstructor(conflictExceptionClass);
+    // Object object = ctor.newInstance(new Object[] { String message });
+    // }
     return new PostCommitFunction(this, postCommit);
   }