diff options
| -rw-r--r-- | wqflask/wqflask/collect.py | 5 | ||||
| -rw-r--r-- | wqflask/wqflask/user_manager.py | 49 |
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) |