Neo4j Schema using py2neo

neo4j_schema.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-


"""
Self-contained script that exemplifies the implementation of a pseudo-schema
for Neo4j using py2neo. Specifically:

1. Defines some helper functions for typechecking and CRUD operations
2. Defines the base clase for the schema. Basically, it is a Node and
   Relationship factory with added typechecks and parameter checks.
3. Extends the base class to implement a basic schema, showcasing the features
4. Creates some nodes and relations using the extended schema. If any of the
   expected parameters/nodes is missing or has the wrong type it will raise
   an exception upon creation.

With this software pattern, we are able to introduce typechecks and parameter
checks to enforce consistency in the creation of nodes and relations.
Updating the schema is a matter of changing very few lines of code in the
extended class. Relevant information is contained purely in the class attributes
and the assert method.

Whereas the system is reasonably flexible and powerful, it has one limitation:
the data types provided in the signatures *must* match the datatypes that are
sent to the Neo4j database. Why is this a problem? For example, if I have an
attribute of type ``date``, Neo4j only accepts the ``neotime.Date`` type. This
means I must provide ``neotime.Date`` in the signature, and the input to the
corresponding method must match. The system is not able to receive a convenient
``(year, month, day)`` tuple and transform it to ``neotime.Date``. So data
conversions must happen before using the schema.

.. note::

  In order to run, ths script requires a Neo4j database listening for driver
  interaction on localhost:7687, and the corresponding password (this can be
  changed at the bottom).
"""


from py2neo import Node, Relationship, Graph


# ##############################################################################
# # CRUD OPERATIONS
# ##############################################################################
def add_nodes_and_relations(graph, nodes=None, relations=None):
    """
    :param py2neo.Graph graph: A handle to the Neo4j database
    :param list nodes: contains the ``py2neo.Node``s to be added
    :param list relations: contains the ``py2neo.Relationship``s to be added
    """
    t = graph.begin()
    if nodes is not None:
        for n in nodes:
            t.create(n)
    if relations is not None:
        for r in relations:
            t.create(r)
    t.commit()


# ##############################################################################
# # TYPECHECK
# ##############################################################################
def assert_types(x, types=None):
    """
    :param types: A list with types to test via ``isinstance`` like str, list...

    Slightly complex assert:
    If types is None, does nothing.
    Otherwise, if x is a subclass of any of the given types, does nothing.
    Otherwise, raises an exception and tells that x is not in [types].
    """
    if types is not None:
        is_type = [isinstance(x, t) for t in types]
        assert any(is_type), f"Input {x} not of type {types}"


def assert_opt_int(x):
    """
    :raises: If x is not of type int or None
    """
    if x is not None:
        assert_types(x, [int])


# ##############################################################################
# # SCHEMA DEFINITION
# ##############################################################################
class Neo4jSchemaBase:
    """
    Base class for defining and executing Neo4j pseudo-schemata.

    :cvar _NODE: Node signatures in the form of a dict
      {k: [(att1, type1), (att2, type2), ...]}, where k is the name of the
      node, and the tuples in the list contain the *compulsory*
      attributes for that node/relation. The attribute name must be given as
      a string. The attribute type can be the type itself (e.g. str, list)
      or a string representing a custom type (e.g. "liststr"). All types will
      be processed by the ``_assert_type`` method, so to create a custom type
      simply extend ``_assert_type`` to handle it adequately.

    :cvar _REL: Relationship signatures, in the form of a list
      [[(ori1, dest1), ...], {k: [(att1, type1), ...]}], where the second
      element is identical to the node signatures. The first element is a list
      of ``origin->destiny`` pairs, designing all the allowed node types for
      that relationship. They must be given as strings, holding the name of
      the class attribute for the corresponding node signatures. If empty, no
      restrictions will be applied.

    See extended classes for usage example.
    """

    @classmethod
    def _assert_type(cls, x, typecode):
        """
        This method can be arbitrarily extended for custom typecodes
        """
        assert_types(x, [typecode])

    @classmethod
    def _make_node(cls, node_sig, **kwargs):
        """
        First checks that the given ``**kwargs`` fulfill the given ``node_sig``,
        and if so creates and returns the corresponding node. Otherwise raises
        an exception.
        """
        node_name, node_signature = next(iter(node_sig.items()))
        # SIGNATURE CHECK
        for sig_name, sig_type in node_signature:
            # check that all compulsory signatures are present
            assert (sig_name in kwargs), \
                f"Node {node_name} must have signature: {node_signature}!"
            # check that compulsory signatures have the proper type
            cls._assert_type(kwargs[sig_name], sig_type)
        # IF SIGNATURE OK, CREATE AND RETURN NODE
        n = Node(node_name, **kwargs)
        return n

    @classmethod
    def _make_rel(cls, node_ori, node_dest, rel_sig, **kwargs):
        """
        First checks that the given ``node_ori, node_dest, **kwargs`` fulfill
        the given ``rel_sig``, and if so creates and returns the corresponding
        relationship. Otherwise raises an exception
        """
        node_types, sig_dict = rel_sig
        # NODE TYPE CHECK: if node_types given, we want to assert that
        # (ori_type, dest_type) appears at least once
        if node_types:
            ori_type = next(iter(node_ori.labels))
            dest_type = next(iter(node_dest.labels))
            node_types = [(next(iter(getattr(cls, o))),
                           next(iter(getattr(cls, d))))
                          for o, d in node_types]
            typematches = [(ori_type, dest_type) == od
                           for od in node_types]
            assert any(typematches), \
                f"Relation ({ori_type}, {dest_type}) isn't one of {node_types}!"
        # SIGNATURE CHECK
        rel_name, rel_signature = next(iter(sig_dict.items()))
        for sig_name, sig_type in rel_signature:
            # check that all compulsory signatures are present
            assert (sig_name in kwargs), \
                f"Relation {rel_name} must have signature: {rel_signature}!"
            # check that compulsory signatures have the proper type
            cls._assert_type(kwargs[sig_name], sig_type)
        # IF SIGNATURE OK, CREATE AND RETURN RELATIONSHIP
        r = Relationship(node_ori, rel_name, node_dest, **kwargs)
        return r


class ExampleSchema(Neo4jSchemaBase):
    """
    This class exemplifies the extension of ``Neo4jSchemabase``. Check the
    corresponding docstring for more details. Usage example::

      cat = ExampleSchema.make_cat_node(name="Dr. Whiskers")
      person = ExampleSchema.make_person_node(name="John Doe", age=100)
      house = ExampleSchema.make_house_node(address="Broad Way, 100")
      r2 = ExampleSchema.make_owns_relation(cat, house)  # cat owns house
      r1 = ExampleSchema.make_owns_relation(cat, person)  # cat owns person
      r3 = ExampleSchema.make_lives_in_relation(person, house, role="Loyal pet",
                                                room_number=3)
      r4 = ExampleSchema.make_lives_in_relation(cat, house, role="Owner",
                                                room_number=None)
    """
    CAT_NODE = {"Cat": [("name", str)]}
    PERSON_NODE = {"Person": [("name", str), ("age", int)]}
    HOUSE_NODE = {"House": [("address", str)]}
    LIVES_IN_REL = [[("CAT_NODE", "HOUSE_NODE"),
                     ("PERSON_NODE", "HOUSE_NODE")],
                    {"LIVES_IN": [("role", str), ("room_number", "optint")]}]
    OWNS_REL = [[("CAT_NODE", "HOUSE_NODE"),
                 ("CAT_NODE", "PERSON_NODE")],
                {"OWNS": []}]

    @classmethod
    def _assert_type(cls, x, typecode):
        """
        Extended method to handle 'optint' typecode separately
        """
        if typecode == "optint":
            assert_opt_int(x)
        else:
            assert_types(x, [typecode])

    @classmethod
    def make_cat_node(cls, **kwargs):
        return cls._make_node(cls.CAT_NODE, **kwargs)

    @classmethod
    def make_person_node(cls, **kwargs):
        return cls._make_node(cls.PERSON_NODE, **kwargs)

    @classmethod
    def make_house_node(cls, **kwargs):
        return cls._make_node(cls.HOUSE_NODE, **kwargs)

    @classmethod
    def make_lives_in_relation(cls, node_ori, node_dest, **kwargs):
        return cls._make_rel(node_ori, node_dest, cls.LIVES_IN_REL, **kwargs)

    @classmethod
    def make_owns_relation(cls, node_ori, node_dest, **kwargs):
        return cls._make_rel(node_ori, node_dest, cls.OWNS_REL, **kwargs)


# ##############################################################################
# # MAIN ROUTINE
# ##############################################################################

# create the nodes and relationships
cat = ExampleSchema.make_cat_node(name="Dr. Whiskers")
person = ExampleSchema.make_person_node(name="John Doe", age=100)
house = ExampleSchema.make_house_node(address="Broad Way, 123")
r1 = ExampleSchema.make_owns_relation(cat, person)
r2 = ExampleSchema.make_owns_relation(cat, house)
r3 = ExampleSchema.make_lives_in_relation(person, house, role="Loyal pet",
                                          room_number=3)
r4 = ExampleSchema.make_lives_in_relation(cat, house, role="Owner",
                                          room_number=None)

# THIS CODE WILL TRY TO LOG INTO THE DATABASE AND ADD THE NODES AND RELATIONS
# PASSWORD = "LOGIN123"
# g = Graph(password=PASSWORD)
# add_nodes_and_relations(g, nodes=[cat, person, house],
#                         relations=[r1, r2, r3, r4])
import pdb; pdb.set_trace()