package io.mediachain.util

import cats.data.Xor
import gremlin.scala._
import io.mediachain.GraphError
import io.mediachain.GraphError.TransactionFailed

object GremlinUtils {

  def withTransaction[T](graph: Graph)(f: => T): Xor[GraphError, T] = {
    if (!graph.tx.isOpen) graph.tx.open()

    val result: Xor[GraphError, T] =
      Xor.catchNonFatal(f).leftMap(TransactionFailed)

    result match {
      case Xor.Left(_) => graph.tx.rollback()
      case _ => graph.tx.commit()
    }

    result
  }


  def withTransactionXor[T](graph: Graph)(f: => Xor[GraphError, T])
  : Xor[GraphError, T] = {
    if (!graph.tx.isOpen) graph.tx.open()

    val result: Xor[GraphError, T] =
      Xor.catchNonFatal(f)
        .leftMap(TransactionFailed)
        .flatMap(res => res)

    result match {
      case Xor.Left(er) => graph.tx.rollback()
      case _ => graph.tx.commit()
    }

    result
  }
}

package io.mediachain.util

import cats.data.Xor
import gremlin.scala._
import io.mediachain.GraphError
import io.mediachain.GraphError.TransactionFailed

object GremlinUtils {

  def withTransaction[T](graph: Graph)(f: => T): Xor[GraphError, T] = {
    if (!graph.tx.isOpen) {
      graph.tx.open()
    }
    Xor.catchNonFatal {
      f
    }.leftMap(ex => {
      graph.tx.rollback()
      TransactionFailed(ex)
    })
      .map(result => {
        graph.tx.commit()
        result
      })
  }

}

--- 
+++ 
@@ -8,19 +8,34 @@
 object GremlinUtils {
 
   def withTransaction[T](graph: Graph)(f: => T): Xor[GraphError, T] = {
-    if (!graph.tx.isOpen) {
-      graph.tx.open()
+    if (!graph.tx.isOpen) graph.tx.open()
+
+    val result: Xor[GraphError, T] =
+      Xor.catchNonFatal(f).leftMap(TransactionFailed)
+
+    result match {
+      case Xor.Left(_) => graph.tx.rollback()
+      case _ => graph.tx.commit()
     }
-    Xor.catchNonFatal {
-      f
-    }.leftMap(ex => {
-      graph.tx.rollback()
-      TransactionFailed(ex)
-    })
-      .map(result => {
-        graph.tx.commit()
-        result
-      })
+
+    result
   }
 
+
+  def withTransactionXor[T](graph: Graph)(f: => Xor[GraphError, T])
+  : Xor[GraphError, T] = {
+    if (!graph.tx.isOpen) graph.tx.open()
+
+    val result: Xor[GraphError, T] =
+      Xor.catchNonFatal(f)
+        .leftMap(TransactionFailed)
+        .flatMap(res => res)
+
+    result match {
+      case Xor.Left(er) => graph.tx.rollback()
+      case _ => graph.tx.commit()
+    }
+
+    result
+  }
 }
Add withTransactionXor helper to rollback on GraphError
