Merge pull request #37 from mistercrunch/pessimistic

Pessimistic pool connection handling for master and run

airflow/bin/airflow

@@ -43,6 +43,8 @@ def backfill(args):
 
 def run(args):
 
+    settings.pessimistic_connection_handling()
+
     # Setting up logging
     directory = getconf().get('core', 'BASE_LOG_FOLDER') + \
         "/{args.dag_id}/{args.task_id}".format(args=args)
@@ -145,6 +147,7 @@ def webserver(args):
 
 
 def master(args):
+    settings.pessimistic_connection_handling()
 
     # Sleep time (seconds) between master runs
     MASTER_SLEEP_INTERVAL = 60

airflow/settings.py

@@ -2,10 +2,12 @@ import sys
 
 from sqlalchemy.orm import scoped_session, sessionmaker
 from sqlalchemy import create_engine
+from sqlalchemy import event
+from sqlalchemy import exc
+from sqlalchemy.pool import Pool
 
 from airflow.configuration import getconf
 
-
 HEADER = """\
        .__         _____.__
 _____  |__|_______/ ____\  |   ______  _  __
@@ -14,6 +16,21 @@ _____  |__|_______/ ____\  |   ______  _  __
 (____  /__||__|   |__|  |____/\____/ \/\_/
      \/"""
 
+def pessimistic_connection_handling():
+    @event.listens_for(Pool, "checkout")
+    def ping_connection(dbapi_connection, connection_record, connection_proxy):
+        '''
+        Disconnect Handling - Pessimistic, taken from:
+        http://docs.sqlalchemy.org/en/rel_0_9/core/pooling.html
+        '''
+        cursor = dbapi_connection.cursor()
+        try:
+            cursor.execute("SELECT 1")
+        except:
+            raise exc.DisconnectionError()
+        cursor.close()
+
+
 BASE_FOLDER = getconf().get('core', 'BASE_FOLDER')
 SQL_ALCHEMY_CONN = getconf().get('core', 'SQL_ALCHEMY_CONN')
 if BASE_FOLDER not in sys.path:

dags/examples/example1.py

@@ -28,7 +28,7 @@ for i in range(9):
     i = str(i)
     task = BashOperator(
             task_id='runme_'+i,
-            bash_command='sleep {{ 10 + macros.random() * 10 }}',
+            bash_command='sleep 20',
             **default_args)
     task.set_downstream(run_this)
     dag.add_task(task)