diff options
| -rw-r--r-- | inventory/api.py | 66 | ||||
| -rw-r--r-- | inventory/schema.py | 2 |
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 |