This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/v1-10-test by this push:
     new 1249b50  [AIRFLOW-8902] Fix Dag Run UI execution date with timezone 
cannot be saved issue (#8902)
1249b50 is described below

commit 1249b501ee315e579b462f30970aef74f3c99121
Author: Jacob Shao <41271167+realradi...@users.noreply.github.com>
AuthorDate: Mon May 25 09:09:02 2020 -0400

    [AIRFLOW-8902] Fix Dag Run UI execution date with timezone cannot be saved 
issue (#8902)
    
    Closes #8842
---
 airflow/www_rbac/forms.py    | 61 +++++++++++++++++++++++++++++-------
 airflow/www_rbac/utils.py    |  7 +++--
 tests/www_rbac/test_views.py | 73 +++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 123 insertions(+), 18 deletions(-)

diff --git a/airflow/www_rbac/forms.py b/airflow/www_rbac/forms.py
index c4f9796..48b9142 100644
--- a/airflow/www_rbac/forms.py
+++ b/airflow/www_rbac/forms.py
@@ -23,33 +23,75 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 import json
+from datetime import datetime as dt
 
+import pendulum
 from flask_appbuilder.fieldwidgets import (
     BS3PasswordFieldWidget, BS3TextAreaFieldWidget, BS3TextFieldWidget, 
Select2Widget,
 )
 from flask_appbuilder.forms import DynamicForm
 from flask_babel import lazy_gettext
 from flask_wtf import FlaskForm
+from wtforms import validators, widgets
+from wtforms.fields import (
+    BooleanField, Field, IntegerField, PasswordField, SelectField, 
StringField, TextAreaField,
+)
 
-from wtforms import validators
-from wtforms.fields import (IntegerField, SelectField, TextAreaField, 
PasswordField,
-                            StringField, DateTimeField, BooleanField)
+from airflow.configuration import conf
 from airflow.models import Connection
 from airflow.utils import timezone
 from airflow.www_rbac.validators import ValidJson
 from airflow.www_rbac.widgets import AirflowDateTimePickerWidget
 
 
+class DateTimeWithTimezoneField(Field):
+    """
+    A text field which stores a `datetime.datetime` matching a format.
+    """
+    widget = widgets.TextInput()
+
+    def __init__(self, label=None, validators=None, format=None, **kwargs):
+        super(DateTimeWithTimezoneField, self).__init__(label, validators, 
**kwargs)
+        self.format = format or "%Y-%m-%d %H:%M:%S%Z"
+
+    def _value(self):
+        if self.raw_data:
+            return ' '.join(self.raw_data)
+        else:
+            return self.data and self.data.strftime(self.format) or ''
+
+    def process_formdata(self, valuelist):
+        if valuelist:
+            date_str = ' '.join(valuelist)
+            try:
+                # Check if the datetime string is in the format without 
timezone, if so convert it to the
+                # default timezone
+                if len(date_str) == 19:
+                    parsed_datetime = dt.strptime(date_str, '%Y-%m-%d 
%H:%M:%S')
+                    defualt_timezone = pendulum.timezone('UTC')
+                    tz = conf.get("core", "default_timezone")
+                    if tz == "system":
+                        defualt_timezone = pendulum.local_timezone()
+                    else:
+                        defualt_timezone = pendulum.timezone(tz)
+                    self.data = defualt_timezone.convert(parsed_datetime)
+                else:
+                    self.data = pendulum.parse(date_str)
+            except ValueError:
+                self.data = None
+                raise ValueError(self.gettext('Not a valid datetime value'))
+
+
 class DateTimeForm(FlaskForm):
     # Date filter form needed for task views
-    execution_date = DateTimeField(
+    execution_date = DateTimeWithTimezoneField(
         "Execution date", widget=AirflowDateTimePickerWidget())
 
 
 class DateTimeWithNumRunsForm(FlaskForm):
     # Date time and number of runs form for tree view, task duration
     # and landing times
-    base_date = DateTimeField(
+    base_date = DateTimeWithTimezoneField(
         "Anchor date", widget=AirflowDateTimePickerWidget(), 
default=timezone.utcnow())
     num_runs = SelectField("Number of runs", default=25, choices=(
         (5, "5"),
@@ -70,10 +112,10 @@ class DagRunForm(DynamicForm):
         lazy_gettext('Dag Id'),
         validators=[validators.DataRequired()],
         widget=BS3TextFieldWidget())
-    start_date = DateTimeField(
+    start_date = DateTimeWithTimezoneField(
         lazy_gettext('Start Date'),
         widget=AirflowDateTimePickerWidget())
-    end_date = DateTimeField(
+    end_date = DateTimeWithTimezoneField(
         lazy_gettext('End Date'),
         widget=AirflowDateTimePickerWidget())
     run_id = StringField(
@@ -84,7 +126,7 @@ class DagRunForm(DynamicForm):
         lazy_gettext('State'),
         choices=(('success', 'success'), ('running', 'running'), ('failed', 
'failed'),),
         widget=Select2Widget())
-    execution_date = DateTimeField(
+    execution_date = DateTimeWithTimezoneField(
         lazy_gettext('Execution Date'),
         widget=AirflowDateTimePickerWidget())
     external_trigger = BooleanField(
@@ -95,10 +137,7 @@ class DagRunForm(DynamicForm):
         widget=BS3TextAreaFieldWidget())
 
     def populate_obj(self, item):
-        # TODO: This is probably better done as a custom field type so we can
-        # set TZ at parse time
         super(DagRunForm, self).populate_obj(item)
-        item.execution_date = timezone.make_aware(item.execution_date)
         if item.conf:
             item.conf = json.loads(item.conf)
 
diff --git a/airflow/www_rbac/utils.py b/airflow/www_rbac/utils.py
index ffd79b1..f493490 100644
--- a/airflow/www_rbac/utils.py
+++ b/airflow/www_rbac/utils.py
@@ -32,8 +32,8 @@ from past.builtins import basestring
 
 from pygments import highlight, lexers
 from pygments.formatters import HtmlFormatter
-from flask import request, Response, Markup, url_for
-from flask_appbuilder.forms import DateTimeField, FieldConverter
+from flask import Markup, Response, request, url_for
+from flask_appbuilder.forms import FieldConverter
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 import flask_appbuilder.models.sqla.filters as fab_sqlafilters
 import sqlalchemy as sqla
@@ -45,6 +45,7 @@ from airflow.operators.subdag_operator import SubDagOperator
 from airflow.utils import timezone
 from airflow.utils.json import AirflowJsonEncoder
 from airflow.utils.state import State
+from airflow.www_rbac.forms import DateTimeWithTimezoneField
 from airflow.www_rbac.widgets import AirflowDateTimePickerWidget
 
 AUTHENTICATE = conf.getboolean('webserver', 'AUTHENTICATE')
@@ -467,6 +468,6 @@ class CustomSQLAInterface(SQLAInterface):
 # subclass) so we have no other option than to edit the converstion table in
 # place
 FieldConverter.conversion_table = (
-    (('is_utcdatetime', DateTimeField, AirflowDateTimePickerWidget),) +
+    (('is_utcdatetime', DateTimeWithTimezoneField, 
AirflowDateTimePickerWidget),) +
     FieldConverter.conversion_table
 )
diff --git a/tests/www_rbac/test_views.py b/tests/www_rbac/test_views.py
index 9a199f6..fb48cfe 100644
--- a/tests/www_rbac/test_views.py
+++ b/tests/www_rbac/test_views.py
@@ -28,7 +28,7 @@ import sys
 import tempfile
 import unittest
 import urllib
-from datetime import timedelta
+from datetime import datetime as dt, timedelta, timezone as tz
 
 import pytest
 import six
@@ -2489,7 +2489,72 @@ class TestDagRunModelView(TestBase):
     def tearDown(self):
         self.clear_table(models.DagRun)
 
-    def test_create_dagrun(self):
+    def test_create_dagrun_execution_date_with_timezone_utc(self):
+        data = {
+            "state": "running",
+            "dag_id": "example_bash_operator",
+            "execution_date": "2018-07-06 05:04:03Z",
+            "run_id": "test_create_dagrun",
+        }
+        resp = self.client.post('/dagrun/add',
+                                data=data,
+                                follow_redirects=True)
+        self.check_content_in_response('Added Row', resp)
+
+        dr = self.session.query(models.DagRun).one()
+
+        self.assertEqual(dr.execution_date, dt(2018, 7, 6, 5, 4, 3, 
tzinfo=tz.utc))
+
+    def test_create_dagrun_execution_date_with_timezone_edt(self):
+        data = {
+            "state": "running",
+            "dag_id": "example_bash_operator",
+            "execution_date": "2018-07-06 05:04:03-04:00",
+            "run_id": "test_create_dagrun",
+        }
+        resp = self.client.post('/dagrun/add',
+                                data=data,
+                                follow_redirects=True)
+        self.check_content_in_response('Added Row', resp)
+
+        dr = self.session.query(models.DagRun).one()
+
+        self.assertEqual(dr.execution_date, dt(2018, 7, 6, 5, 4, 3, 
tzinfo=tz(timedelta(hours=-4))))
+
+    def test_create_dagrun_execution_date_with_timezone_pst(self):
+        data = {
+            "state": "running",
+            "dag_id": "example_bash_operator",
+            "execution_date": "2018-07-06 05:04:03-08:00",
+            "run_id": "test_create_dagrun",
+        }
+        resp = self.client.post('/dagrun/add',
+                                data=data,
+                                follow_redirects=True)
+        self.check_content_in_response('Added Row', resp)
+
+        dr = self.session.query(models.DagRun).one()
+
+        self.assertEqual(dr.execution_date, dt(2018, 7, 6, 5, 4, 3, 
tzinfo=tz(timedelta(hours=-8))))
+
+    @conf_vars({("core", "default_timezone"): "America/Toronto"})
+    def test_create_dagrun_execution_date_without_timezone_default_edt(self):
+        data = {
+            "state": "running",
+            "dag_id": "example_bash_operator",
+            "execution_date": "2018-07-06 05:04:03",
+            "run_id": "test_create_dagrun",
+        }
+        resp = self.client.post('/dagrun/add',
+                                data=data,
+                                follow_redirects=True)
+        self.check_content_in_response('Added Row', resp)
+
+        dr = self.session.query(models.DagRun).one()
+
+        self.assertEqual(dr.execution_date, dt(2018, 7, 6, 5, 4, 3, 
tzinfo=tz(timedelta(hours=-4))))
+
+    def test_create_dagrun_execution_date_without_timezone_default_utc(self):
         data = {
             "state": "running",
             "dag_id": "example_bash_operator",
@@ -2503,14 +2568,14 @@ class TestDagRunModelView(TestBase):
 
         dr = self.session.query(models.DagRun).one()
 
-        self.assertEqual(dr.execution_date, 
timezone.convert_to_utc(datetime(2018, 7, 6, 5, 4, 3)))
+        self.assertEqual(dr.execution_date, dt(2018, 7, 6, 5, 4, 3, 
tzinfo=tz.utc))
 
     def test_create_dagrun_valid_conf(self):
         conf_value = dict(Valid=True)
         data = {
             "state": "running",
             "dag_id": "example_bash_operator",
-            "execution_date": "2018-07-06 05:05:03",
+            "execution_date": "2018-07-06 05:05:03-02:00",
             "run_id": "test_create_dagrun_valid_conf",
             "conf": json.dumps(conf_value)
         }

Reply via email to