Add exporter to a local MongoDB

.travis.yml

@@ -3,6 +3,7 @@ dist: xenial
 os: linux
 
 services:
+  - mongodb
   - mysql
   - postgresql
 

setup.cfg

@@ -65,6 +65,7 @@ install_requires =
     tabulate
     pandas
     jsonschema
+    pymongo
 
 # Random options
 zip_safe = false

src/pybel/__init__.py

@@ -42,6 +42,7 @@
 from .io.umbrella_nodelink import to_umbrella_nodelink, to_umbrella_nodelink_file, to_umbrella_nodelink_gz
 from .manager import Manager, from_database, to_database
 from .parser.parse_bel import parse
+from .schema import is_valid_edge, is_valid_node
 from .struct import BELGraph, Pipeline, Query
 from .struct.operations import union
 from .version import get_version

src/pybel/io/mongodb.py

@@ -0,0 +1,152 @@
+# -*- coding: utf-8 -*-
+"""An exporter from PyBEL graphs to a local Mongo database."""
+
+import copy
+from typing import Any, List, Mapping, Tuple
+
+import pymongo
+from pymongo.collection import Collection
+
+import pybel
+from pybel.schema import is_valid_edge, is_valid_node
+
+__all__ = [
+    '_entity_to_dict',
+    'to_mongodb',
+    'find_nodes',
+    'get_edges_from_node',
+    'get_edges_from_criteria'
+]
+
+client = pymongo.MongoClient()
+
+
+def _handle_member(elem):
+    """Convert a member element from a complex or composite to a dictionary."""
+    # Convert each member element to a dictionary
+    new_elem = dict(elem)
+    # Remove the 'bel' and 'id' properties (which occur seemingly at random)
+    for prop in ('bel', 'id'):
+        if prop in new_elem.keys():
+            del new_elem[prop]
+    return new_elem
+
+
+def _entity_to_dict(entity: pybel.dsl.Entity) -> Mapping[str, Any]:
+    """Input a pybel Entity and return a dict representing it."""
+    new_node = dict(entity)
+    if new_node['function'] in ['Complex', 'Composite']:
+        new_node['members'] = list(map(_handle_member, new_node['members']))
+    return new_node
+
+
+def to_mongodb(graph: pybel.BELGraph, db_name: str, collection_name: str) -> Collection:
+    """Export the given BELGraph to a MongoDB.
+
+    In order to use this function, MongoDB must already be running locally.
+
+    :param graph: the graph to be exported
+    :param db_name: the name of the MongoDB to which the graph should be exported
+    :param collection_name: the name of the collection within the MongoDB where the graph will be stored.
+    :return: the collection that now stores the graph
+    """
+    # Access (or create) the specified database and collection
+    db = client[db_name]
+    collection = db[collection_name]
+    # If a collection with the same name already exists, drop its contents
+    collection.drop()
+
+    # Add the nodes
+    for node in graph:
+        if not is_valid_node(node):
+            # TODO: Raise/log on invalid node
+            pass
+        # Add a 'type' parameter to avoid confusing nodes and links
+        n = copy.deepcopy(node)
+        n = _entity_to_dict(n)
+        n['type'] = 'node'
+        collection.insert_one(n)
+
+    # Add the edges
+    edges = pybel.to_nodelink(graph)['links']
+    for edge in edges:
+        if not is_valid_edge(edge):
+            # TODO: Raise/log on invalid edge
+            pass
+        # Add a 'type' parameter to avoid confusing nodes and links
+        e = copy.deepcopy(edge)
+        e['type'] = 'link'
+        collection.insert_one(e)
+
+    return collection
+
+
+def find_nodes(
+    collection: Collection,
+    name: str = None,
+    identifier: str = None,
+    variants: List[pybel.dsl.EntityVariant] = None,
+) -> List[Mapping[str, Any]]:
+    """Find all the nodes that match the given criteria from a MongoDB Collection where a graph is stored.
+
+    :param collection: A MongoDB collection within a database where a PyBEL graph has been stored
+    :param name: The name of the desired node
+    :param identifier: The identifier of the desired node
+    :param variants: A list of variants that the desired node should contain. Note: nodes that contain the variants in addition to specified variants will be matched.
+    :return: A list containing all the nodes that match the given criteria
+    """
+    if not (name or identifier):
+        raise ValueError("Either a 'name' or 'identifier' is required to find a node.")
+    filter_ = {'type': 'node'}
+    if name:
+        filter_['concept.name'] = name
+    if identifier:
+        filter_['concept.identifier'] = identifier
+    if variants:
+        filter_['variants'] = variants
+
+    return list(collection.find(filter_))
+
+
+def get_edges_from_node(collection: Collection, node: Mapping[str, Any]) -> List[Mapping[str, Any]]:
+    """Return all the edges for the given node.
+
+    :param collection: A MongoDB collection within a database where a PyBEL graph has been stored
+    :param node: The node whose edges should be returned
+    :return: A list of the edges pointing to or from the given node
+    """
+    if not is_valid_node(node):
+        raise ValueError("Invalid node ", node)
+    # Remove the _id and type properties from the node (since they won't be included in the edge source/target information)
+    n = copy.deepcopy(node)
+    # Convert the pybel node to a dict so it can be matched against entries in the MongoDB
+    n = _entity_to_dict(n)
+    del n['_id'], n['type']
+    # Find all the links where either the source or the target is node n
+    filter_ = {'type': 'link', '$or': [{'source': n}, {'target': n}]}
+    return list(collection.find(filter_))
+
+
+def get_edges_from_criteria(
+    collection: Collection,
+    node_name: str = None,
+    node_identifier: str = None,
+    node_variants: List[pybel.dsl.EntityVariant] = None,
+) -> List[Tuple[Mapping[str, Any], List[Mapping[str, Any]]]]:
+    """Get all the edges for nodes that match the given criteria and return in a list of tuples.
+
+    :param collection: A MongoDB collection within a database where a PyBEL graph has been stored
+    :param name: The name of the desired node
+    :param identifier: The identifier of the desired node
+    :param variants: A list of variants that the desired node should contain. Note: nodes that contain the variants in addition to specified variants will be matched.
+    :return: A list of tuples. The first element of each tuple is the node, and the second element is a list of the edges.
+    """
+    if not (node_name or node_identifier):
+        raise ValueError("Either a name or an identifier is required to get edges.")
+    matching_nodes = find_nodes(collection, name=node_name, identifier=node_identifier, variants=node_variants)
+    edges = []
+    for node in matching_nodes:
+        matching_edges = get_edges_from_node(collection, node)
+        # Append (node, [edge1, edge2...]) to edges
+        edges.append((node, matching_edges))
+    return edges

tests/test_io/test_mongodb.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+
+"""Tests for the jsonschema node validation."""
+
+import copy
+from typing import Any, List, Mapping, Tuple
+import unittest
+
+import pymongo
+from pymongo.collection import Collection
+
+from pybel import BELGraph, to_nodelink
+from pybel.dsl import Entity, ComplexAbundance, Gene, Protein, Rna
+from pybel.io.mongodb import (
+    _entity_to_dict, 
+    to_mongodb,
+    find_nodes,
+    get_edges_from_node,
+    get_edges_from_criteria
+)
+from pybel.testing.utils import n
+
+g1 = Gene('hgnc', '1')
+r1 = Rna('hgnc', '1')
+p1 = Protein('hgnc', '1')
+g2 = Gene('hgnc', '2')
+r2 = Rna('hgnc', '2')
+p2 = Protein('hgnc', '2')
+g3 = Gene('hgnc', '3')
+p3 = Protein('hgnc', '3')
+
+ca = ComplexAbundance([p1, g2], name='ca', namespace='hgnc')
+
+client = pymongo.MongoClient()
+TEST_DB = client['test_db']
+TEST_COLLECTION = TEST_DB['test_collection']
+TEST_COLLECTION.drop()
+
+
+class TestMongoDB(unittest.TestCase):
+    """Tests for the MongoDB exporting and querying."""
+    def setUp(self):
+        """Create and export a test graph."""
+        self.graph = BELGraph()
+        self.graph.add_increases(ca, r2, citation=n(), evidence=n())
+        self.graph.add_increases(p2, p3, citation=n(), evidence=n())
+        self.graph.add_decreases(p3, p1, citation=n(), evidence=n())
+        self.links = to_nodelink(self.graph)['links']
+        self.collection = to_mongodb(self.graph, TEST_DB.name, TEST_COLLECTION.name)
+
+    def assert_node(self, node: Mapping[str, Any], assert_in=True) -> None:
+        """Assert that the given node is present or not present in the MongoDB collection."""
+        assert_fn = self.assertEqual
+        if not assert_in:
+            assert_fn = self.assertNotEqual
+        # Convert the node to a dict so it can be matched against entries in the MongoDB
+        node_dict = _entity_to_dict(node)
+        match = self.collection.find_one(node_dict)
+        # If the passed node has a MongoDB unique _id, check that the two ids match
+        if '_id' in node.keys():
+            assert_fn(node_dict['_id'], match['_id'])
+        else:
+            # Otherwise, delete the parameters type and _id unique to the to_mongodb output
+            if match is not None:
+                del match['_id'], match['type']
+            # And check that the two nodes are equal
+            assert_fn(node_dict, match)
+
+    def assert_edge(self, edge: dict, assert_in=True) -> None:
+        """Assert that the given edge is or isn't present in the Mongo collection."""
+        assert_fn = self.assertEqual
+        if not assert_in:
+            assert_fn = self.assertNotEqual
+        if type(edge) is not dict:
+            raise ValueError('Expeced type(edge) to be dict. Insead found: ', type(edge))
+        # Query the collection for the given edge
+        match = self.collection.find_one(edge)
+        if match is not None:
+            # If a match was found, delete the _id and type keys
+            del match['_id'], match['type']
+        # Assert that the given edge and match are or are not equal (based on assert_in param)
+        assert_fn(edge, match)
+
+    def test_export(self):
+        """Test that the to_mongodb export function operates as expeced."""
+        # Assert that all nodes in self.graph are in self.collection
+        for node in self.graph:
+            self.assert_node(node)
+        # Assert that the two nodes not in self.graph are not in self.collection
+        self.assert_node(g1, assert_in=False)
+        self.assert_node(r1, assert_in=False)
+        # Assert that every edge from self.graph is in self.collection
+        for link in self.links:
+            self.assert_edge(link)
+
+    def test_query_nodes(self):
+        """Test that the find_nodes() function correctly finds the desired nodes."""
+        for node in self.graph:
+            # Convert the node to a dict
+            n = _entity_to_dict(node)
+            # Get the concept entry (where the name / identifier are stored)
+            concept = n['concept']
+            name, identifier, variants = None, None, None
+            # Query based on name, identifier, variants if they are present
+            if 'name' in concept.keys():
+                name = concept['name']
+            if 'identifier' in concept.keys():
+                identifier = concept['identifier']
+            if 'variants' in n.keys():
+                variants = n['variants']
+
+            matches = find_nodes(self.collection, name=name, identifier=identifier, variants=variants)
+            for match in matches:
+                del match['_id'], match['type']
+
+            self.assertIn(n, matches)
+
+
+if __name__ == '__main__':
+    unittest.main()