Added option to add traits to existing collection, which I had forgotten previously

2 files changed, 39 insertions, 15 deletions

diff --git a/wqflask/wqflask/collect.py b/wqflask/wqflask/collect.py
index e1aceac8..ecc02c06 100644
--- a/wqflask/wqflask/collect.py
+++ b/wqflask/wqflask/collect.py
@@ -201,9 +201,12 @@ def collections_new():
         return create_new(collection_name)
     elif "add_to_existing" in params:
         logger.debug("in add to existing")
+        collection_id = params['existing_collection'].split(":")[0]
         collection_name = params['existing_collection'].split(":")[1]
         if g.user_session.logged_in:
-            return UserCollection().add_traits(params, collection_name)
+            traits = list(process_traits(params['traits']))
+            g.user_session.add_traits_to_collection(collection_id, traits)
+            return redirect(url_for('view_collection', uc_id=collection_id))
         else:
             ac = AnonCollection(collection_name)
             ac.add_traits(params)
diff --git a/wqflask/wqflask/user_manager.py b/wqflask/wqflask/user_manager.py
index 3772414c..5d388d66 100644
--- a/wqflask/wqflask/user_manager.py
+++ b/wqflask/wqflask/user_manager.py
@@ -320,14 +320,14 @@ class UserSession(object):
             if user_info['collections'] != [] and user_info['collections'] != "[]":
                 current_collections = json.loads(user_info['collections'])
                 current_collections.append(collection_dict)
-                collections_json = json.dumps(current_collections)
+                self.update_collections(current_collections)
+                #collections_json = json.dumps(current_collections)
             else:
-                collections_json = json.dumps([collection_dict])
+                self.update_collections([collection_dict])
+                #collections_json = json.dumps([collection_dict])
         else:
-            collections_json = json.dumps([collection_dict])
-
-        collection_body = {'doc': {'collections': collections_json}}
-        es.update(index='users', doc_type='local', id=user_id, refresh='wait_for', body=collection_body)
+            self.update_collections([collection_dict])
+            #collections_json = json.dumps([collection_dict])
 
         return collection_dict['id']
 
@@ -341,13 +341,31 @@ class UserSession(object):
             else:
                 updated_collections.append(collection)
 
-        es = get_elasticsearch_connection()
-
-        collection_body = {'doc': {'collections': json.dumps(updated_collections)}}
-        es.update(index='users', doc_type='local', id=self.es_user_id, refresh='wait_for', body=collection_body)
+        self.update_collections(updated_collections)
 
         return collection['name']
 
+    def add_traits_to_collection(self, collection_id, traits_to_add):
+        """Add specified traits to a collection"""
+
+        this_collection = self.get_collection_by_id(collection_id)
+
+        updated_collection = this_collection
+        updated_traits = this_collection['members'] + traits_to_add
+
+        updated_collection['members'] = updated_traits
+        updated_collection['num_members'] = len(updated_traits)
+        updated_collection['changed_timestamp'] = datetime.datetime.utcnow().strftime('%b %d %Y %I:%M%p')
+
+        updated_collections = []
+        for collection in self.user_collections:
+            if collection['id'] == collection_id:
+                updated_collections.append(updated_collection)
+            else:
+                updated_collections.append(collection)
+
+        self.update_collections(updated_collections)
+
     def remove_traits_from_collection(self, collection_id, traits_to_remove):
         """Remove specified traits from a collection"""
 
@@ -372,10 +390,7 @@ class UserSession(object):
             else:
                 updated_collections.append(collection)
 
-        es = get_elasticsearch_connection()
-
-        collection_body = {'doc': {'collections': json.dumps(updated_collections)}}
-        es.update(index='users', doc_type='local', id=self.es_user_id, refresh='wait_for', body=collection_body)
+        self.update_collections(updated_collections)
 
         return updated_traits
 
@@ -391,6 +406,12 @@ class UserSession(object):
 
         return None
 
+    def update_collections(self, updated_collections):
+        es = get_elasticsearch_connection()
+
+        collection_body = {'doc': {'collections': json.dumps(updated_collections)}}
+        es.update(index='users', doc_type='local', id=self.es_user_id, refresh='wait_for', body=collection_body)
+
     def delete_session(self):
         # And more importantly delete the redis record
         Redis.delete(self.cookie_name)