You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ep...@apache.org on 2023/08/09 21:07:06 UTC
[airflow] branch v2-7-test updated (2c43ca3ede -> 803d0134ac)
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a change to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
omit 2c43ca3ede Update RELEASE_NOTES.rst
new a57cc8b0c0 Fix missing related PRs when cherry-picking (#33261)
new 494edad3e9 Enable pools to consider deferred tasks (#32709)
new e9c82d9319 wOrkaround failing dedlock when running backfill (#32991)
new 4e817c46b4 Fix flaky sqlite tests with `test_xcom_map_nest` hopefully (#33145)
new 4e00fe8ec0 Refactor: Simplify code in dag_processing (#33161)
new 14ee535270 Gantt chart: Use earliest/oldest ti dates if different than dag run start/end (#33215)
new b099803893 Remove add_task from TaskGroup (#33262)
new 2ff8d4d613 Change type of defualts for include_deferred (#33271)
new 803d0134ac Update RELEASE_NOTES.rst
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (2c43ca3ede)
\
N -- N -- N refs/heads/v2-7-test (803d0134ac)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 9 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
RELEASE_NOTES.rst | 3 +
airflow/api/client/api_client.py | 3 +-
airflow/api/client/json_client.py | 9 +-
airflow/api/client/local_client.py | 10 +-
airflow/api/common/experimental/pool.py | 2 +-
airflow/api_connexion/endpoints/pool_endpoint.py | 4 +-
airflow/api_connexion/openapi/v1.yaml | 15 ++-
airflow/api_connexion/schemas/pool_schema.py | 8 ++
airflow/cli/cli_config.py | 16 ++-
airflow/cli/commands/pool_command.py | 18 +++-
airflow/dag_processing/manager.py | 8 +-
airflow/dag_processing/processor.py | 12 +--
airflow/jobs/backfill_job_runner.py | 117 ++++++++++++--------
airflow/jobs/scheduler_job_runner.py | 2 +
...8_2_7_0_add_include_deferred_column_to_pool.py} | 30 +++---
airflow/models/abstractoperator.py | 8 --
airflow/models/pool.py | 58 ++++++++--
airflow/models/taskmixin.py | 5 -
airflow/models/xcom_arg.py | 10 --
airflow/ti_deps/deps/pool_slots_available_dep.py | 11 +-
airflow/utils/db.py | 1 +
airflow/utils/edgemodifier.py | 9 --
airflow/utils/task_group.py | 20 ----
.../static/js/dag/details/gantt/GanttTooltip.tsx | 3 +
airflow/www/static/js/dag/details/gantt/Row.tsx | 20 ++--
airflow/www/static/js/dag/details/gantt/index.tsx | 32 +++++-
airflow/www/static/js/types/api-generated.ts | 14 ++-
airflow/www/views.py | 29 ++++-
dev/airflow-github | 2 +-
.../logging-monitoring/metrics.rst | 1 +
.../administration-and-deployment/pools.rst | 2 +-
.../authoring-and-scheduling/deferring.rst | 2 +-
docs/apache-airflow/img/airflow_erd.sha256 | 2 +-
docs/apache-airflow/img/airflow_erd.svg | 86 ++++++++-------
docs/apache-airflow/migrations-ref.rst | 4 +-
.../api_connexion/endpoints/test_pool_endpoint.py | 120 ++++++++++++++++-----
tests/api_connexion/schemas/test_pool_schemas.py | 14 ++-
tests/api_connexion/test_auth.py | 4 +
tests/api_experimental/client/test_local_client.py | 20 ++--
.../common/experimental/test_pool.py | 7 +-
tests/cli/commands/test_pool_command.py | 67 ++++++++----
tests/cli/commands/test_task_command.py | 2 +-
tests/jobs/test_backfill_job.py | 3 +-
tests/jobs/test_scheduler_job.py | 31 +++---
tests/models/test_pool.py | 70 +++++++++++-
tests/models/test_taskinstance.py | 2 +-
tests/models/test_xcom_arg_map.py | 4 +-
tests/operators/test_subdag_operator.py | 8 +-
tests/providers/daskexecutor/test_dask_executor.py | 1 -
.../ti_deps/deps/test_pool_slots_available_dep.py | 13 ++-
tests/www/api/experimental/test_endpoints.py | 2 +
tests/www/views/test_views_pool.py | 5 +
52 files changed, 646 insertions(+), 303 deletions(-)
copy airflow/migrations/versions/{0126_2_7_0_add_index_to_task_instance_table.py => 0128_2_7_0_add_include_deferred_column_to_pool.py} (63%)
[airflow] 03/09: wOrkaround failing dedlock when running backfill (#32991)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e9c82d9319c0f27d40f49aca612eb1713e06f4ea
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Wed Aug 9 13:00:48 2023 +0200
wOrkaround failing dedlock when running backfill (#32991)
The dask_executor backfill tests started to fail recently more often due
to backfill exception, and the likely cause for it is that it is now
better parallelise execution and triggering of the deadlocks because of
contention betwee dag_run state update and task state update had
become much easier.
While this PR does not fix the underlying issue, it catches the
operational error where the deadlock occured during the backfill.
and rolls back the operation.
This **should** be safe. backfil has a built-in mechanism to loop and
retry failed tasks and the test passed multiple times, completing the
backfill after this fix was applied. It was not easy to reproduce it
locally but it failed every 20-30 times. When extra logging was added,
it was always connected to OperationalException raised (and caught)
right after _per_task_process. The same exception was observed few times
when rollback was added, and despite it backfill job retried and
completed the process successfully every time. We also leave the logs
with exceptions and add reassuring messages that should make it clear
that in case backfill completes, the exceptions can be ignored as
the updates will be retried by the backfill job.
Fixes: #32778
(cherry picked from commit f616ee8dff8e6ba9b37cbce0d22235dc47d4fc93)
---
airflow/jobs/backfill_job_runner.py | 117 +++++++++++++--------
tests/providers/daskexecutor/test_dask_executor.py | 1 -
2 files changed, 73 insertions(+), 45 deletions(-)
diff --git a/airflow/jobs/backfill_job_runner.py b/airflow/jobs/backfill_job_runner.py
index 162f103642..3d2c20c612 100644
--- a/airflow/jobs/backfill_job_runner.py
+++ b/airflow/jobs/backfill_job_runner.py
@@ -588,61 +588,83 @@ class BackfillJobRunner(BaseJobRunner[Job], LoggingMixin):
try:
for task in self.dag.topological_sort(include_subdag_tasks=True):
for key, ti in list(ti_status.to_run.items()):
- if task.task_id != ti.task_id:
- continue
-
- pool = session.scalar(
- select(models.Pool).where(models.Pool.pool == task.pool).limit(1)
- )
- if not pool:
- raise PoolNotFound(f"Unknown pool: {task.pool}")
-
- open_slots = pool.open_slots(session=session)
- if open_slots <= 0:
- raise NoAvailablePoolSlot(
- f"Not scheduling since there are {open_slots} open slots in pool {task.pool}"
- )
-
- num_running_task_instances_in_dag = DAG.get_num_task_instances(
- self.dag_id,
- states=self.STATES_COUNT_AS_RUNNING,
- session=session,
- )
-
- if num_running_task_instances_in_dag >= self.dag.max_active_tasks:
- raise DagConcurrencyLimitReached(
- "Not scheduling since DAG max_active_tasks limit is reached."
+ # Attempt to workaround deadlock on backfill by attempting to commit the transaction
+ # state update few times before giving up
+ max_attempts = 5
+ for i in range(max_attempts):
+ if task.task_id != ti.task_id:
+ continue
+
+ pool = session.scalar(
+ select(models.Pool).where(models.Pool.pool == task.pool).limit(1)
)
+ if not pool:
+ raise PoolNotFound(f"Unknown pool: {task.pool}")
+
+ open_slots = pool.open_slots(session=session)
+ if open_slots <= 0:
+ raise NoAvailablePoolSlot(
+ f"Not scheduling since there are {open_slots} "
+ f"open slots in pool {task.pool}"
+ )
- if task.max_active_tis_per_dag is not None:
- num_running_task_instances_in_task = DAG.get_num_task_instances(
- dag_id=self.dag_id,
- task_ids=[task.task_id],
+ num_running_task_instances_in_dag = DAG.get_num_task_instances(
+ self.dag_id,
states=self.STATES_COUNT_AS_RUNNING,
session=session,
)
- if num_running_task_instances_in_task >= task.max_active_tis_per_dag:
- raise TaskConcurrencyLimitReached(
- "Not scheduling since Task concurrency limit is reached."
+ if num_running_task_instances_in_dag >= self.dag.max_active_tasks:
+ raise DagConcurrencyLimitReached(
+ "Not scheduling since DAG max_active_tasks limit is reached."
)
- if task.max_active_tis_per_dagrun is not None:
- num_running_task_instances_in_task_dagrun = DAG.get_num_task_instances(
- dag_id=self.dag_id,
- run_id=ti.run_id,
- task_ids=[task.task_id],
- states=self.STATES_COUNT_AS_RUNNING,
- session=session,
- )
+ if task.max_active_tis_per_dag is not None:
+ num_running_task_instances_in_task = DAG.get_num_task_instances(
+ dag_id=self.dag_id,
+ task_ids=[task.task_id],
+ states=self.STATES_COUNT_AS_RUNNING,
+ session=session,
+ )
- if num_running_task_instances_in_task_dagrun >= task.max_active_tis_per_dagrun:
- raise TaskConcurrencyLimitReached(
- "Not scheduling since Task concurrency per DAG run limit is reached."
+ if num_running_task_instances_in_task >= task.max_active_tis_per_dag:
+ raise TaskConcurrencyLimitReached(
+ "Not scheduling since Task concurrency limit is reached."
+ )
+
+ if task.max_active_tis_per_dagrun is not None:
+ num_running_task_instances_in_task_dagrun = DAG.get_num_task_instances(
+ dag_id=self.dag_id,
+ run_id=ti.run_id,
+ task_ids=[task.task_id],
+ states=self.STATES_COUNT_AS_RUNNING,
+ session=session,
)
- _per_task_process(key, ti, session)
- session.commit()
+ if (
+ num_running_task_instances_in_task_dagrun
+ >= task.max_active_tis_per_dagrun
+ ):
+ raise TaskConcurrencyLimitReached(
+ "Not scheduling since Task concurrency per DAG run limit is reached."
+ )
+
+ _per_task_process(key, ti, session)
+ try:
+ session.commit()
+ # break the retry loop
+ break
+ except OperationalError:
+ self.log.error(
+ "Failed to commit task state due to operational error. "
+ "The job will retry this operation so if your backfill succeeds, "
+ "you can safely ignore this message.",
+ exc_info=True,
+ )
+ session.rollback()
+ if i == max_attempts - 1:
+ raise
+ # retry the loop
except (NoAvailablePoolSlot, DagConcurrencyLimitReached, TaskConcurrencyLimitReached) as e:
self.log.debug(e)
@@ -939,6 +961,13 @@ class BackfillJobRunner(BaseJobRunner[Job], LoggingMixin):
# TODO: we will need to terminate running task instances and set the
# state to failed.
self._set_unfinished_dag_runs_to_failed(ti_status.active_runs)
+ except OperationalError:
+ self.log.error(
+ "Backfill job dead-locked. The job will retry the job so it is likely "
+ "to heal itself. If your backfill succeeds you can ignore this exception.",
+ exc_info=True,
+ )
+ raise
finally:
session.commit()
executor.end()
diff --git a/tests/providers/daskexecutor/test_dask_executor.py b/tests/providers/daskexecutor/test_dask_executor.py
index 2d559eaa40..c5773bd83d 100644
--- a/tests/providers/daskexecutor/test_dask_executor.py
+++ b/tests/providers/daskexecutor/test_dask_executor.py
@@ -104,7 +104,6 @@ class TestDaskExecutor(TestBaseDask):
# This test is quarantined because it became rather flaky on our CI in July 2023 and reason for this
# is unknown. An issue for that was created: https://github.com/apache/airflow/issues/32778 and the
# marker should be removed while (possibly) the reason for flaky behaviour is found and fixed.
- @pytest.mark.quarantined
@pytest.mark.execution_timeout(180)
def test_backfill_integration(self):
"""
[airflow] 09/09: Update RELEASE_NOTES.rst
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 803d0134ac81427154b9b8521c48a171d5aee0c3
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Aug 9 07:55:58 2023 +0100
Update RELEASE_NOTES.rst
---
RELEASE_NOTES.rst | 25 +++++++++++++++++++++++--
1 file changed, 23 insertions(+), 2 deletions(-)
diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst
index 24c927764f..e897d038ab 100644
--- a/RELEASE_NOTES.rst
+++ b/RELEASE_NOTES.rst
@@ -128,6 +128,7 @@ New Features
- Trigger Button - Implement Part 2 of AIP-50 (#31583)
- Automatic setup and teardown tasks (`AIP-52 <https://github.com/apache/airflow/pulls?q=is%3Apr+is%3Amerged+label%3AAIP-52+milestone%3A%22Airflow+2.7.0%22>`_)
- OpenLineage in Airflow (`AIP-53 <https://github.com/apache/airflow/pulls?q=is%3Apr+is%3Amerged+label%3AAIP-53+milestone%3A%22Airflow+2.7.0%22>`_)
+- Enable pools to consider deferred tasks (#32709)
- Allows to choose SSL context for SMTP connection (#33070)
- New gantt tab (#31806)
- Load plugins from providers (#32692)
@@ -165,11 +166,14 @@ New Features
Improvements
""""""""""""
+- Configurable health check threshold for triggerer (#33089, #33084)
+- Add documentation generation for CLI commands from executors (#33081, #33143, #33081)
+- add dag_run_ids and task_ids filter for the batch task instance API endpoint (#32705)
- Ensure DAG-level references are filled on unmap (#33083)
- Add support for arrays of different data types in the Trigger Form UI (#32734)
- Always show gantt and code tabs (#33029)
- Move listener success hook to after SQLAlchemy commit (#32988)
-- Rename ``db upgrade`` to ``db migrate`` and add ``connections create-default-connections`` (#32810)
+- Rename ``db upgrade`` to ``db migrate`` and add ``connections create-default-connections`` (#32810, #33136)
- Remove old gantt chart and redirect to grid views gantt tab (#32908)
- Adjust graph zoom based on selected task (#32792)
- Call listener on_task_instance_running after rendering templates (#32716)
@@ -214,6 +218,13 @@ Improvements
Bug Fixes
"""""""""
+- ``Gantt chart:`` Use earliest/oldest ti dates if different than dag run start/end (#33215)
+- Fix ``virtualenv`` detection for Python ``virtualenv`` operator (#33223)
+- Correctly log when there are problems trying to ``chmod`` ``airflow.cfg`` (#33118)
+- Pass app context to webserver_config.py (#32759)
+- Skip served logs for non-running task try (#32561)
+- Fix reload gunicorn workers (#32102)
+- Fix future DagRun rarely triggered by race conditions when ``max_active_runs`` reached its upper limit. (#31414)
- Fix BaseOperator ``get_task_instances`` query (#33054)
- Fix issue with using the various state enum value in logs (#33065)
- Use string concatenation to prepend base URL for log_url (#33063)
@@ -262,6 +273,15 @@ Bug Fixes
Misc/Internal
"""""""""""""
+- Refactor: Simplify code in ``dag_processing`` (#33161)
+- For now limit ``Pydantic`` to ``< 2.0.0`` (#33235)
+- Refactor: Simplify code in models (#33181)
+- Add elasticsearch group to pre-2.7 defaults (#33166)
+- Refactor: Simplify dict manipulation in airflow/cli (#33159)
+- Remove redundant dict.keys() call (#33158)
+- Upgrade ruff to latest 0.0.282 version in pre-commits (#33152)
+- Move openlineage configuration to provider (#33124)
+- Replace State by TaskInstanceState in Airflow executors (#32627)
- Get rid of Python 2 numeric relics (#33050)
- Remove legacy dag code (#33058)
- Remove legacy task instance modal (#33060)
@@ -304,7 +324,8 @@ Misc/Internal
Docs only changes
"""""""""""""""""
-- Documentation change to enhance readability (#33006)
+- Add links to ``DAGRun / DAG / Task`` in templates-ref.rst (#33013)
+- Add example for list-import-errors as a CI check (#32811)
- Clean-up of our new security page (#32951)
- Cleans up Extras reference page (#32954)
- Update Dag trigger API and command docs (#32696)
[airflow] 06/09: Gantt chart: Use earliest/oldest ti dates if different than dag run start/end (#33215)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 14ee53527043978a42a1747e53edf81f40907e17
Author: Brent Bovenzi <br...@astronomer.io>
AuthorDate: Wed Aug 9 09:37:54 2023 -0400
Gantt chart: Use earliest/oldest ti dates if different than dag run start/end (#33215)
(cherry picked from commit 1dcf05f1c69e52667064259308ba5950b594b268)
---
.../static/js/dag/details/gantt/GanttTooltip.tsx | 3 ++
airflow/www/static/js/dag/details/gantt/Row.tsx | 20 ++++++++++----
airflow/www/static/js/dag/details/gantt/index.tsx | 32 ++++++++++++++++++++--
3 files changed, 46 insertions(+), 9 deletions(-)
diff --git a/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx b/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
index 2f7d617c8c..2f50120925 100644
--- a/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
+++ b/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
@@ -42,6 +42,9 @@ const GanttTooltip = ({ task, instance }: Props) => {
<Text>
Task{isGroup ? " Group" : ""}: {task.label}
</Text>
+ {!!instance?.tryNumber && instance.tryNumber > 1 && (
+ <Text>Try Number: {instance.tryNumber}</Text>
+ )}
<br />
{instance?.queuedDttm && (
<Text>
diff --git a/airflow/www/static/js/dag/details/gantt/Row.tsx b/airflow/www/static/js/dag/details/gantt/Row.tsx
index 96f079154e..304fb83e22 100644
--- a/airflow/www/static/js/dag/details/gantt/Row.tsx
+++ b/airflow/www/static/js/dag/details/gantt/Row.tsx
@@ -24,24 +24,31 @@ import { getDuration } from "src/datetime_utils";
import { SimpleStatus } from "src/dag/StatusBox";
import { useContainerRef } from "src/context/containerRef";
import { hoverDelay } from "src/utils";
-import type { DagRun, Task } from "src/types";
+import type { Task } from "src/types";
import GanttTooltip from "./GanttTooltip";
interface Props {
ganttWidth?: number;
openGroupIds: string[];
- dagRun: DagRun;
task: Task;
+ ganttStartDate?: string | null;
+ ganttEndDate?: string | null;
}
-const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
+const Row = ({
+ ganttWidth = 500,
+ openGroupIds,
+ task,
+ ganttStartDate,
+ ganttEndDate,
+}: Props) => {
const {
selected: { runId, taskId },
onSelect,
} = useSelection();
const containerRef = useContainerRef();
- const runDuration = getDuration(dagRun?.startDate, dagRun?.endDate);
+ const runDuration = getDuration(ganttStartDate, ganttEndDate);
const instance = task.instances.find((ti) => ti.runId === runId);
const isSelected = taskId === instance?.taskId;
@@ -54,7 +61,7 @@ const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
? getDuration(instance?.queuedDttm, instance?.startDate)
: 0;
const taskStartOffset = getDuration(
- dagRun.startDate,
+ ganttStartDate,
instance?.queuedDttm || instance?.startDate
);
@@ -127,7 +134,8 @@ const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
<Row
ganttWidth={ganttWidth}
openGroupIds={openGroupIds}
- dagRun={dagRun}
+ ganttStartDate={ganttStartDate}
+ ganttEndDate={ganttEndDate}
task={c}
key={`gantt-${c.id}`}
/>
diff --git a/airflow/www/static/js/dag/details/gantt/index.tsx b/airflow/www/static/js/dag/details/gantt/index.tsx
index 0e4b7adb84..83200ac2ff 100644
--- a/airflow/www/static/js/dag/details/gantt/index.tsx
+++ b/airflow/www/static/js/dag/details/gantt/index.tsx
@@ -106,10 +106,35 @@ const Gantt = ({ openGroupIds, gridScrollRef, ganttScrollRef }: Props) => {
const dagRun = dagRuns.find((dr) => dr.runId === runId);
- const startDate = dagRun?.startDate;
+ let startDate = dagRun?.queuedAt || dagRun?.startDate;
+ let endDate = dagRun?.endDate;
+
+ // Check if any task instance dates are outside the bounds of the dag run dates and update our min start and max end
+ groups.children?.forEach((task) => {
+ const taskInstance = task.instances.find((ti) => ti.runId === runId);
+ if (
+ taskInstance?.queuedDttm &&
+ (!startDate ||
+ Date.parse(taskInstance.queuedDttm) < Date.parse(startDate))
+ ) {
+ startDate = taskInstance.queuedDttm;
+ } else if (
+ taskInstance?.startDate &&
+ (!startDate || Date.parse(taskInstance.startDate) < Date.parse(startDate))
+ ) {
+ startDate = taskInstance.startDate;
+ }
+
+ if (
+ taskInstance?.endDate &&
+ (!endDate || Date.parse(taskInstance.endDate) > Date.parse(endDate))
+ ) {
+ endDate = taskInstance.endDate;
+ }
+ });
const numBars = Math.round(width / 100);
- const runDuration = getDuration(dagRun?.startDate, dagRun?.endDate);
+ const runDuration = getDuration(startDate, endDate);
const intervals = runDuration / numBars;
return (
@@ -166,8 +191,9 @@ const Gantt = ({ openGroupIds, gridScrollRef, ganttScrollRef }: Props) => {
<Row
ganttWidth={width}
openGroupIds={openGroupIds}
- dagRun={dagRun}
task={c}
+ ganttStartDate={startDate}
+ ganttEndDate={endDate}
key={`gantt-${c.id}`}
/>
))}
[airflow] 05/09: Refactor: Simplify code in dag_processing (#33161)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 4e00fe8ec01c9a2bbfea02ac7841d42eb45ea7f5
Author: Miroslav Šedivý <67...@users.noreply.github.com>
AuthorDate: Wed Aug 9 18:49:28 2023 +0000
Refactor: Simplify code in dag_processing (#33161)
(cherry picked from commit 2f16a46db51b67e67b8406fbafe9e9add6eee235)
---
airflow/dag_processing/manager.py | 8 ++++----
airflow/dag_processing/processor.py | 12 ++++--------
2 files changed, 8 insertions(+), 12 deletions(-)
diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py
index 7e44e89816..dd252db9bd 100644
--- a/airflow/dag_processing/manager.py
+++ b/airflow/dag_processing/manager.py
@@ -1001,7 +1001,7 @@ class DagFileProcessorManager(LoggingMixin):
processor.terminate()
self._file_stats.pop(file_path)
- to_remove = set(self._file_stats.keys()) - set(self._file_paths)
+ to_remove = set(self._file_stats).difference(self._file_paths)
for key in to_remove:
# Remove the stats for any dag files that don't exist anymore
del self._file_stats[key]
@@ -1076,7 +1076,7 @@ class DagFileProcessorManager(LoggingMixin):
while self._parallelism - len(self._processors) > 0 and self._file_path_queue:
file_path = self._file_path_queue.popleft()
# Stop creating duplicate processor i.e. processor with the same filepath
- if file_path in self._processors.keys():
+ if file_path in self._processors:
continue
callback_to_execute_for_file = self._callback_to_execute[file_path]
@@ -1115,7 +1115,7 @@ class DagFileProcessorManager(LoggingMixin):
self._parsing_start_time = time.perf_counter()
# If the file path is already being processed, or if a file was
# processed recently, wait until the next batch
- file_paths_in_progress = self._processors.keys()
+ file_paths_in_progress = set(self._processors)
now = timezone.utcnow()
# Sort the file paths by the parsing order mode
@@ -1173,7 +1173,7 @@ class DagFileProcessorManager(LoggingMixin):
file_path for file_path, stat in self._file_stats.items() if stat.run_count == self._max_runs
]
- file_paths_to_exclude = set(file_paths_in_progress).union(
+ file_paths_to_exclude = file_paths_in_progress.union(
file_paths_recently_processed,
files_paths_at_run_limit,
)
diff --git a/airflow/dag_processing/processor.py b/airflow/dag_processing/processor.py
index 369f676878..8b0ecb3640 100644
--- a/airflow/dag_processing/processor.py
+++ b/airflow/dag_processing/processor.py
@@ -573,9 +573,9 @@ class DagFileProcessor(LoggingMixin):
for task in tasks_missed_sla:
if task.email:
if isinstance(task.email, str):
- emails |= set(get_email_address_list(task.email))
+ emails.update(get_email_address_list(task.email))
elif isinstance(task.email, (list, tuple)):
- emails |= set(task.email)
+ emails.update(task.email)
if emails:
try:
send_email(emails, f"[airflow] SLA miss on DAG={dag.dag_id}", email_content)
@@ -652,9 +652,7 @@ class DagFileProcessor(LoggingMixin):
task_pools = {task.pool for task in dag.tasks}
nonexistent_pools = task_pools - pools
if nonexistent_pools:
- return (
- f"Dag '{dag.dag_id}' references non-existent pools: {list(sorted(nonexistent_pools))!r}"
- )
+ return f"Dag '{dag.dag_id}' references non-existent pools: {sorted(nonexistent_pools)!r}"
pools = {p.pool for p in Pool.get_pools(session)}
for dag in dagbag.dags.values():
@@ -679,9 +677,7 @@ class DagFileProcessor(LoggingMixin):
"""
self._validate_task_pools(dagbag=dagbag)
- stored_warnings = set(
- session.query(DagWarning).filter(DagWarning.dag_id.in_(dagbag.dags.keys())).all()
- )
+ stored_warnings = set(session.query(DagWarning).filter(DagWarning.dag_id.in_(dagbag.dags)).all())
for warning_to_delete in stored_warnings - self.dag_warnings:
session.delete(warning_to_delete)
[airflow] 07/09: Remove add_task from TaskGroup (#33262)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit b0998038936d290e1ca5e237930c884fec1ce9e4
Author: Daniel Standish <15...@users.noreply.github.com>
AuthorDate: Wed Aug 9 11:07:27 2023 -0700
Remove add_task from TaskGroup (#33262)
This was added during work for AIP-52, taking inspiration from the method on the setupteardowncontext object. But it causes problems because it's assumed in the task_id label logic that if the group is set to prefix task ids then the task id has been prefixed. This results in bad graph labeling in UI. Rather than fix it now, at release time, better to revert.
(cherry picked from commit 2d2a1d699b8069f59604c6d6ea3e29853faa7945)
---
airflow/models/abstractoperator.py | 8 --------
airflow/models/taskmixin.py | 5 -----
airflow/models/xcom_arg.py | 10 ----------
airflow/utils/edgemodifier.py | 9 ---------
airflow/utils/task_group.py | 20 --------------------
5 files changed, 52 deletions(-)
diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index 11e9184735..458d55ef5d 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -360,14 +360,6 @@ class AbstractOperator(Templater, DAGNode):
yield parent
parent = parent.task_group
- def add_to_taskgroup(self, task_group: TaskGroup) -> None:
- """Add the task to the given task group.
-
- :meta private:
- """
- if self.node_id not in task_group.children:
- task_group.add(self)
-
def get_closest_mapped_task_group(self) -> MappedTaskGroup | None:
"""Get the mapped task group "closest" to this task in the DAG.
diff --git a/airflow/models/taskmixin.py b/airflow/models/taskmixin.py
index 8c19749104..c5e7cea3b9 100644
--- a/airflow/models/taskmixin.py
+++ b/airflow/models/taskmixin.py
@@ -113,11 +113,6 @@ class DependencyMixin:
self.__rshift__(other)
return self
- @abstractmethod
- def add_to_taskgroup(self, task_group: TaskGroup) -> None:
- """Add the task to the given task group."""
- raise NotImplementedError()
-
@classmethod
def _iter_references(cls, obj: Any) -> Iterable[tuple[DependencyMixin, str]]:
from airflow.models.baseoperator import AbstractOperator
diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index e8cde16db8..0a54d18ff7 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -43,7 +43,6 @@ from airflow.utils.xcom import XCOM_RETURN_KEY
if TYPE_CHECKING:
from airflow.models.dag import DAG
from airflow.models.operator import Operator
- from airflow.utils.task_group import TaskGroup
# Callable objects contained by MapXComArg. We only accept callables from
# the user, but deserialize them into strings in a serialized XComArg for
@@ -209,15 +208,6 @@ class XComArg(ResolveMixin, DependencyMixin):
"""
raise NotImplementedError()
- def add_to_taskgroup(self, task_group: TaskGroup) -> None:
- """Add the task to the given task group.
-
- :meta private:
- """
- for op, _ in self.iter_references():
- if op.node_id not in task_group.children:
- task_group.add(op)
-
def __enter__(self):
if not self.operator.is_setup and not self.operator.is_teardown:
raise AirflowException("Only setup/teardown tasks can be used as context managers.")
diff --git a/airflow/utils/edgemodifier.py b/airflow/utils/edgemodifier.py
index 14aa3ef25e..e760beb50a 100644
--- a/airflow/utils/edgemodifier.py
+++ b/airflow/utils/edgemodifier.py
@@ -169,15 +169,6 @@ class EdgeModifier(DependencyMixin):
"""
dag.set_edge_info(upstream_id, downstream_id, {"label": self.label})
- def add_to_taskgroup(self, task_group: TaskGroup) -> None:
- """No-op, since we're not a task.
-
- We only add tasks to TaskGroups and not EdgeModifiers, but we need
- this to satisfy the interface.
-
- :meta private:
- """
-
# Factory functions
def Label(label: str):
diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py
index b6c40a14a9..deee607dd1 100644
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -551,26 +551,6 @@ class TaskGroup(DAGNode):
f"Encountered a DAGNode that is not a TaskGroup or an AbstractOperator: {type(child)}"
)
- def add_task(self, task: AbstractOperator) -> None:
- """Add a task to the task group.
-
- :param task: the task to add
- """
- if not TaskGroupContext.active:
- raise AirflowException(
- "Using this method on a task group that's not a context manager is not supported."
- )
- task.add_to_taskgroup(self)
-
- def add_to_taskgroup(self, task_group: TaskGroup) -> None:
- """No-op, since we're not a task.
-
- We only add tasks to TaskGroups and not TaskGroup, but we need
- this to satisfy the interface.
-
- :meta private:
- """
-
class MappedTaskGroup(TaskGroup):
"""A mapped task group.
[airflow] 02/09: Enable pools to consider deferred tasks (#32709)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 494edad3e9241fbc4e6500a7fb792e04f3896493
Author: Usiel Riedl <us...@gmail.com>
AuthorDate: Wed Aug 9 20:34:37 2023 +0800
Enable pools to consider deferred tasks (#32709)
* Makes pools respect deferrable tasks (with extra setting)
See https://github.com/apache/airflow/discussions/21243
This commit makes pools consider deferred tasks if the `include_deferred` flag is set. By default a pool will not consider deferred tasks as occupied slots, but still show the number of deferred tasks in its stats.
---------
Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
(cherry picked from commit 70a050b1c3a46e3eabae7f6052c95718800dc604)
---
airflow/api/client/api_client.py | 3 +-
airflow/api/client/json_client.py | 9 +-
airflow/api/client/local_client.py | 10 +-
airflow/api/common/experimental/pool.py | 2 +-
airflow/api_connexion/endpoints/pool_endpoint.py | 4 +-
airflow/api_connexion/openapi/v1.yaml | 15 ++-
airflow/api_connexion/schemas/pool_schema.py | 8 ++
airflow/cli/cli_config.py | 16 ++-
airflow/cli/commands/pool_command.py | 18 +++-
airflow/jobs/scheduler_job_runner.py | 2 +
...28_2_7_0_add_include_deferred_column_to_pool.py | 50 +++++++++
airflow/models/pool.py | 58 ++++++++--
airflow/ti_deps/deps/pool_slots_available_dep.py | 11 +-
airflow/utils/db.py | 1 +
airflow/www/static/js/types/api-generated.ts | 14 ++-
airflow/www/views.py | 29 ++++-
.../logging-monitoring/metrics.rst | 1 +
.../administration-and-deployment/pools.rst | 2 +-
.../authoring-and-scheduling/deferring.rst | 2 +-
docs/apache-airflow/img/airflow_erd.sha256 | 2 +-
docs/apache-airflow/img/airflow_erd.svg | 86 ++++++++-------
docs/apache-airflow/migrations-ref.rst | 4 +-
.../api_connexion/endpoints/test_pool_endpoint.py | 120 ++++++++++++++++-----
tests/api_connexion/schemas/test_pool_schemas.py | 14 ++-
tests/api_connexion/test_auth.py | 4 +
tests/api_experimental/client/test_local_client.py | 20 ++--
.../common/experimental/test_pool.py | 7 +-
tests/cli/commands/test_pool_command.py | 67 ++++++++----
tests/cli/commands/test_task_command.py | 2 +-
tests/jobs/test_backfill_job.py | 3 +-
tests/jobs/test_scheduler_job.py | 31 +++---
tests/models/test_pool.py | 70 +++++++++++-
tests/models/test_taskinstance.py | 2 +-
tests/operators/test_subdag_operator.py | 8 +-
.../ti_deps/deps/test_pool_slots_available_dep.py | 13 ++-
tests/www/api/experimental/test_endpoints.py | 2 +
tests/www/views/test_views_pool.py | 5 +
37 files changed, 549 insertions(+), 166 deletions(-)
diff --git a/airflow/api/client/api_client.py b/airflow/api/client/api_client.py
index 1fdbc7027b..d4a6d2357a 100644
--- a/airflow/api/client/api_client.py
+++ b/airflow/api/client/api_client.py
@@ -60,12 +60,13 @@ class Client:
"""Get all pools."""
raise NotImplementedError()
- def create_pool(self, name, slots, description):
+ def create_pool(self, name, slots, description, include_deferred):
"""Create a pool.
:param name: pool name
:param slots: pool slots amount
:param description: pool description
+ :param include_deferred: include deferred tasks in pool calculations
"""
raise NotImplementedError()
diff --git a/airflow/api/client/json_client.py b/airflow/api/client/json_client.py
index 2eb91b3a42..ac73813136 100644
--- a/airflow/api/client/json_client.py
+++ b/airflow/api/client/json_client.py
@@ -108,23 +108,26 @@ class Client(api_client.Client):
pools = self._request(url)
return [(p["pool"], p["slots"], p["description"]) for p in pools]
- def create_pool(self, name: str, slots: int, description: str):
+ def create_pool(self, name: str, slots: int, description: str, include_deferred: bool):
"""Create a new pool.
:param name: The name of the pool to create.
:param slots: The number of slots in the pool.
:param description: A description of the pool.
+ :param include_deferred: include deferred tasks in pool calculations
+
:return: A tuple containing the name of the pool, the number of slots in the pool,
- and a description of the pool.
+ a description of the pool and the include_deferred flag.
"""
endpoint = "/api/experimental/pools"
data = {
"name": name,
"slots": slots,
"description": description,
+ "include_deferred": include_deferred,
}
response = self._request(urljoin(self._api_base_url, endpoint), method="POST", json=data)
- return response["pool"], response["slots"], response["description"]
+ return response["pool"], response["slots"], response["description"], response["include_deferred"]
def delete_pool(self, name: str):
"""Delete a pool.
diff --git a/airflow/api/client/local_client.py b/airflow/api/client/local_client.py
index b90039608f..d8b9565099 100644
--- a/airflow/api/client/local_client.py
+++ b/airflow/api/client/local_client.py
@@ -63,12 +63,12 @@ class Client(api_client.Client):
pool = Pool.get_pool(pool_name=name)
if not pool:
raise PoolNotFound(f"Pool {name} not found")
- return pool.pool, pool.slots, pool.description
+ return pool.pool, pool.slots, pool.description, pool.include_deferred
def get_pools(self):
- return [(p.pool, p.slots, p.description) for p in Pool.get_pools()]
+ return [(p.pool, p.slots, p.description, p.include_deferred) for p in Pool.get_pools()]
- def create_pool(self, name, slots, description):
+ def create_pool(self, name, slots, description, include_deferred):
if not (name and name.strip()):
raise AirflowBadRequest("Pool name shouldn't be empty")
pool_name_length = Pool.pool.property.columns[0].type.length
@@ -78,7 +78,9 @@ class Client(api_client.Client):
slots = int(slots)
except ValueError:
raise AirflowBadRequest(f"Bad value for `slots`: {slots}")
- pool = Pool.create_or_update_pool(name=name, slots=slots, description=description)
+ pool = Pool.create_or_update_pool(
+ name=name, slots=slots, description=description, include_deferred=include_deferred
+ )
return pool.pool, pool.slots, pool.description
def delete_pool(self, name):
diff --git a/airflow/api/common/experimental/pool.py b/airflow/api/common/experimental/pool.py
index cff21d4d56..34e35cd435 100644
--- a/airflow/api/common/experimental/pool.py
+++ b/airflow/api/common/experimental/pool.py
@@ -68,7 +68,7 @@ def create_pool(name, slots, description, session: Session = NEW_SESSION):
session.expire_on_commit = False
pool = session.scalar(select(Pool).filter_by(pool=name).limit(1))
if pool is None:
- pool = Pool(pool=name, slots=slots, description=description)
+ pool = Pool(pool=name, slots=slots, description=description, include_deferred=False)
session.add(pool)
else:
pool.slots = slots
diff --git a/airflow/api_connexion/endpoints/pool_endpoint.py b/airflow/api_connexion/endpoints/pool_endpoint.py
index 6503705d85..fb59ab5ea2 100644
--- a/airflow/api_connexion/endpoints/pool_endpoint.py
+++ b/airflow/api_connexion/endpoints/pool_endpoint.py
@@ -88,10 +88,10 @@ def patch_pool(
) -> APIResponse:
"""Update a pool."""
request_dict = get_json_request_dict()
- # Only slots can be modified in 'default_pool'
+ # Only slots and include_deferred can be modified in 'default_pool'
try:
if pool_name == Pool.DEFAULT_POOL_NAME and request_dict["name"] != Pool.DEFAULT_POOL_NAME:
- if update_mask and len(update_mask) == 1 and update_mask[0].strip() == "slots":
+ if update_mask and all(mask.strip() in {"slots", "include_deferred"} for mask in update_mask):
pass
else:
raise BadRequest(detail="Default Pool's name can't be modified")
diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index a003788062..ca6d38921f 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -3158,7 +3158,7 @@ components:
occupied_slots:
type: integer
readOnly: true
- description: The number of slots used by running/queued tasks at the moment.
+ description: The number of slots used by running/queued tasks at the moment. May include deferred tasks if 'include_deferred' is set to true.
running_slots:
type: integer
readOnly: true
@@ -3175,6 +3175,13 @@ components:
type: integer
readOnly: true
description: The number of slots used by scheduled tasks at the moment.
+ deferred_slots:
+ type: integer
+ readOnly: true
+ description: |
+ The number of slots used by deferred tasks at the moment. Relevant if 'include_deferred' is set to true.
+
+ *New in version 2.7.0*
description:
type: string
description: |
@@ -3182,6 +3189,12 @@ components:
*New in version 2.3.0*
nullable: true
+ include_deferred:
+ type: boolean
+ description: |
+ If set to true, deferred tasks are considered when calculating open pool slots.
+
+ *New in version 2.7.0*
PoolCollection:
type: object
diff --git a/airflow/api_connexion/schemas/pool_schema.py b/airflow/api_connexion/schemas/pool_schema.py
index 4e25287d1d..e18548cee8 100644
--- a/airflow/api_connexion/schemas/pool_schema.py
+++ b/airflow/api_connexion/schemas/pool_schema.py
@@ -39,7 +39,10 @@ class PoolSchema(SQLAlchemySchema):
queued_slots = fields.Method("get_queued_slots", dump_only=True)
scheduled_slots = fields.Method("get_scheduled_slots", dump_only=True)
open_slots = fields.Method("get_open_slots", dump_only=True)
+ deferred_slots = fields.Method("get_deferred_slots", dump_only=True)
description = auto_field()
+ # we skip auto_field() here to be compatible with the manual validation in the pool_endpoint module
+ include_deferred = fields.Boolean(load_default=False)
@staticmethod
def get_occupied_slots(obj: Pool) -> int:
@@ -61,6 +64,11 @@ class PoolSchema(SQLAlchemySchema):
"""Returns the scheduled slots of the pool."""
return obj.scheduled_slots()
+ @staticmethod
+ def get_deferred_slots(obj: Pool) -> int:
+ """Returns the deferred slots of the pool."""
+ return obj.deferred_slots()
+
@staticmethod
def get_open_slots(obj: Pool) -> float:
"""Returns the open slots of the pool."""
diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py
index 52ed700db7..1314ffa0e3 100644
--- a/airflow/cli/cli_config.py
+++ b/airflow/cli/cli_config.py
@@ -513,6 +513,9 @@ ARG_DB_RETRY_DELAY = Arg(
ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
ARG_POOL_SLOTS = Arg(("slots",), type=int, help="Pool slots")
ARG_POOL_DESCRIPTION = Arg(("description",), help="Pool description")
+ARG_POOL_INCLUDE_DEFERRED = Arg(
+ ("--include-deferred",), help="Include deferred tasks in calculations for Pool", action="store_true"
+)
ARG_POOL_IMPORT = Arg(
("file",),
metavar="FILEPATH",
@@ -521,8 +524,8 @@ ARG_POOL_IMPORT = Arg(
textwrap.dedent(
"""
{
- "pool_1": {"slots": 5, "description": ""},
- "pool_2": {"slots": 10, "description": "test"}
+ "pool_1": {"slots": 5, "description": "", "include_deferred": true},
+ "pool_2": {"slots": 10, "description": "test", "include_deferred": false}
}"""
),
" " * 4,
@@ -1456,7 +1459,14 @@ POOLS_COMMANDS = (
name="set",
help="Configure pool",
func=lazy_load_command("airflow.cli.commands.pool_command.pool_set"),
- args=(ARG_POOL_NAME, ARG_POOL_SLOTS, ARG_POOL_DESCRIPTION, ARG_OUTPUT, ARG_VERBOSE),
+ args=(
+ ARG_POOL_NAME,
+ ARG_POOL_SLOTS,
+ ARG_POOL_DESCRIPTION,
+ ARG_POOL_INCLUDE_DEFERRED,
+ ARG_OUTPUT,
+ ARG_VERBOSE,
+ ),
),
ActionCommand(
name="delete",
diff --git a/airflow/cli/commands/pool_command.py b/airflow/cli/commands/pool_command.py
index 8d9e206f1b..7e7c63640c 100644
--- a/airflow/cli/commands/pool_command.py
+++ b/airflow/cli/commands/pool_command.py
@@ -38,6 +38,7 @@ def _show_pools(pools, output):
"pool": x[0],
"slots": x[1],
"description": x[2],
+ "include_deferred": x[3],
},
)
@@ -69,7 +70,9 @@ def pool_get(args):
def pool_set(args):
"""Creates new pool with a given name and slots."""
api_client = get_current_api_client()
- api_client.create_pool(name=args.pool, slots=args.slots, description=args.description)
+ api_client.create_pool(
+ name=args.pool, slots=args.slots, description=args.description, include_deferred=args.include_deferred
+ )
print(f"Pool {args.pool} created")
@@ -119,8 +122,15 @@ def pool_import_helper(filepath):
pools = []
failed = []
for k, v in pools_json.items():
- if isinstance(v, dict) and len(v) == 2:
- pools.append(api_client.create_pool(name=k, slots=v["slots"], description=v["description"]))
+ if isinstance(v, dict) and "slots" in v and "description" in v:
+ pools.append(
+ api_client.create_pool(
+ name=k,
+ slots=v["slots"],
+ description=v["description"],
+ include_deferred=v.get("include_deferred", False),
+ )
+ )
else:
failed.append(k)
return pools, failed
@@ -132,7 +142,7 @@ def pool_export_helper(filepath):
pool_dict = {}
pools = api_client.get_pools()
for pool in pools:
- pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2]}
+ pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2], "include_deferred": pool[3]}
with open(filepath, "w") as poolfile:
poolfile.write(json.dumps(pool_dict, sort_keys=True, indent=4))
return pools
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
index e47fc7bd86..0ed7aae5ce 100644
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -1571,10 +1571,12 @@ class SchedulerJobRunner(BaseJobRunner[Job], LoggingMixin):
Stats.gauge(f"pool.open_slots.{pool_name}", slot_stats["open"])
Stats.gauge(f"pool.queued_slots.{pool_name}", slot_stats["queued"])
Stats.gauge(f"pool.running_slots.{pool_name}", slot_stats["running"])
+ Stats.gauge(f"pool.deferred_slots.{pool_name}", slot_stats["deferred"])
# Same metrics with tagging
Stats.gauge("pool.open_slots", slot_stats["open"], tags={"pool_name": pool_name})
Stats.gauge("pool.queued_slots", slot_stats["queued"], tags={"pool_name": pool_name})
Stats.gauge("pool.running_slots", slot_stats["running"], tags={"pool_name": pool_name})
+ Stats.gauge("pool.deferred_slots", slot_stats["deferred"], tags={"pool_name": pool_name})
@provide_session
def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
diff --git a/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
new file mode 100644
index 0000000000..ea10320480
--- /dev/null
+++ b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
@@ -0,0 +1,50 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""add include_deferred column to pool
+
+Revision ID: 405de8318b3a
+Revises: 788397e78828
+Create Date: 2023-07-20 04:22:21.007342
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "405de8318b3a"
+down_revision = "788397e78828"
+branch_labels = None
+depends_on = None
+airflow_version = "2.7.0"
+
+
+def upgrade():
+ """Apply add include_deferred column to pool"""
+ with op.batch_alter_table("slot_pool") as batch_op:
+ batch_op.add_column(
+ sa.Column("include_deferred", sa.Boolean, nullable=False, server_default=sa.false())
+ )
+
+
+def downgrade():
+ """Unapply add include_deferred column to pool"""
+ with op.batch_alter_table("slot_pool") as batch_op:
+ batch_op.drop_column("include_deferred")
diff --git a/airflow/models/pool.py b/airflow/models/pool.py
index 3b5708b121..13a2e15e66 100644
--- a/airflow/models/pool.py
+++ b/airflow/models/pool.py
@@ -19,7 +19,7 @@ from __future__ import annotations
from typing import Any
-from sqlalchemy import Column, Integer, String, Text, func, select
+from sqlalchemy import Boolean, Column, Integer, String, Text, func, select
from sqlalchemy.orm.session import Session
from airflow.exceptions import AirflowException, PoolNotFound
@@ -37,6 +37,7 @@ class PoolStats(TypedDict):
total: int
running: int
+ deferred: int
queued: int
open: int
@@ -51,6 +52,7 @@ class Pool(Base):
# -1 for infinite
slots = Column(Integer, default=0)
description = Column(Text)
+ include_deferred = Column(Boolean, nullable=False)
DEFAULT_POOL_NAME = "default_pool"
@@ -108,6 +110,7 @@ class Pool(Base):
name: str,
slots: int,
description: str,
+ include_deferred: bool,
session: Session = NEW_SESSION,
) -> Pool:
"""Create a pool with given parameters or update it if it already exists."""
@@ -116,11 +119,12 @@ class Pool(Base):
pool = session.scalar(select(Pool).filter_by(pool=name))
if pool is None:
- pool = Pool(pool=name, slots=slots, description=description)
+ pool = Pool(pool=name, slots=slots, description=description, include_deferred=include_deferred)
session.add(pool)
else:
pool.slots = slots
pool.description = description
+ pool.include_deferred = include_deferred
session.commit()
return pool
@@ -161,21 +165,26 @@ class Pool(Base):
from airflow.models.taskinstance import TaskInstance # Avoid circular import
pools: dict[str, PoolStats] = {}
+ pool_includes_deferred: dict[str, bool] = {}
- query = select(Pool.pool, Pool.slots)
+ query = select(Pool.pool, Pool.slots, Pool.include_deferred)
if lock_rows:
query = with_row_locks(query, session=session, **nowait(session))
pool_rows = session.execute(query)
- for (pool_name, total_slots) in pool_rows:
+ for (pool_name, total_slots, include_deferred) in pool_rows:
if total_slots == -1:
total_slots = float("inf") # type: ignore
- pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0)
+ pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0, deferred=0)
+ pool_includes_deferred[pool_name] = include_deferred
+ allowed_execution_states = EXECUTION_STATES | {
+ TaskInstanceState.DEFERRED,
+ }
state_count_by_pool = session.execute(
select(TaskInstance.pool, TaskInstance.state, func.sum(TaskInstance.pool_slots))
- .filter(TaskInstance.state.in_(EXECUTION_STATES))
+ .filter(TaskInstance.state.in_(allowed_execution_states))
.group_by(TaskInstance.pool, TaskInstance.state)
)
@@ -192,12 +201,16 @@ class Pool(Base):
stats_dict["running"] = count
elif state == TaskInstanceState.QUEUED:
stats_dict["queued"] = count
+ elif state == TaskInstanceState.DEFERRED:
+ stats_dict["deferred"] = count
else:
- raise AirflowException(f"Unexpected state. Expected values: {EXECUTION_STATES}.")
+ raise AirflowException(f"Unexpected state. Expected values: {allowed_execution_states}.")
# calculate open metric
for pool_name, stats_dict in pools.items():
stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"]
+ if pool_includes_deferred[pool_name]:
+ stats_dict["open"] -= stats_dict["deferred"]
return pools
@@ -212,6 +225,7 @@ class Pool(Base):
"pool": self.pool,
"slots": self.slots,
"description": self.description,
+ "include_deferred": self.include_deferred,
}
@provide_session
@@ -224,15 +238,24 @@ class Pool(Base):
"""
from airflow.models.taskinstance import TaskInstance # Avoid circular import
+ occupied_states = self.get_occupied_states()
+
return int(
session.scalar(
select(func.sum(TaskInstance.pool_slots))
.filter(TaskInstance.pool == self.pool)
- .filter(TaskInstance.state.in_(EXECUTION_STATES))
+ .filter(TaskInstance.state.in_(occupied_states))
)
or 0
)
+ def get_occupied_states(self):
+ if self.include_deferred:
+ return EXECUTION_STATES | {
+ TaskInstanceState.DEFERRED,
+ }
+ return EXECUTION_STATES
+
@provide_session
def running_slots(self, session: Session = NEW_SESSION) -> int:
"""
@@ -290,6 +313,25 @@ class Pool(Base):
or 0
)
+ @provide_session
+ def deferred_slots(self, session: Session = NEW_SESSION) -> int:
+ """
+ Get the number of slots deferred at the moment.
+
+ :param session: SQLAlchemy ORM Session
+ :return: the number of deferred slots
+ """
+ from airflow.models.taskinstance import TaskInstance # Avoid circular import
+
+ return int(
+ session.scalar(
+ select(func.sum(TaskInstance.pool_slots)).where(
+ TaskInstance.pool == self.pool, TaskInstance.state == TaskInstanceState.DEFERRED
+ )
+ )
+ or 0
+ )
+
@provide_session
def open_slots(self, session: Session = NEW_SESSION) -> float:
"""
diff --git a/airflow/ti_deps/deps/pool_slots_available_dep.py b/airflow/ti_deps/deps/pool_slots_available_dep.py
index 38bfe31936..e095cfb68e 100644
--- a/airflow/ti_deps/deps/pool_slots_available_dep.py
+++ b/airflow/ti_deps/deps/pool_slots_available_dep.py
@@ -18,7 +18,6 @@
"""This module defines dep for pool slots availability."""
from __future__ import annotations
-from airflow.ti_deps.dependencies_states import EXECUTION_STATES
from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
from airflow.utils.session import provide_session
@@ -49,12 +48,12 @@ class PoolSlotsAvailableDep(BaseTIDep):
reason=f"Tasks using non-existent pool '{pool_name}' will not be scheduled"
)
return
- else:
- # Controlled by UNIQUE key in slot_pool table,
- # only one result can be returned.
- open_slots = pools[0].open_slots(session=session)
+ # Controlled by UNIQUE key in slot_pool table,
+ # only one result can be returned.
+ open_slots = pools[0].open_slots(session=session)
- if ti.state in EXECUTION_STATES:
+ occupied_states = pools[0].get_occupied_states()
+ if ti.state in occupied_states:
open_slots += ti.pool_slots
if open_slots <= (ti.pool_slots - 1):
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 185687ad56..68425f8a4c 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -125,6 +125,7 @@ def add_default_pool_if_not_exists(session: Session = NEW_SESSION):
pool=Pool.DEFAULT_POOL_NAME,
slots=conf.getint(section="core", key="default_pool_task_slot_count"),
description="Default pool",
+ include_deferred=False,
)
session.add(default_pool)
session.commit()
diff --git a/airflow/www/static/js/types/api-generated.ts b/airflow/www/static/js/types/api-generated.ts
index e2d38d0725..8dead88fc7 100644
--- a/airflow/www/static/js/types/api-generated.ts
+++ b/airflow/www/static/js/types/api-generated.ts
@@ -1239,7 +1239,7 @@ export interface components {
name?: string;
/** @description The maximum number of slots that can be assigned to tasks. One job may occupy one or more slots. */
slots?: number;
- /** @description The number of slots used by running/queued tasks at the moment. */
+ /** @description The number of slots used by running/queued tasks at the moment. May include deferred tasks if 'include_deferred' is set to true. */
occupied_slots?: number;
/** @description The number of slots used by running tasks at the moment. */
running_slots?: number;
@@ -1249,12 +1249,24 @@ export interface components {
open_slots?: number;
/** @description The number of slots used by scheduled tasks at the moment. */
scheduled_slots?: number;
+ /**
+ * @description The number of slots used by deferred tasks at the moment. Relevant if 'include_deferred' is set to true.
+ *
+ * *New in version 2.7.0*
+ */
+ deferred_slots?: number;
/**
* @description The description of the pool.
*
* *New in version 2.3.0*
*/
description?: string | null;
+ /**
+ * @description If set to true, deferred tasks are considered when calculating open pool slots.
+ *
+ * *New in version 2.7.0*
+ */
+ include_deferred?: boolean;
};
/**
* @description Collection of pools.
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 3faedd3d8b..5f12a94969 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -70,7 +70,7 @@ from pendulum.parsing.exceptions import ParserError
from sqlalchemy import Date, and_, case, desc, func, inspect, or_, select, union_all
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session, joinedload
-from wtforms import validators
+from wtforms import BooleanField, validators
import airflow
from airflow import models, plugins_manager, settings
@@ -4874,9 +4874,17 @@ class PoolModelView(AirflowModelView):
permissions.ACTION_CAN_ACCESS_MENU,
]
- list_columns = ["pool", "slots", "running_slots", "queued_slots", "scheduled_slots"]
- add_columns = ["pool", "slots", "description"]
- edit_columns = ["pool", "slots", "description"]
+ list_columns = ["pool", "slots", "running_slots", "queued_slots", "scheduled_slots", "deferred_slots"]
+ add_columns = ["pool", "slots", "description", "include_deferred"]
+ edit_columns = ["pool", "slots", "description", "include_deferred"]
+
+ # include_deferred is non-nullable, but as a checkbox in the resulting form we want to allow it unchecked
+ include_deferred_field = BooleanField(
+ validators=[validators.Optional()],
+ description="Check to include deferred tasks when calculating open pool slots.",
+ )
+ edit_form_extra_fields = {"include_deferred": include_deferred_field}
+ add_form_extra_fields = {"include_deferred": include_deferred_field}
base_order = ("pool", "asc")
@@ -4943,11 +4951,24 @@ class PoolModelView(AirflowModelView):
else:
return Markup('<span class="label label-danger">Invalid</span>')
+ def fdeferred_slots(self):
+ """Deferred slots rendering."""
+ pool_id = self.get("pool")
+ deferred_slots = self.get("deferred_slots")
+ if pool_id is not None and deferred_slots is not None:
+ url = url_for("TaskInstanceModelView.list", _flt_3_pool=pool_id, _flt_3_state="deferred")
+ return Markup("<a href='{url}'>{deferred_slots}</a>").format(
+ url=url, deferred_slots=deferred_slots
+ )
+ else:
+ return Markup('<span class="label label-danger">Invalid</span>')
+
formatters_columns = {
"pool": pool_link,
"running_slots": frunning_slots,
"queued_slots": fqueued_slots,
"scheduled_slots": fscheduled_slots,
+ "deferred_slots": fdeferred_slots,
}
validators_columns = {"pool": [validators.DataRequired()], "slots": [validators.NumberRange(min=-1)]}
diff --git a/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst b/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
index e4bb1da68a..b18d05a877 100644
--- a/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
+++ b/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
@@ -214,6 +214,7 @@ Name Description
``pool.open_slots.<pool_name>`` Number of open slots in the pool
``pool.queued_slots.<pool_name>`` Number of queued slots in the pool
``pool.running_slots.<pool_name>`` Number of running slots in the pool
+``pool.deferred_slots.<pool_name>`` Number of deferred slots in the pool
``pool.starving_tasks.<pool_name>`` Number of starving tasks in the pool
``triggers.running.<hostname>`` Number of triggers currently running for a triggerer (described by hostname)
=================================================== ========================================================================
diff --git a/docs/apache-airflow/administration-and-deployment/pools.rst b/docs/apache-airflow/administration-and-deployment/pools.rst
index 5395ccfe38..5a906d3f6f 100644
--- a/docs/apache-airflow/administration-and-deployment/pools.rst
+++ b/docs/apache-airflow/administration-and-deployment/pools.rst
@@ -22,7 +22,7 @@ Pools
Some systems can get overwhelmed when too many processes hit them at the same time. Airflow pools can be used to
**limit the execution parallelism** on arbitrary sets of tasks. The list of pools is managed in the UI
-(``Menu -> Admin -> Pools``) by giving the pools a name and assigning it a number of worker slots.
+(``Menu -> Admin -> Pools``) by giving the pools a name and assigning it a number of worker slots. There you can also decide whether the pool should include :doc:`deferred tasks </authoring-and-scheduling/deferring>` in its calculation of occupied slots.
Tasks can then be associated with one of the existing pools by using the ``pool`` parameter when creating tasks:
diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index 6b5f124e8f..c44ac2c1dd 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -20,7 +20,7 @@ Deferrable Operators & Triggers
Standard :doc:`Operators </core-concepts/operators>` and :doc:`Sensors <../core-concepts/sensors>` take up a full *worker slot* for the entire time they are running, even if they are idle; for example, if you only have 100 worker slots available to run Tasks, and you have 100 DAGs waiting on a Sensor that's currently running but idle, then you *cannot run anything else* - even though your entire Airflow cluster is essentially idle. ``reschedule`` mode for Sensors solves some of this, all [...]
-This is where *Deferrable Operators* come in. A deferrable operator is one that is written with the ability to suspend itself and free up the worker when it knows it has to wait, and hand off the job of resuming it to something called a *Trigger*. As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a lot less resources wasted on idle Operators or Sensors.
+This is where *Deferrable Operators* come in. A deferrable operator is one that is written with the ability to suspend itself and free up the worker when it knows it has to wait, and hand off the job of resuming it to something called a *Trigger*. As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a lot less resources wasted on idle Operators or Sensors. Note that by default deferred tasks will not use up pool slots, if you would l [...]
*Triggers* are small, asynchronous pieces of Python code designed to be run all together in a single Python process; because they are asynchronous, they are able to all co-exist efficiently. As an overview of how this process works:
diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256
index f10af90f5b..5f9515fc19 100644
--- a/docs/apache-airflow/img/airflow_erd.sha256
+++ b/docs/apache-airflow/img/airflow_erd.sha256
@@ -1 +1 @@
-7ff18b1eafa528dbdfb62d75151a526de280f73e1edf7fea18da7c64644f0da9
\ No newline at end of file
+3420f98a9dfa230d76be4da57bc4a3641cb929b33b9eafe7ffa41fb4fcc4d283
\ No newline at end of file
diff --git a/docs/apache-airflow/img/airflow_erd.svg b/docs/apache-airflow/img/airflow_erd.svg
index 18e5f83b11..ba548c60fd 100644
--- a/docs/apache-airflow/img/airflow_erd.svg
+++ b/docs/apache-airflow/img/airflow_erd.svg
@@ -4,11 +4,11 @@
<!-- Generated by graphviz version 2.43.0 (0)
-->
<!-- Title: %3 Pages: 1 -->
-<svg width="1559pt" height="5959pt"
- viewBox="0.00 0.00 1559.00 5958.50" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
-<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 5954.5)">
+<svg width="1559pt" height="5983pt"
+ viewBox="0.00 0.00 1559.00 5982.50" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
+<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 5978.5)">
<title>%3</title>
-<polygon fill="white" stroke="transparent" points="-4,4 -4,-5954.5 1555,-5954.5 1555,4 -4,4"/>
+<polygon fill="white" stroke="transparent" points="-4,4 -4,-5978.5 1555,-5978.5 1555,4 -4,4"/>
<!-- ab_permission -->
<g id="node1" class="node">
<title>ab_permission</title>
@@ -1228,28 +1228,28 @@
<g id="edge48" class="edge">
<title>task_instance--xcom</title>
<path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-482.51C1190.54,-462.01 1215.43,-441.43 1239,-422 1245.31,-416.8 1251.76,-411.48 1258.29,-406.14"/>
-<text text-anchor="start" x="1248.29" y="-394.94" font-family="Times,serif" font-size="14.00">1</text>
+<text text-anchor="start" x="1227.29" y="-394.94" font-family="Times,serif" font-size="14.00">0..N</text>
<text text-anchor="start" x="1166.17" y="-471.31" font-family="Times,serif" font-size="14.00">1</text>
</g>
<!-- task_instance--xcom -->
<g id="edge49" class="edge">
<title>task_instance--xcom</title>
<path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-500.3C1190.54,-480.01 1215.43,-459.43 1239,-440 1245.31,-434.8 1251.76,-429.48 1258.29,-424.12"/>
-<text text-anchor="start" x="1227.29" y="-412.92" font-family="Times,serif" font-size="14.00">0..N</text>
+<text text-anchor="start" x="1248.29" y="-412.92" font-family="Times,serif" font-size="14.00">1</text>
<text text-anchor="start" x="1166.17" y="-489.1" font-family="Times,serif" font-size="14.00">1</text>
</g>
<!-- task_instance--xcom -->
<g id="edge50" class="edge">
<title>task_instance--xcom</title>
<path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-518.08C1190.54,-498.01 1215.43,-477.43 1239,-458 1246.11,-452.14 1253.42,-446.12 1260.79,-440.05"/>
-<text text-anchor="start" x="1250.79" y="-443.85" font-family="Times,serif" font-size="14.00">1</text>
+<text text-anchor="start" x="1229.79" y="-443.85" font-family="Times,serif" font-size="14.00">0..N</text>
<text text-anchor="start" x="1166.17" y="-506.88" font-family="Times,serif" font-size="14.00">1</text>
</g>
<!-- task_instance--xcom -->
<g id="edge51" class="edge">
<title>task_instance--xcom</title>
<path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-535.87C1190.54,-516.01 1215.43,-495.43 1239,-476 1252.96,-464.49 1267.64,-452.42 1282.18,-440.26"/>
-<text text-anchor="start" x="1282.18" y="-444.06" font-family="Times,serif" font-size="14.00">0..N</text>
+<text text-anchor="start" x="1272.18" y="-444.06" font-family="Times,serif" font-size="14.00">1</text>
<text text-anchor="start" x="1166.17" y="-524.67" font-family="Times,serif" font-size="14.00">1</text>
</g>
<!-- log_template -->
@@ -1582,43 +1582,47 @@
<!-- slot_pool -->
<g id="node41" class="node">
<title>slot_pool</title>
-<polygon fill="none" stroke="black" points="103.5,-5739 103.5,-5767 279.5,-5767 279.5,-5739 103.5,-5739"/>
-<text text-anchor="start" x="150.5" y="-5750.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">slot_pool</text>
-<polygon fill="none" stroke="black" points="103.5,-5714 103.5,-5739 279.5,-5739 279.5,-5714 103.5,-5714"/>
-<text text-anchor="start" x="108.5" y="-5723.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
-<text text-anchor="start" x="121.5" y="-5723.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
-<text text-anchor="start" x="198.5" y="-5723.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
-<polygon fill="none" stroke="black" points="103.5,-5689 103.5,-5714 279.5,-5714 279.5,-5689 103.5,-5689"/>
-<text text-anchor="start" x="108.5" y="-5698.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
-<text text-anchor="start" x="186.5" y="-5698.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
-<polygon fill="none" stroke="black" points="103.5,-5664 103.5,-5689 279.5,-5689 279.5,-5664 103.5,-5664"/>
-<text text-anchor="start" x="108.5" y="-5673.8" font-family="Helvetica,sans-Serif" font-size="14.00">pool</text>
-<text text-anchor="start" x="138.5" y="-5673.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(256)]</text>
-<polygon fill="none" stroke="black" points="103.5,-5639 103.5,-5664 279.5,-5664 279.5,-5639 103.5,-5639"/>
-<text text-anchor="start" x="108.5" y="-5648.8" font-family="Helvetica,sans-Serif" font-size="14.00">slots</text>
-<text text-anchor="start" x="141.5" y="-5648.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<polygon fill="none" stroke="black" points="48.5,-5763 48.5,-5791 334.5,-5791 334.5,-5763 48.5,-5763"/>
+<text text-anchor="start" x="150.5" y="-5774.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">slot_pool</text>
+<polygon fill="none" stroke="black" points="48.5,-5738 48.5,-5763 334.5,-5763 334.5,-5738 48.5,-5738"/>
+<text text-anchor="start" x="53.5" y="-5747.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
+<text text-anchor="start" x="66.5" y="-5747.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<text text-anchor="start" x="143.5" y="-5747.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="48.5,-5713 48.5,-5738 334.5,-5738 334.5,-5713 48.5,-5713"/>
+<text text-anchor="start" x="53.5" y="-5722.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
+<text text-anchor="start" x="131.5" y="-5722.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="48.5,-5688 48.5,-5713 334.5,-5713 334.5,-5688 48.5,-5688"/>
+<text text-anchor="start" x="53.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00">include_deferred</text>
+<text text-anchor="start" x="169.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
+<text text-anchor="start" x="253.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="48.5,-5663 48.5,-5688 334.5,-5688 334.5,-5663 48.5,-5663"/>
+<text text-anchor="start" x="53.5" y="-5672.8" font-family="Helvetica,sans-Serif" font-size="14.00">pool</text>
+<text text-anchor="start" x="83.5" y="-5672.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(256)]</text>
+<polygon fill="none" stroke="black" points="48.5,-5638 48.5,-5663 334.5,-5663 334.5,-5638 48.5,-5638"/>
+<text text-anchor="start" x="53.5" y="-5647.8" font-family="Helvetica,sans-Serif" font-size="14.00">slots</text>
+<text text-anchor="start" x="86.5" y="-5647.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
</g>
<!-- variable -->
<g id="node42" class="node">
<title>variable</title>
-<polygon fill="none" stroke="black" points="100.5,-5918 100.5,-5946 283.5,-5946 283.5,-5918 100.5,-5918"/>
-<text text-anchor="start" x="155.5" y="-5929.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">variable</text>
-<polygon fill="none" stroke="black" points="100.5,-5893 100.5,-5918 283.5,-5918 283.5,-5893 100.5,-5893"/>
-<text text-anchor="start" x="105.5" y="-5902.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
-<text text-anchor="start" x="118.5" y="-5902.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
-<text text-anchor="start" x="195.5" y="-5902.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
-<polygon fill="none" stroke="black" points="100.5,-5868 100.5,-5893 283.5,-5893 283.5,-5868 100.5,-5868"/>
-<text text-anchor="start" x="105.5" y="-5877.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
-<text text-anchor="start" x="183.5" y="-5877.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
-<polygon fill="none" stroke="black" points="100.5,-5843 100.5,-5868 283.5,-5868 283.5,-5843 100.5,-5843"/>
-<text text-anchor="start" x="105.5" y="-5852.8" font-family="Helvetica,sans-Serif" font-size="14.00">is_encrypted</text>
-<text text-anchor="start" x="194.5" y="-5852.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
-<polygon fill="none" stroke="black" points="100.5,-5818 100.5,-5843 283.5,-5843 283.5,-5818 100.5,-5818"/>
-<text text-anchor="start" x="105.5" y="-5827.8" font-family="Helvetica,sans-Serif" font-size="14.00">key</text>
-<text text-anchor="start" x="130.5" y="-5827.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(250)]</text>
-<polygon fill="none" stroke="black" points="100.5,-5793 100.5,-5818 283.5,-5818 283.5,-5793 100.5,-5793"/>
-<text text-anchor="start" x="105.5" y="-5802.8" font-family="Helvetica,sans-Serif" font-size="14.00">val</text>
-<text text-anchor="start" x="126.5" y="-5802.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="100.5,-5942 100.5,-5970 283.5,-5970 283.5,-5942 100.5,-5942"/>
+<text text-anchor="start" x="155.5" y="-5953.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">variable</text>
+<polygon fill="none" stroke="black" points="100.5,-5917 100.5,-5942 283.5,-5942 283.5,-5917 100.5,-5917"/>
+<text text-anchor="start" x="105.5" y="-5926.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
+<text text-anchor="start" x="118.5" y="-5926.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<text text-anchor="start" x="195.5" y="-5926.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="100.5,-5892 100.5,-5917 283.5,-5917 283.5,-5892 100.5,-5892"/>
+<text text-anchor="start" x="105.5" y="-5901.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
+<text text-anchor="start" x="183.5" y="-5901.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="100.5,-5867 100.5,-5892 283.5,-5892 283.5,-5867 100.5,-5867"/>
+<text text-anchor="start" x="105.5" y="-5876.8" font-family="Helvetica,sans-Serif" font-size="14.00">is_encrypted</text>
+<text text-anchor="start" x="194.5" y="-5876.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
+<polygon fill="none" stroke="black" points="100.5,-5842 100.5,-5867 283.5,-5867 283.5,-5842 100.5,-5842"/>
+<text text-anchor="start" x="105.5" y="-5851.8" font-family="Helvetica,sans-Serif" font-size="14.00">key</text>
+<text text-anchor="start" x="130.5" y="-5851.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(250)]</text>
+<polygon fill="none" stroke="black" points="100.5,-5817 100.5,-5842 283.5,-5842 283.5,-5817 100.5,-5817"/>
+<text text-anchor="start" x="105.5" y="-5826.8" font-family="Helvetica,sans-Serif" font-size="14.00">val</text>
+<text text-anchor="start" x="126.5" y="-5826.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
</g>
</g>
</svg>
diff --git a/docs/apache-airflow/migrations-ref.rst b/docs/apache-airflow/migrations-ref.rst
index be796354da..a3a42654f6 100644
--- a/docs/apache-airflow/migrations-ref.rst
+++ b/docs/apache-airflow/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
+---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description |
+=================================+===================+===================+==============================================================+
-| ``788397e78828`` (head) | ``937cbd173ca1`` | ``2.7.0`` | Add custom_operator_name column |
+| ``405de8318b3a`` (head) | ``788397e78828`` | ``2.7.0`` | add include_deferred column to pool |
++---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
+| ``788397e78828`` | ``937cbd173ca1`` | ``2.7.0`` | Add custom_operator_name column |
+---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
| ``937cbd173ca1`` | ``c804e5c76e3e`` | ``2.7.0`` | Add index to task_instance table |
+---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
diff --git a/tests/api_connexion/endpoints/test_pool_endpoint.py b/tests/api_connexion/endpoints/test_pool_endpoint.py
index c12e2473f2..a0121fccd1 100644
--- a/tests/api_connexion/endpoints/test_pool_endpoint.py
+++ b/tests/api_connexion/endpoints/test_pool_endpoint.py
@@ -63,7 +63,7 @@ class TestBasePoolEndpoints:
class TestGetPools(TestBasePoolEndpoints):
def test_response_200(self, session):
- pool_model = Pool(pool="test_pool_a", slots=3)
+ pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
session.add(pool_model)
session.commit()
result = session.query(Pool).all()
@@ -79,8 +79,10 @@ class TestGetPools(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 128,
"description": "Default pool",
+ "include_deferred": False,
},
{
"name": "test_pool_a",
@@ -89,15 +91,17 @@ class TestGetPools(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": None,
+ "include_deferred": True,
},
],
"total_entries": 2,
} == response.json
def test_response_200_with_order_by(self, session):
- pool_model = Pool(pool="test_pool_a", slots=3)
+ pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
session.add(pool_model)
session.commit()
result = session.query(Pool).all()
@@ -113,8 +117,10 @@ class TestGetPools(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": None,
+ "include_deferred": True,
},
{
"name": "default_pool",
@@ -123,8 +129,10 @@ class TestGetPools(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 128,
"description": "Default pool",
+ "include_deferred": False,
},
],
"total_entries": 2,
@@ -164,7 +172,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
)
@provide_session
def test_limit_and_offset(self, url, expected_pool_ids, session):
- pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+ pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
session.add_all(pools)
session.commit()
result = session.query(Pool).count()
@@ -175,7 +183,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
assert pool_ids == expected_pool_ids
def test_should_respect_page_size_limit_default(self, session):
- pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+ pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
session.add_all(pools)
session.commit()
result = session.query(Pool).count()
@@ -185,7 +193,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
assert len(response.json["pools"]) == 100
def test_should_raise_400_for_invalid_orderby(self, session):
- pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+ pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
session.add_all(pools)
session.commit()
result = session.query(Pool).count()
@@ -199,7 +207,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
@conf_vars({("api", "maximum_page_limit"): "150"})
def test_should_return_conf_max_if_req_max_above_conf(self, session):
- pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 200)]
+ pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 200)]
session.add_all(pools)
session.commit()
result = session.query(Pool).count()
@@ -211,7 +219,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
class TestGetPool(TestBasePoolEndpoints):
def test_response_200(self, session):
- pool_model = Pool(pool="test_pool_a", slots=3)
+ pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
session.add(pool_model)
session.commit()
response = self.client.get("/api/v1/pools/test_pool_a", environ_overrides={"REMOTE_USER": "test"})
@@ -223,8 +231,10 @@ class TestGetPool(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": None,
+ "include_deferred": True,
} == response.json
def test_response_404(self):
@@ -246,7 +256,7 @@ class TestGetPool(TestBasePoolEndpoints):
class TestDeletePool(TestBasePoolEndpoints):
def test_response_204(self, session):
pool_name = "test_pool"
- pool_instance = Pool(pool=pool_name, slots=3)
+ pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
session.add(pool_instance)
session.commit()
@@ -268,7 +278,7 @@ class TestDeletePool(TestBasePoolEndpoints):
def test_should_raises_401_unauthenticated(self, session):
pool_name = "test_pool"
- pool_instance = Pool(pool=pool_name, slots=3)
+ pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
session.add(pool_instance)
session.commit()
@@ -285,7 +295,7 @@ class TestPostPool(TestBasePoolEndpoints):
def test_response_200(self):
response = self.client.post(
"api/v1/pools",
- json={"name": "test_pool_a", "slots": 3, "description": "test pool"},
+ json={"name": "test_pool_a", "slots": 3, "description": "test pool", "include_deferred": True},
environ_overrides={"REMOTE_USER": "test"},
)
assert response.status_code == 200
@@ -296,18 +306,20 @@ class TestPostPool(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": "test pool",
+ "include_deferred": True,
} == response.json
def test_response_409(self, session):
pool_name = "test_pool_a"
- pool_instance = Pool(pool=pool_name, slots=3)
+ pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
session.add(pool_instance)
session.commit()
response = self.client.post(
"api/v1/pools",
- json={"name": "test_pool_a", "slots": 3},
+ json={"name": "test_pool_a", "slots": 3, "include_deferred": False},
environ_overrides={"REMOTE_USER": "test"},
)
assert response.status_code == 409
@@ -334,7 +346,7 @@ class TestPostPool(TestBasePoolEndpoints):
pytest.param(
{},
"Missing required property(ies): ['name', 'slots']",
- id="for missing pool name AND slots",
+ id="for missing pool name AND slots AND include_deferred",
),
pytest.param(
{"name": "invalid_pool", "slots": 3, "extra_field_1": "extra"},
@@ -363,12 +375,12 @@ class TestPostPool(TestBasePoolEndpoints):
class TestPatchPool(TestBasePoolEndpoints):
def test_response_200(self, session):
- pool = Pool(pool="test_pool", slots=2)
+ pool = Pool(pool="test_pool", slots=2, include_deferred=True)
session.add(pool)
session.commit()
response = self.client.patch(
"api/v1/pools/test_pool",
- json={"name": "test_pool_a", "slots": 3},
+ json={"name": "test_pool_a", "slots": 3, "include_deferred": False},
environ_overrides={"REMOTE_USER": "test"},
)
assert response.status_code == 200
@@ -379,8 +391,10 @@ class TestPatchPool(TestBasePoolEndpoints):
"open_slots": 3,
"running_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"slots": 3,
"description": None,
+ "include_deferred": False,
} == response.json
@pytest.mark.parametrize(
@@ -393,13 +407,13 @@ class TestPatchPool(TestBasePoolEndpoints):
# Extra properties
(
"{'extra_field': ['Unknown field.']}",
- {"name": "test_pool_a", "slots": 3, "extra_field": "extra"},
+ {"name": "test_pool_a", "slots": 3, "include_deferred": True, "extra_field": "extra"},
),
],
)
@provide_session
def test_response_400(self, error_detail, request_json, session):
- pool = Pool(pool="test_pool", slots=2)
+ pool = Pool(pool="test_pool", slots=2, include_deferred=False)
session.add(pool)
session.commit()
response = self.client.patch(
@@ -414,7 +428,7 @@ class TestPatchPool(TestBasePoolEndpoints):
} == response.json
def test_should_raises_401_unauthenticated(self, session):
- pool = Pool(pool="test_pool", slots=2)
+ pool = Pool(pool="test_pool", slots=2, include_deferred=False)
session.add(pool)
session.commit()
@@ -443,7 +457,7 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
pytest.param(
400,
"api/v1/pools/default_pool",
- {"name": "test_pool_a", "slots": 3},
+ {"name": "test_pool_a", "slots": 3, "include_deferred": False},
{
"detail": "Default Pool's name can't be modified",
"status": 400,
@@ -455,7 +469,7 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
pytest.param(
400,
"api/v1/pools/default_pool?update_mask=name, slots",
- {"name": "test_pool_a", "slots": 3},
+ {"name": "test_pool_a", "slots": 3, "include_deferred": False},
{
"detail": "Default Pool's name can't be modified",
"status": 400,
@@ -475,11 +489,49 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
"open_slots": 3,
"running_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"slots": 3,
"description": "Default pool",
+ "include_deferred": False,
},
id="200 Update mask with slots",
),
+ pytest.param(
+ 200,
+ "api/v1/pools/default_pool?update_mask=include_deferred",
+ {"name": "test_pool_a", "include_deferred": True},
+ {
+ "occupied_slots": 0,
+ "queued_slots": 0,
+ "name": "default_pool",
+ "open_slots": 128,
+ "running_slots": 0,
+ "scheduled_slots": 0,
+ "deferred_slots": 0,
+ "slots": 128,
+ "description": "Default pool",
+ "include_deferred": True,
+ },
+ id="200 Update mask with include_deferred",
+ ),
+ pytest.param(
+ 200,
+ "api/v1/pools/default_pool?update_mask=slots,include_deferred",
+ {"name": "test_pool_a", "slots": 3, "include_deferred": True},
+ {
+ "occupied_slots": 0,
+ "queued_slots": 0,
+ "name": "default_pool",
+ "open_slots": 3,
+ "running_slots": 0,
+ "scheduled_slots": 0,
+ "deferred_slots": 0,
+ "slots": 3,
+ "description": "Default pool",
+ "include_deferred": True,
+ },
+ id="200 Update mask with slots AND include_deferred",
+ ),
pytest.param(
200,
"api/v1/pools/default_pool?update_mask=name,slots",
@@ -491,8 +543,10 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
"open_slots": 3,
"running_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"slots": 3,
"description": "Default pool",
+ "include_deferred": False,
},
id="200 Update mask with slots and name",
),
@@ -502,6 +556,7 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
{
"name": "default_pool",
"slots": 3,
+ "include_deferred": True,
},
{
"occupied_slots": 0,
@@ -510,8 +565,10 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
"open_slots": 3,
"running_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"slots": 3,
"description": "Default pool",
+ "include_deferred": True,
},
id="200 no update mask",
),
@@ -525,37 +582,50 @@ class TestModifyDefaultPool(TestBasePoolEndpoints):
class TestPatchPoolWithUpdateMask(TestBasePoolEndpoints):
@pytest.mark.parametrize(
- "url, patch_json, expected_name, expected_slots",
+ "url, patch_json, expected_name, expected_slots, expected_include_deferred",
[
(
"api/v1/pools/test_pool?update_mask=name, slots",
{"name": "test_pool_a", "slots": 2},
"test_pool_a",
2,
+ False,
),
(
"api/v1/pools/test_pool?update_mask=name",
{"name": "test_pool_a", "slots": 2},
"test_pool_a",
3,
+ False,
),
(
"api/v1/pools/test_pool?update_mask=slots",
{"name": "test_pool_a", "slots": 2},
"test_pool",
2,
+ False,
),
(
"api/v1/pools/test_pool?update_mask=slots",
{"slots": 2},
"test_pool",
2,
+ False,
+ ),
+ (
+ "api/v1/pools/test_pool?update_mask=include_deferred",
+ {"include_deferred": True},
+ "test_pool",
+ 3,
+ True,
),
],
)
@provide_session
- def test_response_200(self, url, patch_json, expected_name, expected_slots, session):
- pool = Pool(pool="test_pool", slots=3)
+ def test_response_200(
+ self, url, patch_json, expected_name, expected_slots, expected_include_deferred, session
+ ):
+ pool = Pool(pool="test_pool", slots=3, include_deferred=False)
session.add(pool)
session.commit()
response = self.client.patch(url, json=patch_json, environ_overrides={"REMOTE_USER": "test"})
@@ -567,8 +637,10 @@ class TestPatchPoolWithUpdateMask(TestBasePoolEndpoints):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": expected_slots,
"description": None,
+ "include_deferred": expected_include_deferred,
} == response.json
@pytest.mark.parametrize(
@@ -602,7 +674,7 @@ class TestPatchPoolWithUpdateMask(TestBasePoolEndpoints):
)
@provide_session
def test_response_400(self, error_detail, url, patch_json, session):
- pool = Pool(pool="test_pool", slots=3)
+ pool = Pool(pool="test_pool", slots=3, include_deferred=False)
session.add(pool)
session.commit()
response = self.client.patch(url, json=patch_json, environ_overrides={"REMOTE_USER": "test"})
diff --git a/tests/api_connexion/schemas/test_pool_schemas.py b/tests/api_connexion/schemas/test_pool_schemas.py
index f0eb0c0a49..9c5a78994d 100644
--- a/tests/api_connexion/schemas/test_pool_schemas.py
+++ b/tests/api_connexion/schemas/test_pool_schemas.py
@@ -31,7 +31,7 @@ class TestPoolSchema:
@provide_session
def test_serialize(self, session):
- pool_model = Pool(pool="test_pool", slots=2)
+ pool_model = Pool(pool="test_pool", slots=2, include_deferred=False)
session.add(pool_model)
session.commit()
pool_instance = session.query(Pool).filter(Pool.pool == pool_model.pool).first()
@@ -43,13 +43,15 @@ class TestPoolSchema:
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 2,
"description": None,
+ "include_deferred": False,
}
@provide_session
def test_deserialize(self, session):
- pool_dict = {"name": "test_pool", "slots": 3}
+ pool_dict = {"name": "test_pool", "slots": 3, "include_deferred": True}
deserialized_pool = pool_schema.load(pool_dict, session=session)
assert not isinstance(deserialized_pool, Pool) # Checks if load_instance is set to True
@@ -62,8 +64,8 @@ class TestPoolCollectionSchema:
clear_db_pools()
def test_serialize(self):
- pool_model_a = Pool(pool="test_pool_a", slots=3)
- pool_model_b = Pool(pool="test_pool_b", slots=3)
+ pool_model_a = Pool(pool="test_pool_a", slots=3, include_deferred=False)
+ pool_model_b = Pool(pool="test_pool_b", slots=3, include_deferred=True)
instance = PoolCollection(pools=[pool_model_a, pool_model_b], total_entries=2)
assert {
"pools": [
@@ -74,8 +76,10 @@ class TestPoolCollectionSchema:
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": None,
+ "include_deferred": False,
},
{
"name": "test_pool_b",
@@ -84,8 +88,10 @@ class TestPoolCollectionSchema:
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 3,
"description": None,
+ "include_deferred": True,
},
],
"total_entries": 2,
diff --git a/tests/api_connexion/test_auth.py b/tests/api_connexion/test_auth.py
index 35491ad7b2..b4cdc4a2bc 100644
--- a/tests/api_connexion/test_auth.py
+++ b/tests/api_connexion/test_auth.py
@@ -78,8 +78,10 @@ class TestBasicAuth(BaseTestAuth):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 128,
"description": "Default pool",
+ "include_deferred": False,
},
],
"total_entries": 1,
@@ -153,8 +155,10 @@ class TestSessionAuth(BaseTestAuth):
"running_slots": 0,
"queued_slots": 0,
"scheduled_slots": 0,
+ "deferred_slots": 0,
"open_slots": 128,
"description": "Default pool",
+ "include_deferred": False,
},
],
"total_entries": 1,
diff --git a/tests/api_experimental/client/test_local_client.py b/tests/api_experimental/client/test_local_client.py
index 575400a115..6724610670 100644
--- a/tests/api_experimental/client/test_local_client.py
+++ b/tests/api_experimental/client/test_local_client.py
@@ -181,22 +181,26 @@ class TestLocalClient:
assert session.query(DagModel).filter(DagModel.dag_id == key).count() == 0
def test_get_pool(self):
- self.client.create_pool(name="foo", slots=1, description="")
+ self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
pool = self.client.get_pool(name="foo")
- assert pool == ("foo", 1, "")
+ assert pool == ("foo", 1, "", False)
def test_get_pool_non_existing_raises(self):
with pytest.raises(PoolNotFound):
self.client.get_pool(name="foo")
def test_get_pools(self):
- self.client.create_pool(name="foo1", slots=1, description="")
- self.client.create_pool(name="foo2", slots=2, description="")
+ self.client.create_pool(name="foo1", slots=1, description="", include_deferred=False)
+ self.client.create_pool(name="foo2", slots=2, description="", include_deferred=True)
pools = sorted(self.client.get_pools(), key=lambda p: p[0])
- assert pools == [("default_pool", 128, "Default pool"), ("foo1", 1, ""), ("foo2", 2, "")]
+ assert pools == [
+ ("default_pool", 128, "Default pool", False),
+ ("foo1", 1, "", False),
+ ("foo2", 2, "", True),
+ ]
def test_create_pool(self):
- pool = self.client.create_pool(name="foo", slots=1, description="")
+ pool = self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
assert pool == ("foo", 1, "")
with create_session() as session:
assert session.query(Pool).count() == 2
@@ -207,6 +211,7 @@ class TestLocalClient:
name="foo",
slots="foo",
description="",
+ include_deferred=True,
)
def test_create_pool_name_too_long(self):
@@ -219,10 +224,11 @@ class TestLocalClient:
name=long_name,
slots=5,
description="",
+ include_deferred=False,
)
def test_delete_pool(self):
- self.client.create_pool(name="foo", slots=1, description="")
+ self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
with create_session() as session:
assert session.query(Pool).count() == 2
self.client.delete_pool(name="foo")
diff --git a/tests/api_experimental/common/experimental/test_pool.py b/tests/api_experimental/common/experimental/test_pool.py
index 4e6bdda7ac..d00f318ffa 100644
--- a/tests/api_experimental/common/experimental/test_pool.py
+++ b/tests/api_experimental/common/experimental/test_pool.py
@@ -45,6 +45,7 @@ class TestPool:
pool=name,
slots=i,
description=name,
+ include_deferred=True,
)
self.pools.append(pool)
with create_session() as session:
@@ -77,10 +78,12 @@ class TestPool:
assert session.query(models.Pool).count() == self.TOTAL_POOL_COUNT + 1
def test_create_pool_existing(self):
- pool = pool_api.create_pool(name=self.pools[0].pool, slots=5, description="")
- assert pool.pool == self.pools[0].pool
+ pool = pool_api.create_pool(name=self.pools[1].pool, slots=5, description="")
+ assert pool.pool == self.pools[1].pool
assert pool.slots == 5
assert pool.description == ""
+ # do not change include_deferred on an existing pool
+ assert pool.include_deferred is True
with create_session() as session:
assert session.query(models.Pool).count() == self.TOTAL_POOL_COUNT
diff --git a/tests/cli/commands/test_pool_command.py b/tests/cli/commands/test_pool_command.py
index ec175912d8..ca76a96514 100644
--- a/tests/cli/commands/test_pool_command.py
+++ b/tests/cli/commands/test_pool_command.py
@@ -19,7 +19,6 @@ from __future__ import annotations
import io
import json
-import os
from contextlib import redirect_stdout
import pytest
@@ -67,6 +66,18 @@ class TestCliPools:
pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
assert self.session.query(Pool).count() == 2
+ def test_pool_update_deferred(self):
+ pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
+ assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
+ pool_command.pool_set(
+ self.parser.parse_args(["pools", "set", "foo", "1", "test", "--include-deferred"])
+ )
+ assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is True
+
+ pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
+ assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
def test_pool_get(self):
pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
pool_command.pool_get(self.parser.parse_args(["pools", "get", "foo"]))
@@ -80,39 +91,57 @@ class TestCliPools:
with pytest.raises(SystemExit):
pool_command.pool_import(self.parser.parse_args(["pools", "import", "nonexistent.json"]))
- def test_pool_import_invalid_json(self):
- with open("pools_import_invalid.json", mode="w") as file:
+ def test_pool_import_invalid_json(self, tmp_path):
+ invalid_pool_import_file_path = tmp_path / "pools_import_invalid.json"
+ with open(invalid_pool_import_file_path, mode="w") as file:
file.write("not valid json")
with pytest.raises(SystemExit):
- pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import_invalid.json"]))
-
- def test_pool_import_invalid_pools(self):
- pool_config_input = {"foo": {"description": "foo_test"}}
- with open("pools_import_invalid.json", mode="w") as file:
+ pool_command.pool_import(
+ self.parser.parse_args(["pools", "import", str(invalid_pool_import_file_path)])
+ )
+
+ def test_pool_import_invalid_pools(self, tmp_path):
+ invalid_pool_import_file_path = tmp_path / "pools_import_invalid.json"
+ pool_config_input = {"foo": {"description": "foo_test", "include_deferred": False}}
+ with open(invalid_pool_import_file_path, mode="w") as file:
json.dump(pool_config_input, file)
with pytest.raises(SystemExit):
- pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import_invalid.json"]))
+ pool_command.pool_import(
+ self.parser.parse_args(["pools", "import", str(invalid_pool_import_file_path)])
+ )
- def test_pool_import_export(self):
- # Create two pools first
+ def test_pool_import_backwards_compatibility(self, tmp_path):
+ pool_import_file_path = tmp_path / "pools_import.json"
pool_config_input = {
+ # JSON before version 2.7.0 does not contain `include_deferred`
"foo": {"description": "foo_test", "slots": 1},
- "default_pool": {"description": "Default pool", "slots": 128},
- "baz": {"description": "baz_test", "slots": 2},
}
- with open("pools_import.json", mode="w") as file:
+ with open(pool_import_file_path, mode="w") as file:
+ json.dump(pool_config_input, file)
+
+ pool_command.pool_import(self.parser.parse_args(["pools", "import", str(pool_import_file_path)]))
+
+ assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
+ def test_pool_import_export(self, tmp_path):
+ pool_import_file_path = tmp_path / "pools_import.json"
+ pool_export_file_path = tmp_path / "pools_export.json"
+ pool_config_input = {
+ "foo": {"description": "foo_test", "slots": 1, "include_deferred": True},
+ "default_pool": {"description": "Default pool", "slots": 128, "include_deferred": False},
+ "baz": {"description": "baz_test", "slots": 2, "include_deferred": False},
+ }
+ with open(pool_import_file_path, mode="w") as file:
json.dump(pool_config_input, file)
# Import json
- pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import.json"]))
+ pool_command.pool_import(self.parser.parse_args(["pools", "import", str(pool_import_file_path)]))
# Export json
- pool_command.pool_export(self.parser.parse_args(["pools", "export", "pools_export.json"]))
+ pool_command.pool_export(self.parser.parse_args(["pools", "export", str(pool_export_file_path)]))
- with open("pools_export.json") as file:
+ with open(pool_export_file_path) as file:
pool_config_output = json.load(file)
assert pool_config_input == pool_config_output, "Input and output pool files are not same"
- os.remove("pools_import.json")
- os.remove("pools_export.json")
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
index a96c5cbddf..eac5fe6802 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -813,7 +813,7 @@ class TestLogsfromTaskRunCommand:
clear_db_pools()
with create_session() as session:
- pool = Pool(pool=pool_name, slots=1)
+ pool = Pool(pool=pool_name, slots=1, include_deferred=False)
session.add(pool)
session.commit()
diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py
index a11b8d5688..e4efd8682e 100644
--- a/tests/jobs/test_backfill_job.py
+++ b/tests/jobs/test_backfill_job.py
@@ -598,6 +598,7 @@ class TestBackfillJob:
pool = Pool(
pool="pool_with_two_slots",
slots=slots,
+ include_deferred=False,
)
session.add(pool)
session.commit()
@@ -961,7 +962,7 @@ class TestBackfillJob:
Test that queued tasks are executed by BackfillJobRunner
"""
session = settings.Session()
- pool = Pool(pool="test_backfill_pooled_task_pool", slots=1)
+ pool = Pool(pool="test_backfill_pooled_task_pool", slots=1, include_deferred=False)
session.add(pool)
session.commit()
session.close()
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 3a3d3f55b4..ed09d9cd18 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -605,8 +605,8 @@ class TestSchedulerJob:
for ti in tis:
ti.state = State.SCHEDULED
session.merge(ti)
- pool = Pool(pool="a", slots=1, description="haha")
- pool2 = Pool(pool="b", slots=100, description="haha")
+ pool = Pool(pool="a", slots=1, description="haha", include_deferred=False)
+ pool2 = Pool(pool="b", slots=100, description="haha", include_deferred=False)
session.add(pool)
session.add(pool2)
session.flush()
@@ -732,8 +732,8 @@ class TestSchedulerJob:
dag_id = "SchedulerJobTest.test_find_executable_task_instances_order_priority_with_pools"
- session.add(Pool(pool="pool1", slots=32))
- session.add(Pool(pool="pool2", slots=32))
+ session.add(Pool(pool="pool1", slots=32, include_deferred=False))
+ session.add(Pool(pool="pool2", slots=32, include_deferred=False))
with dag_maker(dag_id=dag_id, max_active_tasks=2):
op1 = EmptyOperator(task_id="dummy1", priority_weight=1, pool="pool1")
@@ -888,7 +888,12 @@ class TestSchedulerJob:
ti = dr.task_instances[0]
ti.state = State.SCHEDULED
session.merge(ti)
- infinite_pool = Pool(pool="infinite_pool", slots=-1, description="infinite pool")
+ infinite_pool = Pool(
+ pool="infinite_pool",
+ slots=-1,
+ description="infinite pool",
+ include_deferred=False,
+ )
session.add(infinite_pool)
session.commit()
@@ -913,7 +918,7 @@ class TestSchedulerJob:
ti = dr.task_instances[1]
ti.state = State.SCHEDULED
session.merge(ti)
- some_pool = Pool(pool="some_pool", slots=2, description="my pool")
+ some_pool = Pool(pool="some_pool", slots=2, description="my pool", include_deferred=False)
session.add(some_pool)
session.commit()
with caplog.at_level(logging.WARNING):
@@ -1336,8 +1341,8 @@ class TestSchedulerJob:
self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
session = settings.Session()
- pool1 = Pool(pool="pool1", slots=1)
- pool2 = Pool(pool="pool2", slots=1)
+ pool1 = Pool(pool="pool1", slots=1, include_deferred=False)
+ pool2 = Pool(pool="pool2", slots=1, include_deferred=False)
session.add(pool1)
session.add(pool2)
@@ -2578,7 +2583,7 @@ class TestSchedulerJob:
)
session = settings.Session()
- pool = Pool(pool="test_scheduler_verify_pool_full", slots=1)
+ pool = Pool(pool="test_scheduler_verify_pool_full", slots=1, include_deferred=False)
session.add(pool)
session.flush()
@@ -2619,7 +2624,7 @@ class TestSchedulerJob:
bash_command="echo hi",
)
- pool = Pool(pool="test_scheduler_verify_pool_full_2_slots_per_task", slots=6)
+ pool = Pool(pool="test_scheduler_verify_pool_full_2_slots_per_task", slots=6, include_deferred=False)
session.add(pool)
session.flush()
@@ -2673,8 +2678,8 @@ class TestSchedulerJob:
dag_d2 = dag_maker.dag
session = settings.Session()
- pool_p1 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p1", slots=1)
- pool_p2 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p2", slots=10)
+ pool_p1 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p1", slots=1, include_deferred=False)
+ pool_p2 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p2", slots=10, include_deferred=False)
session.add(pool_p1)
session.add(pool_p2)
session.flush()
@@ -2751,7 +2756,7 @@ class TestSchedulerJob:
)
session = settings.Session()
- pool = Pool(pool="test_scheduler_verify_priority_and_slots", slots=2)
+ pool = Pool(pool="test_scheduler_verify_priority_and_slots", slots=2, include_deferred=False)
session.add(pool)
session.flush()
diff --git a/tests/models/test_pool.py b/tests/models/test_pool.py
index 2e142814c6..32d30b99bb 100644
--- a/tests/models/test_pool.py
+++ b/tests/models/test_pool.py
@@ -55,6 +55,7 @@ class TestPool:
pool=name,
slots=i,
description=name,
+ include_deferred=False,
)
self.pools.append(pool)
with create_session() as session:
@@ -64,23 +65,27 @@ class TestPool:
self.clean_db()
def test_open_slots(self, dag_maker):
- pool = Pool(pool="test_pool", slots=5)
+ pool = Pool(pool="test_pool", slots=5, include_deferred=False)
with dag_maker(
dag_id="test_open_slots",
start_date=DEFAULT_DATE,
):
op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+ op3 = EmptyOperator(task_id="dummy3", pool="test_pool")
dag_maker.create_dagrun()
ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
+ ti3 = TI(task=op3, execution_date=DEFAULT_DATE)
ti1.state = State.RUNNING
ti2.state = State.QUEUED
+ ti3.state = State.DEFERRED
session = settings.Session()
session.add(pool)
session.merge(ti1)
session.merge(ti2)
+ session.merge(ti3)
session.commit()
session.close()
@@ -88,23 +93,69 @@ class TestPool:
assert 1 == pool.running_slots()
assert 1 == pool.queued_slots()
assert 2 == pool.occupied_slots()
+ assert 1 == pool.deferred_slots()
assert {
"default_pool": {
"open": 128,
"queued": 0,
"total": 128,
"running": 0,
+ "deferred": 0,
},
"test_pool": {
"open": 3,
"queued": 1,
"running": 1,
+ "deferred": 1,
+ "total": 5,
+ },
+ } == pool.slots_stats()
+
+ def test_open_slots_including_deferred(self, dag_maker):
+ pool = Pool(pool="test_pool", slots=5, include_deferred=True)
+ with dag_maker(
+ dag_id="test_open_slots_including_deferred",
+ start_date=DEFAULT_DATE,
+ ):
+ op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
+ op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+ dag_maker.create_dagrun()
+ ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
+ ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
+ ti1.state = State.RUNNING
+ ti2.state = State.DEFERRED
+
+ session = settings.Session()
+ session.add(pool)
+ session.merge(ti1)
+ session.merge(ti2)
+ session.commit()
+ session.close()
+
+ assert 3 == pool.open_slots()
+ assert 1 == pool.running_slots()
+ assert 0 == pool.queued_slots()
+ assert 1 == pool.deferred_slots()
+ assert 2 == pool.occupied_slots()
+ assert {
+ "default_pool": {
+ "open": 128,
+ "queued": 0,
+ "total": 128,
+ "running": 0,
+ "deferred": 0,
+ },
+ "test_pool": {
+ "open": 3,
+ "queued": 0,
+ "running": 1,
+ "deferred": 1,
"total": 5,
},
} == pool.slots_stats()
def test_infinite_slots(self, dag_maker):
- pool = Pool(pool="test_pool", slots=-1)
+ pool = Pool(pool="test_pool", slots=-1, include_deferred=False)
with dag_maker(
dag_id="test_infinite_slots",
):
@@ -133,12 +184,14 @@ class TestPool:
"queued": 0,
"total": 128,
"running": 0,
+ "deferred": 0,
},
"test_pool": {
"open": float("inf"),
"queued": 1,
"running": 1,
"total": float("inf"),
+ "deferred": 0,
},
} == pool.slots_stats()
@@ -170,6 +223,7 @@ class TestPool:
"queued": 2,
"total": 5,
"running": 1,
+ "deferred": 0,
}
} == Pool.slots_stats()
@@ -194,18 +248,22 @@ class TestPool:
def test_create_pool(self, session):
self.add_pools()
- pool = Pool.create_or_update_pool(name="foo", slots=5, description="")
+ pool = Pool.create_or_update_pool(name="foo", slots=5, description="", include_deferred=True)
assert pool.pool == "foo"
assert pool.slots == 5
assert pool.description == ""
+ assert pool.include_deferred is True
assert session.query(Pool).count() == self.TOTAL_POOL_COUNT + 1
def test_create_pool_existing(self, session):
self.add_pools()
- pool = Pool.create_or_update_pool(name=self.pools[0].pool, slots=5, description="")
+ pool = Pool.create_or_update_pool(
+ name=self.pools[0].pool, slots=5, description="", include_deferred=False
+ )
assert pool.pool == self.pools[0].pool
assert pool.slots == 5
assert pool.description == ""
+ assert pool.include_deferred is False
assert session.query(Pool).count() == self.TOTAL_POOL_COUNT
def test_delete_pool(self, session):
@@ -223,7 +281,9 @@ class TestPool:
Pool.delete_pool(Pool.DEFAULT_POOL_NAME)
def test_is_default_pool(self):
- pool = Pool.create_or_update_pool(name="not_default_pool", slots=1, description="test")
+ pool = Pool.create_or_update_pool(
+ name="not_default_pool", slots=1, description="test", include_deferred=False
+ )
default_pool = Pool.get_default_pool()
assert not Pool.is_default_pool(id=pool.id)
assert Pool.is_default_pool(str(default_pool.id))
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9e7f5eb48a..ac2f7f7174 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -96,7 +96,7 @@ from tests.test_utils.mock_operators import MockOperator
@pytest.fixture
def test_pool():
with create_session() as session:
- test_pool = Pool(pool="test_pool", slots=1)
+ test_pool = Pool(pool="test_pool", slots=1, include_deferred=False)
session.add(test_pool)
session.flush()
yield test_pool
diff --git a/tests/operators/test_subdag_operator.py b/tests/operators/test_subdag_operator.py
index 8df1ee9fe3..c814e64a7e 100644
--- a/tests/operators/test_subdag_operator.py
+++ b/tests/operators/test_subdag_operator.py
@@ -88,8 +88,8 @@ class TestSubDagOperator:
subdag = DAG("parent.child", default_args=default_args)
session = airflow.settings.Session()
- pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1)
- pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10)
+ pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False)
+ pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False)
session.add(pool_1)
session.add(pool_10)
session.commit()
@@ -116,8 +116,8 @@ class TestSubDagOperator:
subdag = DAG("parent.child", default_args=default_args)
session = airflow.settings.Session()
- pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1)
- pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10)
+ pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False)
+ pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False)
session.add(pool_1)
session.add(pool_10)
session.commit()
diff --git a/tests/ti_deps/deps/test_pool_slots_available_dep.py b/tests/ti_deps/deps/test_pool_slots_available_dep.py
index cba53ef65b..08abba29e4 100644
--- a/tests/ti_deps/deps/test_pool_slots_available_dep.py
+++ b/tests/ti_deps/deps/test_pool_slots_available_dep.py
@@ -23,6 +23,7 @@ from airflow.models import Pool
from airflow.ti_deps.dependencies_states import EXECUTION_STATES
from airflow.ti_deps.deps.pool_slots_available_dep import PoolSlotsAvailableDep
from airflow.utils.session import create_session
+from airflow.utils.state import TaskInstanceState
from tests.test_utils import db
@@ -30,8 +31,9 @@ class TestPoolSlotsAvailableDep:
def setup_method(self):
db.clear_db_pools()
with create_session() as session:
- test_pool = Pool(pool="test_pool")
- session.add(test_pool)
+ test_pool = Pool(pool="test_pool", include_deferred=False)
+ test_includes_deferred_pool = Pool(pool="test_includes_deferred_pool", include_deferred=True)
+ session.add_all([test_pool, test_includes_deferred_pool])
session.commit()
def teardown_method(self):
@@ -53,6 +55,13 @@ class TestPoolSlotsAvailableDep:
ti = Mock(pool="test_pool", state=state, pool_slots=1)
assert PoolSlotsAvailableDep().is_met(ti=ti)
+ @patch("airflow.models.Pool.open_slots", return_value=0)
+ def test_deferred_pooled_task_pass(self, mock_open_slots):
+ ti = Mock(pool="test_includes_deferred_pool", state=TaskInstanceState.DEFERRED, pool_slots=1)
+ assert PoolSlotsAvailableDep().is_met(ti=ti)
+ ti_to_fail = Mock(pool="test_pool", state=TaskInstanceState.DEFERRED, pool_slots=1)
+ assert not PoolSlotsAvailableDep().is_met(ti=ti_to_fail)
+
def test_task_with_nonexistent_pool(self):
ti = Mock(pool="nonexistent_pool", pool_slots=1)
assert not PoolSlotsAvailableDep().is_met(ti=ti)
diff --git a/tests/www/api/experimental/test_endpoints.py b/tests/www/api/experimental/test_endpoints.py
index 86a145f5b8..d42cfb2237 100644
--- a/tests/www/api/experimental/test_endpoints.py
+++ b/tests/www/api/experimental/test_endpoints.py
@@ -383,6 +383,7 @@ class TestPoolApiExperimental(TestBase):
pool=name,
slots=i,
description=name,
+ include_deferred=False,
)
self.session.add(pool)
self.pools.append(pool)
@@ -434,6 +435,7 @@ class TestPoolApiExperimental(TestBase):
assert pool["pool"] == "foo"
assert pool["slots"] == 1
assert pool["description"] == ""
+ assert pool["include_deferred"] is False
assert self._get_pool_count() == self.TOTAL_POOL_COUNT + 1
def test_create_pool_with_bad_name(self):
diff --git a/tests/www/views/test_views_pool.py b/tests/www/views/test_views_pool.py
index aebcd2c844..e79ec46c56 100644
--- a/tests/www/views/test_views_pool.py
+++ b/tests/www/views/test_views_pool.py
@@ -29,6 +29,7 @@ POOL = {
"pool": "test-pool",
"slots": 777,
"description": "test-pool-description",
+ "include_deferred": False,
}
@@ -90,9 +91,13 @@ def test_list(app, admin_client, pool_factory):
url = flask.url_for("TaskInstanceModelView.list", _flt_3_pool="test-pool", _flt_3_state="scheduled")
scheduled_tag = markupsafe.Markup("<a href='{url}'>{slots}</a>").format(url=url, slots=0)
+ url = flask.url_for("TaskInstanceModelView.list", _flt_3_pool="test-pool", _flt_3_state="deferred")
+ deferred_tag = markupsafe.Markup("<a href='{url}'>{slots}</a>").format(url=url, slots=0)
+
check_content_in_response(used_tag, resp)
check_content_in_response(queued_tag, resp)
check_content_in_response(scheduled_tag, resp)
+ check_content_in_response(deferred_tag, resp)
def test_pool_muldelete(session, admin_client, pool_factory):
[airflow] 04/09: Fix flaky sqlite tests with `test_xcom_map_nest` hopefully (#33145)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 4e817c46b4c94d44b5c50665f23f19d7284b42be
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Sat Aug 5 23:15:49 2023 +0200
Fix flaky sqlite tests with `test_xcom_map_nest` hopefully (#33145)
Recently sqlite started to fail randomly during teardown of
`test_xcom_map_nest` or `test_xcom_map_zip_nest`. This happpened
after adding `test_xcom_map_zip_nest` . It looked very strange:
```
sqlite3.ProgrammingError: SQLite objects created in a thread can only be
used in that same thread
```
By analysing possible reasons it seems that it was a side effect
of the existing `test_xcom_map_nest` that allocated a new
session in the run method of task instance rather than pass
the sesion that is created and torrn down in the pytest fixture.
The hypothesis is that the session created in the ``test_xcom_map_nest``
were being reclaimed and closed while the `test_xcom_map_zip_nest` test
was already starting in a different thread started by Pytest.
The fix is to pass the session object to run method of the taskinstance
in the ``test_xcom_map_nest`` test.
(cherry picked from commit d1d6fc994d46aaed9c801162595cae91a1ffc19c)
---
tests/models/test_xcom_arg_map.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/tests/models/test_xcom_arg_map.py b/tests/models/test_xcom_arg_map.py
index e88c1a81d8..92e14da567 100644
--- a/tests/models/test_xcom_arg_map.py
+++ b/tests/models/test_xcom_arg_map.py
@@ -251,7 +251,7 @@ def test_xcom_map_nest(dag_maker, session):
# Run "push".
decision = dr.task_instance_scheduling_decisions(session=session)
for ti in decision.schedulable_tis:
- ti.run()
+ ti.run(session=session)
session.flush()
session.commit()
@@ -262,7 +262,7 @@ def test_xcom_map_nest(dag_maker, session):
# Now "pull" should apply the mapping functions in order.
decision = dr.task_instance_scheduling_decisions(session=session)
for ti in decision.schedulable_tis:
- ti.run()
+ ti.run(session=session)
assert results == {"aa", "bb", "cc"}
[airflow] 08/09: Change type of defualts for include_deferred (#33271)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 2ff8d4d6133300b1c8017ed8f5d9e981921b95c9
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Wed Aug 9 23:00:21 2023 +0200
Change type of defualts for include_deferred (#33271)
We should change the defaults to client defaults for MSSQL compatibility
The change in #32709 was not liked by MSSQL
Co-authored-by: Ephraim Anierobi <sp...@gmail.com>
(cherry picked from commit d989e9dba5899e87780df9a8c5994ed1e3f8a776)
---
.../versions/0128_2_7_0_add_include_deferred_column_to_pool.py | 4 +---
docs/apache-airflow/img/airflow_erd.sha256 | 2 +-
2 files changed, 2 insertions(+), 4 deletions(-)
diff --git a/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
index ea10320480..bbafc68057 100644
--- a/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
+++ b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
@@ -39,9 +39,7 @@ airflow_version = "2.7.0"
def upgrade():
"""Apply add include_deferred column to pool"""
with op.batch_alter_table("slot_pool") as batch_op:
- batch_op.add_column(
- sa.Column("include_deferred", sa.Boolean, nullable=False, server_default=sa.false())
- )
+ batch_op.add_column(sa.Column("include_deferred", sa.Boolean, default=False))
def downgrade():
diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256
index 5f9515fc19..a7704c8fc5 100644
--- a/docs/apache-airflow/img/airflow_erd.sha256
+++ b/docs/apache-airflow/img/airflow_erd.sha256
@@ -1 +1 @@
-3420f98a9dfa230d76be4da57bc4a3641cb929b33b9eafe7ffa41fb4fcc4d283
\ No newline at end of file
+52b0e45032fb1c60274374e0114cf199dfb759adc0a4fe32964136347457506a
\ No newline at end of file
[airflow] 01/09: Fix missing related PRs when cherry-picking (#33261)
Posted by ep...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-7-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit a57cc8b0c05c1fa0121cdc0321e7ae75c82d05a5
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Wed Aug 9 17:41:06 2023 +0200
Fix missing related PRs when cherry-picking (#33261)
The `airflow-github` tooling had a check if PR has been already
merged, but the PR did not work well when there was already a PR
cherry-picked that referred the PR in question in their log message.
The check was done via `git log --grep=#PR` essentially and so
it happenes that `--grep` matches not only the first line of
the message but the whole log message. So if the log message
contained reference to the cherry-picked PR it was seen as merged.
This change adds ( ) around and $ to the regexp to match not
only the PR number but also parentheses around.
(cherry picked from commit eb07618d29ce957ef4401a98f25449a23d0135ad)
---
dev/airflow-github | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/dev/airflow-github b/dev/airflow-github
index 33e543a77f..115ed48d19 100755
--- a/dev/airflow-github
+++ b/dev/airflow-github
@@ -96,7 +96,7 @@ def get_commit_in_main_associated_with_pr(repo: git.Repo, issue: Issue) -> str |
def is_cherrypicked(repo: git.Repo, issue: Issue, previous_version: str | None = None) -> bool:
"""Check if a given issue is cherry-picked in the current branch or not"""
- log_args = ["--format=%H", f"--grep=#{issue.number}"]
+ log_args = ["--format=%H", f"--grep=(#{issue.number})$"]
if previous_version:
log_args.append(previous_version + "..")
log = repo.git.log(*log_args)