diff --git a/airflow/models.py b/airflow/models.py index 76ff26752a..ba8226a2c8 100644 --- a/airflow/models.py +++ b/airflow/models.py @@ -5,7 +5,7 @@ import imp import jinja2 import logging import os -import pickle +import dill as pickle import re import signal import socket @@ -14,7 +14,6 @@ from sqlalchemy import ( Column, Integer, String, DateTime, Text, Boolean, ForeignKey) from sqlalchemy import func from sqlalchemy.ext.declarative import declarative_base -from sqlalchemy.ext.serializer import loads, dumps from sqlalchemy.dialects.mysql import LONGTEXT from sqlalchemy.orm import relationship @@ -230,10 +229,10 @@ class DagPickle(Base): self.job = job for t in dag.tasks: t.materialize_files() - self.pickle = dumps(dag) + self.pickle = pickle.dumps(dag) def get_object(self): - return loads(self.pickle) + return pickle.loads(self.pickle) class TaskInstance(Base): diff --git a/requirements.txt b/requirements.txt index e3902ce0e0..4a1fa7cef3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ celery chartkick +dill flake8 flask flask-admin diff --git a/setup.py b/setup.py index 0d158e3a53..33a51e1ae9 100644 --- a/setup.py +++ b/setup.py @@ -14,6 +14,7 @@ setup( install_requires=[ 'celery>=3.1.17', 'chartkick>=0.4.2', + 'dill>=0.2.2', 'flask>=0.10.1', 'flask-admin>=1.0.9', 'flask-bootstrap>=3.3.0.1',