Move node formatting to helper function

2 files changed, 49 insertions, 19 deletions

diff --git a/inventory/api.py b/inventory/api.py
index 388a698..91872f5 100644
--- a/inventory/api.py
+++ b/inventory/api.py
@@ -73,12 +73,58 @@ def auth_required(f):
     return wrapper
 
 
+def format_node(node):
+    parents = dict((parent['_id'], parent) for parent in node.pop('parents'))
+    current = node
+    # Recursively assign each node their respective parent
+    while current:
+        # NOTE: Python assigns from left to right
+        current['parent'] = current = parents.get(current.get('parent_id'))
+
+    children = dict((child['_id'], child) for child in node.pop('children'))
+    children_list = node['children'] = []
+    for child in children.values():
+        # Direct children should be assigned to the top-most list
+        if child['parent_id'] == node['_id']:
+            children_list.append(child)
+            continue
+        # Otherwise assign to their respective parent
+        parent = children[child['parent_id']]
+        parent.setdefault('children', []).append(child)
+
+    return node
+
+def format_nodes(nodes):
+    return (format_node(node) for node in nodes)
+
 # Routes
 @app.route('/nodes')
 @auth_required
 def root_nodes():
+    result = mongo.db.nodes.aggregate([
+        {'$match': {'parent_id': None, 'user_id': g.user['_id']}},
+        {
+            '$graphLookup': {
+                'from': 'nodes',
+                'startWith': '$parent_id',
+                'connectFromField': 'parent_id',
+                'connectToField': '_id',
+                'as': 'parents',
+            },
+        },
+        {
+            '$graphLookup': {
+                'from': 'nodes',
+                'startWith': '$_id',
+                'connectFromField': '_id',
+                'connectToField': 'parent_id',
+                'as': 'children',
+                'maxDepth': 1,
+            },
+        },
+    ])
     schema = NodeSchema(many=True)
-    data = schema.dump(mongo.db.nodes.find({'parent_id': None}))
+    data = schema.dump(format_nodes(result))
     return jsonify(data)
 
 
@@ -129,23 +175,7 @@ def node(node_id):
     except StopIteration:
         abort(404, 'No node found')
 
-    parents = dict((parent['_id'], parent) for parent in node.pop('parents'))
-    current = node
-    # Recursively assign each node their respective parent
-    while current:
-        # NOTE: Python assigns from left to right
-        current['parent'] = current = parents.get(current.get('parent_id'))
-
-    children = dict((child['_id'], child) for child in node.pop('children'))
-    children_list = node['children'] = []
-    for child in children.values():
-        # Direct children should be assigned to the top-most list
-        if child['parent_id'] == node['_id']:
-            children_list.append(child)
-            continue
-        # Otherwise assign to their respective parent
-        parent = children[child['parent_id']]
-        parent.setdefault('children', []).append(child)
+    node = format_node(node)
 
     schema = NodeSchema()
     return jsonify(schema.dump(node))
diff --git a/inventory/schema.py b/inventory/schema.py
index 4837d58..0728aea 100644
--- a/inventory/schema.py
+++ b/inventory/schema.py
@@ -27,7 +27,7 @@ class FieldSchema(Schema):
 class NodeSchema(Schema):
     _id = ObjectId(dump_only=True, data_key='id')
     name = fields.String(required=True)
-    parent_id = ObjectId(default=None)
+    parent_id = ObjectId(default=None, allow_none=True)
     _fields = fields.List(fields.Nested(FieldSchema()), default=[], attribute='fields', data_key='fields')
 
     # These are not set by the caller, but by the API endpoint