arangodb#80 allow multiple transaction calls unless there are no additional collections allowed

Author: aburmeis

src/main/java/com/arangodb/springframework/repository/query/AbstractArangoQuery.java

@@ -78,7 +78,7 @@ public Object execute(final Object[] parameters) {
 
 		final Pair<String, ? extends Collection<String>> queryAndCollection = createQuery(accessor, bindVars);
 		if (options.getStreamTransactionId() == null && transactionBridge != null) {
-			options.streamTransactionId(transactionBridge.beginCurrentTransaction(queryAndCollection.getSecond()));
+			options.streamTransactionId(transactionBridge.getCurrentTransaction(queryAndCollection.getSecond()));
 		}
 
 

src/main/java/com/arangodb/springframework/repository/query/QueryTransactionBridge.java

@@ -10,21 +10,21 @@
  */
 public class QueryTransactionBridge {
     private static final Function<Collection<String>, String> NO_TRANSACTION = any -> null;
-    private static final ThreadLocal<Function<Collection<String>, String>> CURRENT_TRANSACTION_BEGIN = new NamedInheritableThreadLocal<>("ArangoTransactionBegin");
+    private static final ThreadLocal<Function<Collection<String>, String>> CURRENT_TRANSACTION = new NamedInheritableThreadLocal<>("ArangoTransactionBegin");
 
     public QueryTransactionBridge() {
-        CURRENT_TRANSACTION_BEGIN.set(NO_TRANSACTION);
+        CURRENT_TRANSACTION.set(NO_TRANSACTION);
     }
 
-    public void setCurrentTransactionBegin(Function<Collection<String>, String> begin) {
-        CURRENT_TRANSACTION_BEGIN.set(begin);
+    public void setCurrentTransaction(Function<Collection<String>, String> begin) {
+        CURRENT_TRANSACTION.set(begin);
     }
 
-    public void clearCurrentTransactionBegin() {
-        CURRENT_TRANSACTION_BEGIN.set(NO_TRANSACTION);
+    public void clearCurrentTransaction() {
+        CURRENT_TRANSACTION.set(NO_TRANSACTION);
     }
 
-    public String beginCurrentTransaction(Collection<String> collections) {
-        return CURRENT_TRANSACTION_BEGIN.get().apply(collections);
+    public String getCurrentTransaction(Collection<String> collections) {
+        return CURRENT_TRANSACTION.get().apply(collections);
     }
 }

src/main/java/com/arangodb/springframework/transaction/ArangoTransaction.java

@@ -16,6 +16,8 @@
 class ArangoTransaction implements SmartTransactionObject {
 
     private final ArangoDatabase database;
+
+    private final Set<String> writeCollections = new HashSet<>();
     private TransactionDefinition definition;
     private StreamTransactionEntity transaction;
 
@@ -29,19 +31,22 @@ boolean exists() {
 
     void configure(TransactionDefinition definition) {
         this.definition = definition;
+        if (definition instanceof TransactionAttribute) {
+            writeCollections.addAll(((TransactionAttribute) definition).getLabels());
+        }
     }
 
-    String begin(Collection<String> collections) {
+    String getOrBegin(Collection<String> collections) {
         if (transaction != null) {
-            throw new IllegalTransactionStateException("Stream transaction already started");
-        }
-        Set<String> allCollections = new HashSet<>(collections);
-        if (definition instanceof TransactionAttribute) {
-            allCollections.addAll(((TransactionAttribute) definition).getLabels());
+            if (!writeCollections.containsAll(collections)) {
+                throw new IllegalTransactionStateException("Stream transaction already started, no additional collections allowed");
+            }
+            return transaction.getId();
         }
+        writeCollections.addAll(collections);
         StreamTransactionOptions options = new StreamTransactionOptions()
                 .allowImplicit(true)
-                .writeCollections(allCollections.toArray(new String[0]))
+                .writeCollections(writeCollections.toArray(new String[0]))
                 .lockTimeout(definition.getTimeout() == -1 ? 0 : definition.getTimeout());
         transaction = database.beginStreamTransaction(options);
         return transaction.getId();

src/main/java/com/arangodb/springframework/transaction/ArangoTransactionManager.java

@@ -43,8 +43,8 @@ protected void doBegin(Object transaction, TransactionDefinition definition) thr
         }
         ArangoTransaction tx = (ArangoTransaction) transaction;
         tx.configure(definition);
-        Function<Collection<String>, String> begin = tx::begin;
-        bridge.setCurrentTransactionBegin(begin.andThen(id -> {
+        Function<Collection<String>, String> begin = tx::getOrBegin;
+        bridge.setCurrentTransaction(begin.andThen(id -> {
             if (logger.isDebugEnabled()) {
                 logger.debug("Began stream transaction " + id);
             }
@@ -59,7 +59,7 @@ protected void doCommit(DefaultTransactionStatus status) throws TransactionExcep
             logger.debug("Commit stream transaction " + tx);
         }
         tx.commit();
-        bridge.clearCurrentTransactionBegin();
+        bridge.clearCurrentTransaction();
     }
 
     @Override
@@ -69,7 +69,7 @@ protected void doRollback(DefaultTransactionStatus status) throws TransactionExc
             logger.debug("Rollback stream transaction " + tx);
         }
         tx.rollback();
-        bridge.clearCurrentTransactionBegin();
+        bridge.clearCurrentTransaction();
     }
 
     @Override

src/test/java/com/arangodb/springframework/repository/query/QueryTransactionBridgeTest.java

@@ -13,18 +13,18 @@ public class QueryTransactionBridgeTest {
     private QueryTransactionBridge underTest = new QueryTransactionBridge();
 
     @Test
-    public void beginCurrentTransactionInitiallyReturnsNull() {
-        assertThat(underTest.beginCurrentTransaction(Collections.singleton("test")), Matchers.nullValue());
+    public void getCurrentTransactionInitiallyReturnsNull() {
+        assertThat(underTest.getCurrentTransaction(Collections.singleton("test")), Matchers.nullValue());
     }
 
     @Test
-    public void setCurrentTransactionBeginIsAppliedOnBeginCurrentTransaction() {
-        underTest.setCurrentTransactionBegin(collections -> collections.iterator().next());
-        assertThat(underTest.beginCurrentTransaction(Collections.singleton("test")), Matchers.is("test"));
+    public void setCurrentTransactionIsAppliedOnGetCurrentTransaction() {
+        underTest.setCurrentTransaction(collections -> collections.iterator().next());
+        assertThat(underTest.getCurrentTransaction(Collections.singleton("test")), Matchers.is("test"));
     }
 
     @After
     public void cleanup() {
-        underTest.clearCurrentTransactionBegin();
+        underTest.clearCurrentTransaction();
     }
 }

src/test/java/com/arangodb/springframework/transaction/ArangoTransactionManagerTest.java

@@ -67,7 +67,7 @@ public void getTransactionReturnsNewTransactionWithoutStreamTransaction() {
         TransactionStatus transaction = underTest.getTransaction(new DefaultTransactionAttribute());
         assertThat(transaction.isNewTransaction(), is(true));
         verify(driver).db(DATABASE_NAME);
-        verify(bridge).setCurrentTransactionBegin(any());
+        verify(bridge).setCurrentTransaction(any());
         verifyNoInteractions(database);
     }
 
@@ -81,7 +81,7 @@ public void getTransactionReturnsTransactionCreatesStreamTransactionOnBridgeBegi
                 .thenReturn("123");
         when(database.beginStreamTransaction(any()))
                 .thenReturn(streamTransaction);
-        verify(bridge).setCurrentTransactionBegin(beginPassed.capture());
+        verify(bridge).setCurrentTransaction(beginPassed.capture());
         beginPassed.getValue().apply(Arrays.asList("foo", "bar"));
         verify(database).beginStreamTransaction(optionsPassed.capture());
         assertThat(optionsPassed.getValue().getAllowImplicit(), is(true));