Repository: incubator-airflow Updated Branches: refs/heads/v1-10-test 2a8358353 -> 5a3618701
[AIRFLOW-1730] Unpickle value of XCom queried from DB Cherry picked from #2701 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/5a361870 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/5a361870 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/5a361870 Branch: refs/heads/v1-10-test Commit: 5a3618701d838a974f532ffbd7658f327a801661 Parents: 2a83583 Author: Shintaro Murakami <[email protected]> Authored: Tue Oct 17 18:40:19 2017 +0900 Committer: Ash Berlin-Taylor <[email protected]> Committed: Fri May 25 10:10:04 2018 +0100 ---------------------------------------------------------------------- airflow/models.py | 80 +++++++++++++++++++++----------------------------- tests/models.py | 59 ++++++++++++++++++++++++------------- 2 files changed, 72 insertions(+), 67 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5a361870/airflow/models.py ---------------------------------------------------------------------- diff --git a/airflow/models.py b/airflow/models.py index 1a66eb1..3077e6d 100755 --- a/airflow/models.py +++ b/airflow/models.py @@ -4349,6 +4349,24 @@ class XCom(Base, LoggingMixin): Index('idx_xcom_dag_task_date', dag_id, task_id, execution_date, unique=False), ) + """ + TODO: "pickling" has been deprecated and JSON is preferred. + "pickling" will be removed in Airflow 2.0. + """ + @reconstructor + def init_on_load(self): + enable_pickling = configuration.getboolean('core', 'enable_xcom_pickling') + if enable_pickling: + self.value = pickle.loads(self.value) + else: + try: + self.value = json.loads(self.value.decode('UTF-8')) + except (UnicodeEncodeError, ValueError): + # For backward-compatibility. + # Preventing errors in webserver + # due to XComs mixed with pickled and unpickled. + self.value = pickle.loads(self.value) + def __repr__(self): return '<XCom "{key}" ({task_id} @ {execution_date})>'.format( key=self.key, @@ -4364,23 +4382,16 @@ class XCom(Base, LoggingMixin): execution_date, task_id, dag_id, - enable_pickling=None, session=None): """ Store an XCom value. - TODO: "pickling" has been deprecated and JSON is preferred. "pickling" will be - removed in Airflow 2.0. :param enable_pickling: If pickling is not enabled, the - XCOM value will be parsed as JSON instead. - + TODO: "pickling" has been deprecated and JSON is preferred. + "pickling" will be removed in Airflow 2.0. :return: None """ session.expunge_all() - if enable_pickling is None: - enable_pickling = configuration.conf.getboolean( - 'core', 'enable_xcom_pickling' - ) - + enable_pickling = configuration.getboolean('core', 'enable_xcom_pickling') if enable_pickling: value = pickle.dumps(value) else: @@ -4422,13 +4433,11 @@ class XCom(Base, LoggingMixin): task_id=None, dag_id=None, include_prior_dates=False, - enable_pickling=None, session=None): """ Retrieve an XCom value, optionally meeting certain criteria. - TODO: "pickling" has been deprecated and JSON is preferred. "pickling" will be removed in Airflow 2.0. - - :param enable_pickling: If pickling is not enabled, the XCOM value will be parsed to JSON instead. + TODO: "pickling" has been deprecated and JSON is preferred. + "pickling" will be removed in Airflow 2.0. :return: XCom value """ filters = [] @@ -4450,11 +4459,7 @@ class XCom(Base, LoggingMixin): result = query.first() if result: - if enable_pickling is None: - enable_pickling = configuration.conf.getboolean( - 'core', 'enable_xcom_pickling' - ) - + enable_pickling = configuration.getboolean('core', 'enable_xcom_pickling') if enable_pickling: return pickle.loads(result.value) else: @@ -4462,7 +4467,7 @@ class XCom(Base, LoggingMixin): return json.loads(result.value.decode('UTF-8')) except ValueError: log = LoggingMixin().log - log.error("Could not serialize the XCOM value into JSON. " + log.error("Could not deserialize the XCOM value from JSON. " "If you are using pickles instead of JSON " "for XCOM, then you need to enable pickle " "support for XCOM in your airflow config.") @@ -4470,16 +4475,14 @@ class XCom(Base, LoggingMixin): @classmethod @provide_session - def get_many( - cls, - execution_date, - key=None, - task_ids=None, - dag_ids=None, - include_prior_dates=False, - limit=100, - enable_pickling=None, - session=None): + def get_many(cls, + execution_date, + key=None, + task_ids=None, + dag_ids=None, + include_prior_dates=False, + limit=100, + session=None): """ Retrieve an XCom value, optionally meeting certain criteria TODO: "pickling" has been deprecated and JSON is preferred. "pickling" will be removed in Airflow 2.0. @@ -4502,23 +4505,6 @@ class XCom(Base, LoggingMixin): .order_by(cls.execution_date.desc(), cls.timestamp.desc()) .limit(limit)) results = query.all() - if enable_pickling is None: - enable_pickling = configuration.conf.getboolean( - 'core', 'enable_xcom_pickling' - ) - for result in results: - if enable_pickling: - result.value = pickle.loads(result.value) - else: - try: - result.value = json.loads(result.value.decode('UTF-8')) - except ValueError: - log = LoggingMixin().log - log.error("Could not serialize the XCOM value into JSON. " - "If you are using pickles instead of JSON " - "for XCOM, then you need to enable pickle " - "support for XCOM in your airflow config.") - raise return results @classmethod http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5a361870/tests/models.py ---------------------------------------------------------------------- diff --git a/tests/models.py b/tests/models.py index c2e54e5..9f09027 100644 --- a/tests/models.py +++ b/tests/models.py @@ -2219,24 +2219,34 @@ class ClearTasksTest(unittest.TestCase): self.assertEqual(ti2.max_tries, 1) def test_xcom_disable_pickle_type(self): + configuration.load_test_config() + json_obj = {"key": "value"} execution_date = timezone.utcnow() key = "xcom_test1" dag_id = "test_dag1" task_id = "test_task1" + configuration.set("core", "enable_xcom_pickling", "False") + XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, - execution_date=execution_date, - enable_pickling=False) + execution_date=execution_date) ret_value = XCom.get_one(key=key, - dag_id=dag_id, - task_id=task_id, - execution_date=execution_date, - enable_pickling=False) + dag_id=dag_id, + task_id=task_id, + execution_date=execution_date) + + self.assertEqual(ret_value, json_obj) + + session = settings.Session() + ret_value = session.query(XCom).filter(XCom.key == key, XCom.dag_id == dag_id, + XCom.task_id == task_id, + XCom.execution_date == execution_date + ).first().value self.assertEqual(ret_value, json_obj) @@ -2247,18 +2257,26 @@ class ClearTasksTest(unittest.TestCase): dag_id = "test_dag2" task_id = "test_task2" + configuration.set("core", "enable_xcom_pickling", "True") + XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, - execution_date=execution_date, - enable_pickling=True) + execution_date=execution_date) ret_value = XCom.get_one(key=key, - dag_id=dag_id, - task_id=task_id, - execution_date=execution_date, - enable_pickling=True) + dag_id=dag_id, + task_id=task_id, + execution_date=execution_date) + + self.assertEqual(ret_value, json_obj) + + session = settings.Session() + ret_value = session.query(XCom).filter(XCom.key == key, XCom.dag_id == dag_id, + XCom.task_id == task_id, + XCom.execution_date == execution_date + ).first().value self.assertEqual(ret_value, json_obj) @@ -2266,13 +2284,15 @@ class ClearTasksTest(unittest.TestCase): class PickleRce(object): def __reduce__(self): return (os.system, ("ls -alt",)) + + configuration.set("core", "xcom_enable_pickling", "False") + self.assertRaises(TypeError, XCom.set, key="xcom_test3", value=PickleRce(), dag_id="test_dag3", task_id="test_task3", - execution_date=timezone.utcnow(), - enable_pickling=False) + execution_date=timezone.utcnow()) def test_xcom_get_many(self): json_obj = {"key": "value"} @@ -2283,23 +2303,22 @@ class ClearTasksTest(unittest.TestCase): dag_id2 = "test_dag5" task_id2 = "test_task5" + configuration.set("core", "xcom_enable_pickling", "True") + XCom.set(key=key, value=json_obj, dag_id=dag_id1, task_id=task_id1, - execution_date=execution_date, - enable_pickling=True) + execution_date=execution_date) XCom.set(key=key, value=json_obj, dag_id=dag_id2, task_id=task_id2, - execution_date=execution_date, - enable_pickling=True) + execution_date=execution_date) results = XCom.get_many(key=key, - execution_date=execution_date, - enable_pickling=True) + execution_date=execution_date) for result in results: self.assertEqual(result.value, json_obj)
