Merge pull request #1219 from airbnb/arthurw_utils_refactor_take2

Refactoring utils to be more sane

airflow/__init__.py

@@ -31,7 +31,7 @@
 from airflow.models import DAG
 from flask_admin import BaseView
 from importlib import import_module
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 
 DAGS_FOLDER = os.path.expanduser(conf.get('core', 'DAGS_FOLDER'))
 if DAGS_FOLDER not in sys.path:

airflow/bin/cli.py

@@ -14,11 +14,14 @@
 import json
 
 import airflow
-from airflow import jobs, settings, utils
+from airflow import jobs, settings
 from airflow import configuration as conf
 from airflow.executors import DEFAULT_EXECUTOR
 from airflow.models import DagModel, DagBag, TaskInstance, DagPickle, DagRun
-from airflow.utils import AirflowException, State
+from airflow.utils import db as db_utils
+from airflow.utils import logging as logging_utils
+from airflow.utils.state import State
+from airflow.exceptions import AirflowException
 
 DAGS_FOLDER = os.path.expanduser(conf.get('core', 'DAGS_FOLDER'))
 
@@ -78,7 +81,8 @@ def backfill(args, dag=None):
             mark_success=args.mark_success,
             include_adhoc=args.include_adhoc,
             local=args.local,
-            donot_pickle=(args.donot_pickle or conf.getboolean('core', 'donot_pickle')),
+            donot_pickle=(args.donot_pickle or
+                          conf.getboolean('core', 'donot_pickle')),
             ignore_dependencies=args.ignore_dependencies,
             pool=args.pool)
 
@@ -133,7 +137,7 @@ def set_is_paused(is_paused, args, dag=None):
 
 def run(args, dag=None):
 
-    utils.pessimistic_connection_handling()
+    db_utils.pessimistic_connection_handling()
     if dag:
         args.dag_id = dag.dag_id
 
@@ -236,10 +240,10 @@ def run(args, dag=None):
         remote_log_location = filename.replace(log_base, remote_base)
         # S3
         if remote_base.startswith('s3:/'):
-            utils.S3Log().write(log, remote_log_location)
+            logging_utils.S3Log().write(log, remote_log_location)
         # GCS
         elif remote_base.startswith('gs:/'):
-            utils.GCSLog().write(
+            logging_utils.GCSLog().write(
                 log,
                 remote_log_location,
                 append=True)
@@ -401,7 +405,7 @@ def worker(args):
 
 def initdb(args):  # noqa
     print("DB: " + repr(settings.engine.url))
-    utils.initdb()
+    db_utils.initdb()
     print("Done.")
 
 
@@ -412,14 +416,14 @@ def resetdb(args):
             "Proceed? (y/n)").upper() == "Y":
         logging.basicConfig(level=settings.LOGGING_LEVEL,
                             format=settings.SIMPLE_LOG_FORMAT)
-        utils.resetdb()
+        db_utils.resetdb()
     else:
         print("Bail.")
 
 
 def upgradedb(args):  # noqa
     print("DB: " + repr(settings.engine.url))
-    utils.upgradedb()
+    db_utils.upgradedb()
 
 
 def version(args):  # noqa

airflow/configuration.py

@@ -16,6 +16,7 @@
 from collections import OrderedDict
 from configparser import ConfigParser
 
+
 class AirflowConfigException(Exception):
     pass
 
@@ -612,6 +613,7 @@ def test_mode():
 def get(section, key, **kwargs):
     return conf.get(section, key, **kwargs)
 
+
 def getboolean(section, key):
     return conf.getboolean(section, key)
 
@@ -644,5 +646,6 @@ def set(section, option, value):  # noqa
 ########################
 # convenience method to access config entries
 
+
 def get_dags_folder():
     return os.path.expanduser(get('core', 'DAGS_FOLDER'))

airflow/contrib/executors/mesos_executor.py

@@ -11,8 +11,8 @@
 from airflow import configuration
 from airflow.executors.base_executor import BaseExecutor
 from airflow.settings import Session
-from airflow.utils import State
-from airflow.utils import AirflowException
+from airflow.utils.state import State
+from airflow.exceptions import AirflowException
 
 
 DEFAULT_FRAMEWORK_NAME = 'Airflow'

airflow/contrib/hooks/__init__.py

@@ -1,6 +1,6 @@
 # Imports the hooks dynamically while keeping the package API clean,
 # abstracting the underlying modules
-from airflow.utils import import_module_attrs as _import_module_attrs
+from airflow.utils.helpers import import_module_attrs as _import_module_attrs
 
 _hooks = {
     'ftp_hook': ['FTPHook'],

airflow/contrib/hooks/gc_base_hook.py

@@ -2,7 +2,7 @@
 import logging
 
 from airflow.hooks.base_hook import BaseHook
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 from oauth2client.client import SignedJwtAssertionCredentials, GoogleCredentials
 
 class GoogleCloudBaseHook(BaseHook):

airflow/contrib/hooks/qubole_hook.py

@@ -3,7 +3,7 @@
 import datetime
 import logging
 
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 from airflow.hooks.base_hook import BaseHook
 from airflow import configuration
 
@@ -151,4 +151,4 @@ def create_cmd_args(self):
             else:
                 args += inplace_args.split(' ')
 
-        return args
+        return args

airflow/contrib/hooks/ssh_hook.py

@@ -20,7 +20,7 @@
 from contextlib import contextmanager
 
 from airflow.hooks.base_hook import BaseHook
-from airflow import AirflowException
+from airflow.exceptions import AirflowException
 
 import logging
 

airflow/contrib/operators/__init__.py

@@ -1,6 +1,6 @@
 # Imports the operators dynamically while keeping the package API clean,
 # abstracting the underlying modules
-from airflow.utils import import_module_attrs as _import_module_attrs
+from airflow.utils.helpers import import_module_attrs as _import_module_attrs
 
 _operators = {
     'ssh_execute_operator': ['SSHExecuteOperator'],

airflow/contrib/operators/bigquery_check_operator.py

@@ -1,6 +1,6 @@
 from airflow.contrib.hooks.bigquery_hook import BigQueryHook
 from airflow.operators import CheckOperator, ValueCheckOperator, IntervalCheckOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 
 class BigQueryCheckOperator(CheckOperator):

airflow/contrib/operators/bigquery_operator.py

@@ -2,7 +2,7 @@
 
 from airflow.contrib.hooks.bigquery_hook import BigQueryHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class BigQueryOperator(BaseOperator):
     """

airflow/contrib/operators/bigquery_to_bigquery.py

@@ -2,7 +2,7 @@
 
 from airflow.contrib.hooks.bigquery_hook import BigQueryHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class BigQueryToBigQueryOperator(BaseOperator):
     """

airflow/contrib/operators/bigquery_to_gcs.py

@@ -2,7 +2,7 @@
 
 from airflow.contrib.hooks.bigquery_hook import BigQueryHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class BigQueryToCloudStorageOperator(BaseOperator):
     """

airflow/contrib/operators/gcs_download_operator.py

@@ -2,7 +2,7 @@
 
 from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class GoogleCloudStorageDownloadOperator(BaseOperator):
     """

airflow/contrib/operators/gcs_to_bq.py

@@ -4,7 +4,7 @@
 from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
 from airflow.contrib.hooks.bigquery_hook import BigQueryHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class GoogleCloudStorageToBigQueryOperator(BaseOperator):
     """

airflow/contrib/operators/mysql_to_gcs.py

@@ -5,7 +5,7 @@
 from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
 from airflow.hooks import MySqlHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 from collections import OrderedDict
 from datetime import date, datetime
 from decimal import Decimal

airflow/contrib/operators/qubole_operator.py

@@ -1,5 +1,5 @@
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 from airflow.contrib.hooks import QuboleHook
 
 

airflow/contrib/operators/ssh_execute_operator.py

@@ -4,8 +4,8 @@
 from subprocess import STDOUT
 
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
-from airflow.utils import AirflowException
+from airflow.utils.decorators import apply_defaults
+from airflow.exceptions import AirflowException
 
 
 class SSHTempFileContent():

airflow/contrib/operators/vertica_operator.py

@@ -2,7 +2,7 @@
 
 from airflow.contrib.hooks import VerticaHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 
 class VerticaOperator(BaseOperator):

airflow/contrib/operators/vertica_to_hive.py

@@ -7,7 +7,7 @@
 from airflow.hooks import HiveCliHook
 from airflow.contrib.hooks import VerticaHook
 from airflow.models import BaseOperator
-from airflow.utils import apply_defaults
+from airflow.utils.decorators import apply_defaults
 
 class VerticaToHiveTransfer(BaseOperator):
     """

airflow/example_dags/example_short_circuit_operator.py

@@ -1,6 +1,6 @@
 from airflow.operators import ShortCircuitOperator, DummyOperator
 from airflow.models import DAG
-import airflow.utils
+import airflow.utils.helpers
 from datetime import datetime, timedelta
 
 seven_days_ago = datetime.combine(datetime.today() - timedelta(7),
@@ -21,5 +21,5 @@
 ds_true = [DummyOperator(task_id='true_' + str(i), dag=dag) for i in [1, 2]]
 ds_false = [DummyOperator(task_id='false_' + str(i), dag=dag) for i in [1, 2]]
 
-airflow.utils.chain(cond_true, *ds_true)
-airflow.utils.chain(cond_false, *ds_false)
+airflow.utils.helpers.chain(cond_true, *ds_true)
+airflow.utils.helpers.chain(cond_false, *ds_false)

airflow/example_dags/example_trigger_controller_dag.py

@@ -1,3 +1,4 @@
+
 """This example illustrates the use of the TriggerDagRunOperator. There are 2
 entities at work in this scenario:
 1. The Controller DAG - the DAG that conditionally executes the trigger
@@ -14,6 +15,7 @@
      state is then made available to the TargetDag
 2. A Target DAG : c.f. example_trigger_target_dag.py
 """
+
 from airflow import DAG
 from airflow.operators import TriggerDagRunOperator
 from datetime import datetime
@@ -35,8 +37,8 @@ def conditionally_trigger(context, dag_run_obj):
 
 # Define the DAG
 dag = DAG(dag_id='example_trigger_controller_dag',
-          default_args={"owner" : "me",
-          "start_date":datetime.now()},
+          default_args={"owner": "me",
+                        "start_date": datetime.now()},
           schedule_interval='@once')
 
 

airflow/example_dags/example_trigger_target_dag.py

@@ -34,7 +34,7 @@
 
 
 def run_this_func(ds, **kwargs):
-    print( "Remotely received value of {} for key=message".format(kwargs['dag_run'].conf['message']))
+    print("Remotely received value of {} for key=message".format(kwargs['dag_run'].conf['message']))
 
 run_this = PythonOperator(
     task_id='run_this',

airflow/exceptions.py

@@ -0,0 +1,10 @@
+class AirflowException(Exception):
+    pass
+
+
+class AirflowSensorTimeout(Exception):
+    pass
+
+
+class AirflowTaskTimeout(Exception):
+    pass

airflow/executors/__init__.py

@@ -10,7 +10,7 @@
 except:
     pass
 
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 
 _EXECUTOR = configuration.get('core', 'EXECUTOR')
 

airflow/executors/base_executor.py

@@ -1,7 +1,8 @@
 from builtins import range
 
 from airflow import configuration
-from airflow.utils import State, LoggingMixin
+from airflow.utils.state import State
+from airflow.utils.logging import LoggingMixin
 
 PARALLELISM = configuration.getint('core', 'PARALLELISM')
 

airflow/executors/celery_executor.py

@@ -6,7 +6,7 @@
 from celery import Celery
 from celery import states as celery_states
 
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 from airflow.executors.base_executor import BaseExecutor
 from airflow import configuration
 

airflow/executors/local_executor.py

@@ -6,7 +6,8 @@
 
 from airflow import configuration
 from airflow.executors.base_executor import BaseExecutor
-from airflow.utils import State, LoggingMixin
+from airflow.utils.state import State
+from airflow.utils.logging import LoggingMixin
 
 PARALLELISM = configuration.get('core', 'PARALLELISM')
 

airflow/executors/sequential_executor.py

@@ -1,9 +1,8 @@
 from builtins import str
-import logging
 import subprocess
 
 from airflow.executors.base_executor import BaseExecutor
-from airflow.utils import State
+from airflow.utils.state import State
 
 
 class SequentialExecutor(BaseExecutor):

airflow/hooks/S3_hook.py

@@ -15,7 +15,7 @@
 boto.set_stream_logger('boto')
 logging.getLogger("boto").setLevel(logging.INFO)
 
-from airflow.utils import AirflowException
+from airflow.exceptions import AirflowException
 from airflow.hooks.base_hook import BaseHook
 
 

airflow/hooks/__init__.py

@@ -1,6 +1,7 @@
 # Imports the hooks dynamically while keeping the package API clean,
 # abstracting the underlying modules
-from airflow.utils import import_module_attrs as _import_module_attrs
+
+from airflow.utils.helpers import import_module_attrs as _import_module_attrs
 from airflow.hooks.base_hook import BaseHook  # noqa to expose in package
 
 _hooks = {