Neo4j and Spray JSON

The title of this post may be a bit confusing–what does Neo4j have to do with Spray JSON? (I know that they both run on the JVM, don’t try to be a smart ar*e.) Neo4j is a graph database and to make it useful for our systems, the vertices and edges can have properties. These properties can be Strings, numbers, Dates and such like. So, how are we going to squeeze rich structures in our objects into these (scalar) properties?

The initial reaction could be serialise them; as in, use the Object[Input|Output]Stream. The trouble with Java serialisation is that it is rather temperamental. All it takes is to add a property and boom!, you are having trouble reading the objects back. It would be nice to save these objects in some slightly more manageable format. JSON is a pretty good match, especially when we have so many different JSON serialisation and deserialisation libraries.

If you are in a Spray project, you are probably using Spray JSON. It would be useful to reuse the various JsonFormat[A] typeclass instances to not only deal with JSON on your API level, but also to use the same JSON format in Neo4j. So, let’s get on to it.

Low-level access

Let’s start with low-level Neo4j code. We will need code that gives us access to the underlying GraphDatabaseService, creates nodes, maintains transactions and other plumbing code.

trait GraphDatabase {

  def graphDatabase: GraphDatabaseService

  /**
   * Performs block ``f`` within a transaction
   *
   * @param f the block to be performed
   * @tparam T the type the inner block returns
   * @return ``f``'s result 
   */
  def withTransaction[T](f: => T): T = {
    val tx = graphDatabase.beginTx
    try {
      val result = f
      tx.success()
      result
    } catch {
      case e: Throwable =>
        tx.failure()
        throw e
    } finally {
      tx.finish()
    }
  }

  /**
   * Creates a new and empty node
   *
   * @return the newly created node
   */
  def newNode(): Node = graphDatabase.createNode()

}

So far, no big surprises. The withTransaction function is a variation of the bracket and I will not even insult your intelligence by describing what the newNode() function does. But we’re in Scala and we have these strong types. Let’s see how we can make the most of it.

/**
 * Modifies the given ``Node``s with values in the instances of ``A``
 *
 * @tparam A the A
 */
trait NodeMarshaller[A] {
  def marshal(node: Node)(a: A): Node
}

/**
 * Unmarshals given ``Node``s to create instances of ``A``
 *
 * @tparam A the A
 */
trait NodeUnmarshaller[A] {
  def unmarshal(node: Node): A
}

/**
 * Provides index for the ``A``s
 *
 * @tparam A the A
 */
trait IndexSource[A] {

  def getIndex(graphDatabase: GraphDatabaseService): Index[Node]

}

First, we define the typeclasses that contain functions to marshal value of type A into a Node, unmarshal value of type A from a Node, and finally to provide an Index for some type A. In Scala-speak, these are the ordinary traits in the listing above. Now, let’s use them:

trait TypedGraphDatabase extends GraphDatabase {

  type Identifiable = { def id: UUID }

  import language.reflectiveCalls

  private def createNode[A](a: A)
                           (implicit ma: NodeMarshaller[A]): Node = 
                           ma.marshal(newNode())(a)

  private def find[A](indexOperation: Index[Node] => IndexHits[Node])
                     (implicit is: IndexSource[A], uma: NodeUnmarshaller[A]): 
                     Option[(A, Node)] = {
    val index = is.getIndex(graphDatabase)
    val hits = indexOperation(index)
    val result = if (hits.size() == 1) {
      val node = hits.getSingle
      Some((uma.unmarshal(node), node))
    } else {
      None
    }
    hits.close()

    result
  }

  private def byIdIndexOpertaion(id: UUID): 
    Index[Node] => IndexHits[Node] = index => index.get("id", id.toString)

  def findOne[A <: Identifiable]
      (id: UUID)
      (implicit is: IndexSource[A], uma: NodeUnmarshaller[A]): 
      Option[(A, Node)] =
    find(byIdIndexOpertaion(id))

  def findOneEntity[A <: Identifiable]
      (id: UUID)
      (implicit is: IndexSource[A], uma: NodeUnmarshaller[A]): 
      Option[A] =
    find(byIdIndexOpertaion(id)).map(_._1)

  def findOneEntityWithIndex[A]
      (indexOperation: Index[Node] => IndexHits[Node])
      (implicit is: IndexSource[A], uma: NodeUnmarshaller[A]): 
      Option[A] =
    find(indexOperation).map(_._1)

  def addOne[A <: Identifiable]
      (a: A)
      (implicit is: IndexSource[A], ma: NodeMarshaller[A]): Node = {
    val node = createNode(a)
    is.getIndex(graphDatabase).putIfAbsent(node, "id", a.id.toString)
    node
  }
}

I am only showing the most important functions here, specifically:

• findOne finds one entity and node for the given identity
• findOneEntity as convenience variant of findOne where we do not want the Node
• findOneEntityWithIndex where we want to look up the Node in some custom way
• addOne that adds a new Node and add is it to the index

Spray JSON

But we’re not using Spray JSON yet. All that I’ve given you is some mechanism of marshalling between instances of some type A and Neo4j Nodes. Let’s complete the picture by having the typeclass instances:

trait SprayJsonNodeMarshalling {

  implicit def sprayJsonNodeMarshaller[A : JsonFormat] = 
    new SprayJsonStringNodeMarshaller[A]

  implicit def sprayJsonNodeUnmarshaller[A : JsonFormat] = 
    new SprayJsonStringNodeUnmarshaller[A]

  class SprayJsonStringNodeMarshaller[A : JsonFormat] extends 
    NodeMarshaller[A] {

    def marshal(node: Node)(a: A) = {
      val formatter = implicitly[JsonFormat[A]]
      val json = formatter.write(a).compactPrint
      node.setProperty("json", json)
      node
    }

  }

  class SprayJsonStringNodeUnmarshaller[A : JsonFormat] extends 
    NodeUnmarshaller[A] {

    def unmarshal(node: Node) = {
      val json = node.getProperty("json").toString
      val parsed = JsonParser.apply(json)
      val formatter = implicitly[JsonFormat[A]]
      formatter.read(parsed)
    }
  }

}

Now, this is better. If we mix in the TypedGraphDatabase with SprayJsonNodeMarshalling and have instances of JsonFormat for the types we are saving, we’re in business!

Verba docent, exempla tranunt, so let’s see some sample code. We will be saving some Customer instances. Let’s start with the definition of the Customer data type.

case class Customer(
  id: UUID, firstName: String, lastName: String, 
  dateOfBirth: Date, theNameOfTheHospitalInWhichHisMotherWasBorn: String)

If we now wish to use it with Spray JSON, we need JsonFormat instance for the type Customer, and because we will be using this instance in both the API and the Neo4j access, we will put it in a trait that we can mix in wherever we need it.

trait CustomerFormats extends DefaultJsonProtocol with 
  UuidFormats with DateFormats {
  implicit val CustomerFormat = jsonFormat5(Customer)
}

Those pesky UuidFormats and DateFormats traits define the JsonFormat instances for types UUID and Date. Now, we’re nearly ready to start manipulating the Customer Neo4j store. The last thing we need is to define which Index the Customer-carrying nodes will live in. This is the job of the IndexSource. It is slightly more involved

We need to get the underlying GraphDatabaseService to create (or get) the Index.

trait CustomerGraphDatabaseIndexes {
  this: GraphDatabase =>

  lazy val customerIndex = graphDatabase.index().forNodes("customers")

  implicit object CustomerIndexSource extends IndexSource[Customer] {
    def getIndex(graphDatabase: GraphDatabaseService) = customerIndex
  }

}

Now we have all the components ready: we can mix in our sample application

object Demo extends 
  Application with
  TypedGraphDatabase with 
  SprayJsonNodeMarshalling with 
  CustomerGraphDatabaseIndexes with
  CustomerFormats {

  // satisfy GraphDatabase.graphDatabase
  val graphDatabase = new GraphDatabaseFactory().newEmbeddedDatabase("path")

  val customer = Customer(...)

  withTransaction {
    addOne(customer)
  }

  val found = findOneEntity[Customer](customer.id)
}

The code is now actually quite pleasant. We have poor man’s ORM on top of graph database. We can persist complex instances in properties of our nodes.

If you like this post, keep an eye on https://github.com/janm399/scalad; the Neo4j functionality is coming there in the next few days.