From f504568088461380edbbe27b28547042614cf6d4 Mon Sep 17 00:00:00 2001
From: Daniel Vaz Gaspar <danielvazgaspar@gmail.com>
Date: Tue, 30 Apr 2019 17:08:07 +0100
Subject: [PATCH] [annotations] Improves UX on annotation validation,
 start_dttm, end_dttm (#7326)

---
 superset/models/annotations.py |  2 +-
 superset/views/annotations.py  | 32 +++++++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/superset/models/annotations.py b/superset/models/annotations.py
index 027f14f9ba..1de9dfd2f9 100644
--- a/superset/models/annotations.py
+++ b/superset/models/annotations.py
@@ -46,7 +46,7 @@ class Annotation(Model, AuditMixinNullable):
     id = Column(Integer, primary_key=True)
     start_dttm = Column(DateTime)
     end_dttm = Column(DateTime)
-    layer_id = Column(Integer, ForeignKey('annotation_layer.id'))
+    layer_id = Column(Integer, ForeignKey('annotation_layer.id'), nullable=False)
     short_descr = Column(String(500))
     long_descr = Column(Text)
     layer = relationship(
diff --git a/superset/views/annotations.py b/superset/views/annotations.py
index 8177efc883..9ef2d65cad 100644
--- a/superset/views/annotations.py
+++ b/superset/views/annotations.py
@@ -18,12 +18,30 @@
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_babel import gettext as __
 from flask_babel import lazy_gettext as _
+from wtforms.validators import StopValidation
 
 from superset import appbuilder
 from superset.models.annotations import Annotation, AnnotationLayer
 from .base import DeleteMixin, SupersetModelView
 
 
+class StartEndDttmValidator(object):
+    """
+    Validates dttm fields.
+    """
+    def __call__(self, form, field):
+        if not form['start_dttm'].data and not form['end_dttm'].data:
+            raise StopValidation(
+                _('annotation start time or end time is required.'),
+            )
+        elif (form['end_dttm'].data and
+                form['start_dttm'].data and
+                form['end_dttm'].data < form['start_dttm'].data):
+            raise StopValidation(
+                _('Annotation end time must be no earlier than start time.'),
+            )
+
+
 class AnnotationModelView(SupersetModelView, DeleteMixin):  # noqa
     datamodel = SQLAInterface(Annotation)
 
@@ -53,17 +71,17 @@ class AnnotationModelView(SupersetModelView, DeleteMixin):  # noqa
          annotation needs to add more context.',
     }
 
+    validators_columns = {
+        'start_dttm': [
+            StartEndDttmValidator(),
+        ],
+    }
+
     def pre_add(self, obj):
-        if not obj.layer:
-            raise Exception('Annotation layer is required.')
-        if not obj.start_dttm and not obj.end_dttm:
-            raise Exception('Annotation start time or end time is required.')
-        elif not obj.start_dttm:
+        if not obj.start_dttm:
             obj.start_dttm = obj.end_dttm
         elif not obj.end_dttm:
             obj.end_dttm = obj.start_dttm
-        elif obj.end_dttm < obj.start_dttm:
-            raise Exception('Annotation end time must be no earlier than start time.')
 
     def pre_update(self, obj):
         self.pre_add(obj)