Dispose unused connection pool (apache#21565)

NayaraWakewski · Feb 15, 2022 · 8155e8a · 8155e8a
1 parent 7864693
commit 8155e8a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
@@ -319,11 +319,11 @@ def task_run(args, dag=None):
 
     settings.MASK_SECRETS_IN_LOGS = True
 
-    # IMPORTANT, have to use the NullPool, otherwise, each "run" command may leave
+    # IMPORTANT, have to re-configure ORM with the NullPool, otherwise, each "run" command may leave
     # behind multiple open sleeping connections while heartbeating, which could
     # easily exceed the database connection limit when
     # processing hundreds of simultaneous tasks.
-    settings.configure_orm(disable_connection_pool=True)
+    settings.reconfigure_orm(disable_connection_pool=True)
 
     if args.pickle:
         print(f'Loading pickle id: {args.pickle}')

diff --git a/airflow/settings.py b/airflow/settings.py
@@ -380,6 +380,12 @@ def dispose_orm():
         engine = None
 
 
+def reconfigure_orm(disable_connection_pool=False):
+    """Properly close database connections and re-configure ORM"""
+    dispose_orm()
+    configure_orm(disable_connection_pool=disable_connection_pool)
+
+
 def configure_adapters():
     """Register Adapters and DB Converters"""
     from pendulum import DateTime as Pendulum