From a96526351bab4432a70c8b0def32aa1206a48cb2 Mon Sep 17 00:00:00 2001
From: Francesco Siddi <francesco.siddi@gmail.com>
Date: Sun, 15 Nov 2015 17:46:32 +0100
Subject: [PATCH] Tweaks to upgrade_node_types

---
 pillar/manage.py                     | 43 +++++++++++-----------------
 pillar/manage/node_types/__init__.py |  8 ++++++
 pillar/manage/node_types/asset.py    | 23 ++++++++++-----
 pillar/manage/node_types/project.py  | 22 +++-----------
 4 files changed, 44 insertions(+), 52 deletions(-)

diff --git a/pillar/manage.py b/pillar/manage.py
index 21434aa6..36df884d 100644
--- a/pillar/manage.py
+++ b/pillar/manage.py
@@ -62,23 +62,17 @@ def clear_db():
 
 @manager.command
 def upgrade_node_types():
-    """Wipes node_types collection
-    and populates it again
-    """
-    from pymongo import MongoClient
-
-    client = MongoClient(MONGO_HOST, 27017)
-    db = client.eve
-    node_types = db.node_types.find({})
+    """Wipes node_types collection and populates it again"""
+    node_types_collection = app.data.driver.db['node_types']
+    node_types = node_types_collection.find({})
     old_ids = {}
-    for nt in node_types:
-        old_ids[nt['name']] = nt['_id']
+    for node_type in node_types:
+        old_ids[node_type['name']] = node_type['_id']
     populate_node_types(old_ids)
 
 
 def get_id(collection, name):
-    """Returns the _id of the given collection
-    and name."""
+    """Returns the _id of the given collection and name"""
     from pymongo import MongoClient
     client = MongoClient(MONGO_HOST, 27017)
     db = client.eve
@@ -164,37 +158,32 @@ def manage_groups():
 
 
 def populate_node_types(old_ids={}):
-
-    from pymongo import MongoClient
-
-    client = MongoClient(MONGO_HOST, 27017)
-    db = client.eve
+    node_types_collection = app.data.driver.db['node_types']
 
     def mix_node_type(old_id, node_type_dict):
         # Take eve parameters
-        node_type = db.node_types.find({'_id':old_id})
-        node_type = node_type[0]
+        node_type = node_types_collection.find_one({'_id': old_id})
         for attr in node_type:
-            if attr[0]=='_':
-                # Mix with node type attributes
-                node_type_dict[attr]=node_type[attr]
+            if attr[0] == '_':
+                # Mix with node eve attributes. This is really not needed since
+                # the attributes are stripped before doing a put_internal.
+                node_type_dict[attr] = node_type[attr]
+            elif attr == 'permissions':
+                node_type_dict['permissions'] = node_type['permissions']
         return node_type_dict
 
     def upgrade(node_type, old_ids):
         print("Node {0}".format(node_type['name']))
         node_name = node_type['name']
         if node_name in old_ids:
-            node_type = mix_node_type(old_ids[node_name], node_type)
-            node_id = node_type['_id']
+            node_id = old_ids[node_name]
+            node_type = mix_node_type(node_id, node_type)
 
             # Removed internal fields that would cause validation error
             internal_fields = ['_id', '_etag', '_updated', '_created']
             for field in internal_fields:
                 node_type.pop(field, None)
-            # Also remove permissions, since they are managed separately
-            node_type.pop('permissions', None)
             p = put_internal('node_types', node_type, **{'_id': node_id})
-
         else:
             print("Making the node")
             print(node_type)
diff --git a/pillar/manage/node_types/__init__.py b/pillar/manage/node_types/__init__.py
index e69de29b..86a05647 100644
--- a/pillar/manage/node_types/__init__.py
+++ b/pillar/manage/node_types/__init__.py
@@ -0,0 +1,8 @@
+_file_embedded_schema = {
+    'type': 'objectid',
+    'data_relation': {
+        'resource': 'files',
+        'field': '_id',
+        'embeddable': True
+    }
+}
diff --git a/pillar/manage/node_types/asset.py b/pillar/manage/node_types/asset.py
index 5df5a1ef..a63b17e8 100644
--- a/pillar/manage/node_types/asset.py
+++ b/pillar/manage/node_types/asset.py
@@ -1,3 +1,5 @@
+from manage.node_types import _file_embedded_schema
+
 node_type_asset = {
     'name': 'asset',
     'description': 'Basic Asset Type',
@@ -23,19 +25,26 @@ node_type_asset = {
         },
         # We point to the original file (and use it to extract any relevant
         # variation useful for our scope).
-        'file': {
-            'type': 'objectid',
-            'data_relation': {
-                'resource': 'files',
-                'field': '_id',
-                'embeddable': True
-            },
+        'file': _file_embedded_schema,
+        'description_attachments': {
+            'type': 'list',
+            'schema': {
+                'type': 'dict',
+                'schema': {
+                    'file': _file_embedded_schema,
+                    'url': {
+                        'type': 'string',
+                        'minlength': 1
+                    }
+                }
+            }
         }
     },
     'form_schema': {
         'status': {},
         'content_type': {'visible': False},
         'file': {'visible': False},
+        'description_attachments': {'visible': False},
     },
     'permissions': {
         # 'groups': [{
diff --git a/pillar/manage/node_types/project.py b/pillar/manage/node_types/project.py
index 4759619c..24a6520f 100644
--- a/pillar/manage/node_types/project.py
+++ b/pillar/manage/node_types/project.py
@@ -1,3 +1,5 @@
+from manage.node_types import _file_embedded_schema
+
 node_type_project = {
     'name': 'project',
     'parent': {},
@@ -60,25 +62,9 @@ node_type_project = {
             ],
         },
         # Logo
-        'picture_square': {
-            'type': 'objectid',
-            'nullable': True,
-            'data_relation': {
-               'resource': 'files',
-               'field': '_id',
-               'embeddable': True
-            },
-        },
+        'picture_square': _file_embedded_schema,
         # Header
-        'picture_header': {
-            'type': 'objectid',
-            'nullable': True,
-            'data_relation': {
-               'resource': 'files',
-               'field': '_id',
-               'embeddable': True
-            },
-        },
+        'picture_header': _file_embedded_schema,
         # Short summary for the project
         'summary': {
             'type': 'string',