From 09d3ae27b62fe7c042e46a8676d75e0d02484170 Mon Sep 17 00:00:00 2001
From: Nicholas Staples <nicholas.staples@solution-tribute.com>
Date: Fri, 29 Apr 2016 12:24:43 +0100
Subject: [PATCH] Update created_by to work like other relationship fields.

---
 app/schemas.py | 35 +++++++++++------------------------
 1 file changed, 11 insertions(+), 24 deletions(-)

diff --git a/app/schemas.py b/app/schemas.py
index e5facc137..c84fee407 100644
--- a/app/schemas.py
+++ b/app/schemas.py
@@ -7,9 +7,7 @@ from marshmallow import (
     validates_schema,
     pre_load
 )
-from sqlalchemy.dialects.postgresql import UUID
 from marshmallow_sqlalchemy import field_for
-from marshmallow_sqlalchemy.convert import ModelConverter
 
 from notifications_utils.recipients import (
     validate_email_address,
@@ -48,25 +46,6 @@ class BaseSchema(ma.ModelSchema):
         return super(BaseSchema, self).make_instance(data)
 
 
-class CreatedBySchema(ma.Schema):
-
-    created_by = fields.Str(required=True, load_only=True)
-
-    @validates_schema
-    def validates_created_by(self, data):
-        try:
-            if not isinstance(data.get('created_by'), models.User):
-                created_by = models.User.query.filter_by(id=data.get('created_by')).one()
-        except:
-            raise ValidationError('Invalid created_by: {}'.format(data))
-
-    @post_load
-    def format_created_by(self, item):
-        if not isinstance(item.get('created_by'), models.User):
-            item['created_by'] = models.User.query.filter_by(id=item.get('created_by')).one()
-        return item
-
-
 class UserSchema(BaseSchema):
 
     permissions = fields.Method("user_permissions", dump_only=True)
@@ -89,7 +68,10 @@ class UserSchema(BaseSchema):
             "_password", "verify_codes")
 
 
-class ServiceSchema(BaseSchema, CreatedBySchema):
+class ServiceSchema(BaseSchema):
+
+    created_by = field_for(models.Service, 'created_by', required=True)
+
     class Meta:
         model = models.Service
         exclude = ("updated_at", "created_at", "api_keys", "templates", "jobs", 'old_id')
@@ -107,7 +89,9 @@ class BaseTemplateSchema(BaseSchema):
         exclude = ("updated_at", "created_at", "service_id", "jobs")
 
 
-class TemplateSchema(BaseTemplateSchema, CreatedBySchema):
+class TemplateSchema(BaseTemplateSchema):
+
+    created_by = field_for(models.Template, 'created_by', required=True)
 
     @validates_schema
     def validate_type(self, data):
@@ -123,7 +107,10 @@ class NotificationsStatisticsSchema(BaseSchema):
         model = models.NotificationStatistics
 
 
-class ApiKeySchema(BaseSchema, CreatedBySchema):
+class ApiKeySchema(BaseSchema):
+
+    created_by = field_for(models.ApiKey, 'created_by', required=True)
+
     class Meta:
         model = models.ApiKey
         exclude = ("service", "secret")