[airflow] 02/17: Add explicit session parameter in PoolSlotsAvailableDep (#18875)

[je...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 8666bf0b7eeac887516884687636bbd900479fa0
Author: Aleksey Kirilishin <54...@users.noreply.github.com>
AuthorDate: Fri Oct 29 09:06:28 2021 +0300

    Add explicit session parameter in PoolSlotsAvailableDep (#18875)
    
    (cherry picked from commit e0aa36ead4bb703abf5702bb1c9b105d60c15b28)
---
 airflow/ti_deps/deps/pool_slots_available_dep.py | 2 +-
 tests/jobs/test_local_task_job.py                | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/ti_deps/deps/pool_slots_available_dep.py b/airflow/ti_deps/deps/pool_slots_available_dep.py
index f8c461e..a37e62a 100644
--- a/airflow/ti_deps/deps/pool_slots_available_dep.py
+++ b/airflow/ti_deps/deps/pool_slots_available_dep.py
@@ -55,7 +55,7 @@ class PoolSlotsAvailableDep(BaseTIDep):
         else:
             # Controlled by UNIQUE key in slot_pool table,
             # only one result can be returned.
-            open_slots = pools[0].open_slots()
+            open_slots = pools[0].open_slots(session=session)
 
         if ti.state in EXECUTION_STATES:
             open_slots += ti.pool_slots
diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py
index 663ba43..f23a94f 100644
--- a/tests/jobs/test_local_task_job.py
+++ b/tests/jobs/test_local_task_job.py
@@ -878,5 +878,5 @@ def test_number_of_queries_single_loop(mock_get_task_runner, return_codes, dag_m
     ti.refresh_from_task(task)
 
     job = LocalTaskJob(task_instance=ti, executor=MockExecutor())
-    with assert_queries_count(20):
+    with assert_queries_count(18):
         job.run()