You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "uranusjr (via GitHub)" <gi...@apache.org> on 2024/09/16 06:43:51 UTC
[PR] Refactor bulk_save_to_db [airflow]
uranusjr opened a new pull request, #42245:
URL: https://github.com/apache/airflow/pull/42245
This function collects information from DAG objects, and creates/updates database rows against them. However, it handles A LOT of information, reading a lot of objects, touching a lot of models. The function is not very readable.
A new module has been introduced in airflow.dag_processing.collection to encapsulate the logic previously in bulk_save_to_db.
Some loops are broken down into multiple loops, so each loop does not do too much (which leads to a lot of long-living variables that reduce readability). Not much is changed aside from that, just mostly splitting things into separate steps to make things clearer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766116173
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
Review Comment:
sounds good to me š
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766112798
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
Review Comment:
Tags are re-added later in the process (now inside `update_orm_dags`). I believe SQLAlchemy is smart enough to figure out what objects are the same and do the appropriate updates automatically.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766113548
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
Review Comment:
Technically they can even be combined into one single function; I tried that, but IMO this is more readable. Some code duplication (and likely a little slower), but worthwhile for maintainability.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766114722
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
+ for dag in dags:
+ for _, dataset in dag.timetable.dataset_condition.iter_datasets():
+ yield dataset
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, Dataset):
+ yield obj
+
+
+def _find_all_dataset_aliases(dags: Iterable[DAG]) -> Iterator[DatasetAlias]:
+ for dag in dags:
+ for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases():
+ yield alias
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, DatasetAlias):
+ yield obj
+
+
+class DatasetCollection(NamedTuple):
+ """Datasets collected from DAGs."""
+
+ schedule_dataset_references: dict[str, list[Dataset]]
+ schedule_dataset_alias_references: dict[str, list[DatasetAlias]]
+ outlet_references: dict[str, list[tuple[str, Dataset]]]
+ datasets: dict[str, Dataset]
+ dataset_aliases: dict[str, DatasetAlias]
+
+ @classmethod
+ def collect(cls, dags: dict[str, DAG]) -> Self:
+ coll = cls(
+ schedule_dataset_references={
+ dag_id: [dataset for _, dataset in dag.timetable.dataset_condition.iter_datasets()]
+ for dag_id, dag in dags.items()
+ },
+ schedule_dataset_alias_references={
+ dag_id: [alias for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases()]
+ for dag_id, dag in dags.items()
+ },
+ outlet_references={
+ dag_id: [
+ (task_id, outlet)
+ for task_id, task in dag.task_dict.items()
+ for outlet in task.outlets
+ if isinstance(outlet, Dataset)
+ ]
+ for dag_id, dag in dags.items()
+ },
+ datasets={dataset.uri: dataset for dataset in _find_all_datasets(dags.values())},
+ dataset_aliases={alias.name: alias for alias in _find_all_dataset_aliases(dags.values())},
+ )
+ return coll
+
+ def add_datasets(self, *, session: Session) -> dict[str, DatasetModel]:
Review Comment:
I use this name since the SQLAlchemy function that does the INSERT is `session.add(model)`. Adding `_to_db` to everything seems a bit cumbersomeā¦ Maybe we can rename the DatasetCollection class instead? DatasetModelOperation?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #42245:
URL: https://github.com/apache/airflow/pull/42245#issuecomment-2359722286
Also maybe we need a new method in DatasetManager for creating aliases. (Or even more for other things?)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr merged PR #42245:
URL: https://github.com/apache/airflow/pull/42245
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #42245:
URL: https://github.com/apache/airflow/pull/42245#issuecomment-2359715193
Note of self: Submit another PR to refactor `_register_dataset_changes`. There are at least a few db optimisations we can do.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766312781
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
Review Comment:
```suggestion
"""Build a select statement to retrieve the last automated run for each dag."""
```
##########
airflow/models/dag.py:
##########
@@ -2648,323 +2643,38 @@ def bulk_write_to_db(
if not dags:
return
- log.info("Sync %s DAGs", len(dags))
- dag_by_ids = {dag.dag_id: dag for dag in dags}
-
- dag_ids = set(dag_by_ids)
- query = (
- select(DagModel)
- .options(joinedload(DagModel.tags, innerjoin=False))
- .where(DagModel.dag_id.in_(dag_ids))
- .options(joinedload(DagModel.schedule_dataset_references))
- .options(joinedload(DagModel.schedule_dataset_alias_references))
- .options(joinedload(DagModel.task_outlet_dataset_references))
- )
- query = with_row_locks(query, of=DagModel, session=session)
- orm_dags: list[DagModel] = session.scalars(query).unique().all()
- existing_dags: dict[str, DagModel] = {x.dag_id: x for x in orm_dags}
- missing_dag_ids = dag_ids.difference(existing_dags.keys())
-
- for missing_dag_id in missing_dag_ids:
- orm_dag = DagModel(dag_id=missing_dag_id)
- dag = dag_by_ids[missing_dag_id]
- if dag.is_paused_upon_creation is not None:
- orm_dag.is_paused = dag.is_paused_upon_creation
- orm_dag.tags = []
- log.info("Creating ORM DAG for %s", dag.dag_id)
- session.add(orm_dag)
- orm_dags.append(orm_dag)
-
- latest_runs: dict[str, DagRun] = {}
- num_active_runs: dict[str, int] = {}
- # Skip these queries entirely if no DAGs can be scheduled to save time.
- if any(dag.timetable.can_be_scheduled for dag in dags):
- # Get the latest automated dag run for each existing dag as a single query (avoid n+1 query)
- query = cls._get_latest_runs_stmt(dags=list(existing_dags.keys()))
- latest_runs = {run.dag_id: run for run in session.scalars(query)}
-
- # Get number of active dagruns for all dags we are processing as a single query.
- num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dags, session=session)
-
- filelocs = []
-
- for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id):
- dag = dag_by_ids[orm_dag.dag_id]
- filelocs.append(dag.fileloc)
- orm_dag.fileloc = dag.fileloc
- orm_dag.owners = dag.owner
- orm_dag.is_active = True
- orm_dag.has_import_errors = False
- orm_dag.last_parsed_time = timezone.utcnow()
- orm_dag.default_view = dag.default_view
- orm_dag._dag_display_property_value = dag._dag_display_property_value
- orm_dag.description = dag.description
- orm_dag.max_active_tasks = dag.max_active_tasks
- orm_dag.max_active_runs = dag.max_active_runs
- orm_dag.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
- orm_dag.has_task_concurrency_limits = any(
- t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None
- for t in dag.tasks
- )
- orm_dag.timetable_summary = dag.timetable.summary
- orm_dag.timetable_description = dag.timetable.description
- orm_dag.dataset_expression = dag.timetable.dataset_condition.as_expression()
-
- orm_dag.processor_subdir = processor_subdir
-
- last_automated_run: DagRun | None = latest_runs.get(dag.dag_id)
- if last_automated_run is None:
- last_automated_data_interval = None
- else:
- last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
- if num_active_runs.get(dag.dag_id, 0) >= orm_dag.max_active_runs:
- orm_dag.next_dagrun_create_after = None
- else:
- orm_dag.calculate_dagrun_date_fields(dag, last_automated_data_interval)
-
- dag_tags = set(dag.tags or {})
- orm_dag_tags = list(orm_dag.tags or [])
- for orm_tag in orm_dag_tags:
- if orm_tag.name not in dag_tags:
- session.delete(orm_tag)
- orm_dag.tags.remove(orm_tag)
- orm_tag_names = {t.name for t in orm_dag_tags}
- for dag_tag in dag_tags:
- if dag_tag not in orm_tag_names:
- dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
- orm_dag.tags.append(dag_tag_orm)
- session.add(dag_tag_orm)
-
- orm_dag_links = orm_dag.dag_owner_links or []
- for orm_dag_link in orm_dag_links:
- if orm_dag_link not in dag.owner_links:
- session.delete(orm_dag_link)
- for owner_name, owner_link in dag.owner_links.items():
- dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
- session.add(dag_owner_orm)
-
- DagCode.bulk_sync_to_db(filelocs, session=session)
-
- from airflow.datasets import Dataset
- from airflow.models.dataset import (
- DagScheduleDatasetAliasReference,
- DagScheduleDatasetReference,
- DatasetModel,
- TaskOutletDatasetReference,
+ from airflow.dag_processing.collection import (
+ DatasetModelOperation,
+ collect_orm_dags,
+ create_orm_dag,
+ update_orm_dags,
)
- dag_references: dict[str, set[tuple[Literal["dataset", "dataset-alias"], str]]] = defaultdict(set)
- outlet_references = defaultdict(set)
- # We can't use a set here as we want to preserve order
- outlet_dataset_models: dict[DatasetModel, None] = {}
- input_dataset_models: dict[DatasetModel, None] = {}
- outlet_dataset_alias_models: set[DatasetAliasModel] = set()
- input_dataset_alias_models: set[DatasetAliasModel] = set()
-
- # here we go through dags and tasks to check for dataset references
- # if there are now None and previously there were some, we delete them
- # if there are now *any*, we add them to the above data structures, and
- # later we'll persist them to the database.
- for dag in dags:
- curr_orm_dag = existing_dags.get(dag.dag_id)
- if not (dataset_condition := dag.timetable.dataset_condition):
- if curr_orm_dag:
- if curr_orm_dag.schedule_dataset_references:
- curr_orm_dag.schedule_dataset_references = []
- if curr_orm_dag.schedule_dataset_alias_references:
- curr_orm_dag.schedule_dataset_alias_references = []
- else:
- for _, dataset in dataset_condition.iter_datasets():
- dag_references[dag.dag_id].add(("dataset", dataset.uri))
- input_dataset_models[DatasetModel.from_public(dataset)] = None
-
- for dataset_alias in dataset_condition.iter_dataset_aliases():
- dag_references[dag.dag_id].add(("dataset-alias", dataset_alias.name))
- input_dataset_alias_models.add(DatasetAliasModel.from_public(dataset_alias))
-
- curr_outlet_references = curr_orm_dag and curr_orm_dag.task_outlet_dataset_references
- for task in dag.tasks:
- dataset_outlets: list[Dataset] = []
- dataset_alias_outlets: list[DatasetAlias] = []
- for outlet in task.outlets:
- if isinstance(outlet, Dataset):
- dataset_outlets.append(outlet)
- elif isinstance(outlet, DatasetAlias):
- dataset_alias_outlets.append(outlet)
-
- if not dataset_outlets:
- if curr_outlet_references:
- this_task_outlet_refs = [
- x
- for x in curr_outlet_references
- if x.dag_id == dag.dag_id and x.task_id == task.task_id
- ]
- for ref in this_task_outlet_refs:
- curr_outlet_references.remove(ref)
-
- for d in dataset_outlets:
- outlet_dataset_models[DatasetModel.from_public(d)] = None
- outlet_references[(task.dag_id, task.task_id)].add(d.uri)
-
- for d_a in dataset_alias_outlets:
- outlet_dataset_alias_models.add(DatasetAliasModel.from_public(d_a))
-
- all_dataset_models = outlet_dataset_models
- all_dataset_models.update(input_dataset_models)
-
- # store datasets
- stored_dataset_models: dict[str, DatasetModel] = {}
- new_dataset_models: list[DatasetModel] = []
- for dataset in all_dataset_models:
- stored_dataset_model = session.scalar(
- select(DatasetModel).where(DatasetModel.uri == dataset.uri).limit(1)
- )
- if stored_dataset_model:
- # Some datasets may have been previously unreferenced, and therefore orphaned by the
- # scheduler. But if we're here, then we have found that dataset again in our DAGs, which
- # means that it is no longer an orphan, so set is_orphaned to False.
- stored_dataset_model.is_orphaned = expression.false()
- stored_dataset_models[stored_dataset_model.uri] = stored_dataset_model
- else:
- new_dataset_models.append(dataset)
- dataset_manager.create_datasets(dataset_models=new_dataset_models, session=session)
- stored_dataset_models.update(
- {dataset_model.uri: dataset_model for dataset_model in new_dataset_models}
- )
-
- del new_dataset_models
- del all_dataset_models
-
- # store dataset aliases
- all_datasets_alias_models = input_dataset_alias_models | outlet_dataset_alias_models
- stored_dataset_alias_models: dict[str, DatasetAliasModel] = {}
- new_dataset_alias_models: set[DatasetAliasModel] = set()
- if all_datasets_alias_models:
- all_dataset_alias_names = {
- dataset_alias_model.name for dataset_alias_model in all_datasets_alias_models
- }
-
- stored_dataset_alias_models = {
- dsa_m.name: dsa_m
- for dsa_m in session.scalars(
- select(DatasetAliasModel).where(DatasetAliasModel.name.in_(all_dataset_alias_names))
- ).fetchall()
- }
-
- if stored_dataset_alias_models:
- new_dataset_alias_models = {
- dataset_alias_model
- for dataset_alias_model in all_datasets_alias_models
- if dataset_alias_model.name not in stored_dataset_alias_models.keys()
- }
- else:
- new_dataset_alias_models = all_datasets_alias_models
-
- session.add_all(new_dataset_alias_models)
- session.flush()
- stored_dataset_alias_models.update(
- {
- dataset_alias_model.name: dataset_alias_model
- for dataset_alias_model in new_dataset_alias_models
- }
+ log.info("Sync %s DAGs", len(dags))
+ dags_by_ids = {dag.dag_id: dag for dag in dags}
+ del dags
+
+ orm_dags = collect_orm_dags(dags_by_ids, session=session)
+ orm_dags.update(
+ (dag_id, create_orm_dag(dag, session=session))
+ for dag_id, dag in dags_by_ids.items()
+ if dag_id not in orm_dags
)
- del new_dataset_alias_models
- del all_datasets_alias_models
+ update_orm_dags(dags_by_ids, orm_dags, processor_subdir=processor_subdir, session=session)
+ DagCode.bulk_sync_to_db((dag.fileloc for dag in dags_by_ids.values()), session=session)
Review Comment:
Should we somehow return the fileloc instead of iterating through the dags again to get the filelocs?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766116777
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
+ for dag in dags:
+ for _, dataset in dag.timetable.dataset_condition.iter_datasets():
+ yield dataset
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, Dataset):
+ yield obj
+
+
+def _find_all_dataset_aliases(dags: Iterable[DAG]) -> Iterator[DatasetAlias]:
+ for dag in dags:
+ for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases():
+ yield alias
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, DatasetAlias):
+ yield obj
+
+
+class DatasetCollection(NamedTuple):
+ """Datasets collected from DAGs."""
+
+ schedule_dataset_references: dict[str, list[Dataset]]
+ schedule_dataset_alias_references: dict[str, list[DatasetAlias]]
+ outlet_references: dict[str, list[tuple[str, Dataset]]]
+ datasets: dict[str, Dataset]
+ dataset_aliases: dict[str, DatasetAlias]
+
+ @classmethod
+ def collect(cls, dags: dict[str, DAG]) -> Self:
+ coll = cls(
+ schedule_dataset_references={
+ dag_id: [dataset for _, dataset in dag.timetable.dataset_condition.iter_datasets()]
+ for dag_id, dag in dags.items()
+ },
+ schedule_dataset_alias_references={
+ dag_id: [alias for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases()]
+ for dag_id, dag in dags.items()
+ },
+ outlet_references={
+ dag_id: [
+ (task_id, outlet)
+ for task_id, task in dag.task_dict.items()
+ for outlet in task.outlets
+ if isinstance(outlet, Dataset)
+ ]
+ for dag_id, dag in dags.items()
+ },
+ datasets={dataset.uri: dataset for dataset in _find_all_datasets(dags.values())},
+ dataset_aliases={alias.name: alias for alias in _find_all_dataset_aliases(dags.values())},
+ )
+ return coll
+
+ def add_datasets(self, *, session: Session) -> dict[str, DatasetModel]:
Review Comment:
> Maybe we can rename the DatasetCollection class instead? DatasetModelOperation?
`DatasetModelOperation` sounds like a great idea!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766358139
##########
airflow/models/dag.py:
##########
@@ -2648,323 +2643,38 @@ def bulk_write_to_db(
if not dags:
return
- log.info("Sync %s DAGs", len(dags))
- dag_by_ids = {dag.dag_id: dag for dag in dags}
-
- dag_ids = set(dag_by_ids)
- query = (
- select(DagModel)
- .options(joinedload(DagModel.tags, innerjoin=False))
- .where(DagModel.dag_id.in_(dag_ids))
- .options(joinedload(DagModel.schedule_dataset_references))
- .options(joinedload(DagModel.schedule_dataset_alias_references))
- .options(joinedload(DagModel.task_outlet_dataset_references))
- )
- query = with_row_locks(query, of=DagModel, session=session)
- orm_dags: list[DagModel] = session.scalars(query).unique().all()
- existing_dags: dict[str, DagModel] = {x.dag_id: x for x in orm_dags}
- missing_dag_ids = dag_ids.difference(existing_dags.keys())
-
- for missing_dag_id in missing_dag_ids:
- orm_dag = DagModel(dag_id=missing_dag_id)
- dag = dag_by_ids[missing_dag_id]
- if dag.is_paused_upon_creation is not None:
- orm_dag.is_paused = dag.is_paused_upon_creation
- orm_dag.tags = []
- log.info("Creating ORM DAG for %s", dag.dag_id)
- session.add(orm_dag)
- orm_dags.append(orm_dag)
-
- latest_runs: dict[str, DagRun] = {}
- num_active_runs: dict[str, int] = {}
- # Skip these queries entirely if no DAGs can be scheduled to save time.
- if any(dag.timetable.can_be_scheduled for dag in dags):
- # Get the latest automated dag run for each existing dag as a single query (avoid n+1 query)
- query = cls._get_latest_runs_stmt(dags=list(existing_dags.keys()))
- latest_runs = {run.dag_id: run for run in session.scalars(query)}
-
- # Get number of active dagruns for all dags we are processing as a single query.
- num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dags, session=session)
-
- filelocs = []
-
- for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id):
- dag = dag_by_ids[orm_dag.dag_id]
- filelocs.append(dag.fileloc)
- orm_dag.fileloc = dag.fileloc
- orm_dag.owners = dag.owner
- orm_dag.is_active = True
- orm_dag.has_import_errors = False
- orm_dag.last_parsed_time = timezone.utcnow()
- orm_dag.default_view = dag.default_view
- orm_dag._dag_display_property_value = dag._dag_display_property_value
- orm_dag.description = dag.description
- orm_dag.max_active_tasks = dag.max_active_tasks
- orm_dag.max_active_runs = dag.max_active_runs
- orm_dag.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
- orm_dag.has_task_concurrency_limits = any(
- t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None
- for t in dag.tasks
- )
- orm_dag.timetable_summary = dag.timetable.summary
- orm_dag.timetable_description = dag.timetable.description
- orm_dag.dataset_expression = dag.timetable.dataset_condition.as_expression()
-
- orm_dag.processor_subdir = processor_subdir
-
- last_automated_run: DagRun | None = latest_runs.get(dag.dag_id)
- if last_automated_run is None:
- last_automated_data_interval = None
- else:
- last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
- if num_active_runs.get(dag.dag_id, 0) >= orm_dag.max_active_runs:
- orm_dag.next_dagrun_create_after = None
- else:
- orm_dag.calculate_dagrun_date_fields(dag, last_automated_data_interval)
-
- dag_tags = set(dag.tags or {})
- orm_dag_tags = list(orm_dag.tags or [])
- for orm_tag in orm_dag_tags:
- if orm_tag.name not in dag_tags:
- session.delete(orm_tag)
- orm_dag.tags.remove(orm_tag)
- orm_tag_names = {t.name for t in orm_dag_tags}
- for dag_tag in dag_tags:
- if dag_tag not in orm_tag_names:
- dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
- orm_dag.tags.append(dag_tag_orm)
- session.add(dag_tag_orm)
-
- orm_dag_links = orm_dag.dag_owner_links or []
- for orm_dag_link in orm_dag_links:
- if orm_dag_link not in dag.owner_links:
- session.delete(orm_dag_link)
- for owner_name, owner_link in dag.owner_links.items():
- dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
- session.add(dag_owner_orm)
-
- DagCode.bulk_sync_to_db(filelocs, session=session)
-
- from airflow.datasets import Dataset
- from airflow.models.dataset import (
- DagScheduleDatasetAliasReference,
- DagScheduleDatasetReference,
- DatasetModel,
- TaskOutletDatasetReference,
+ from airflow.dag_processing.collection import (
+ DatasetModelOperation,
+ collect_orm_dags,
+ create_orm_dag,
+ update_orm_dags,
)
- dag_references: dict[str, set[tuple[Literal["dataset", "dataset-alias"], str]]] = defaultdict(set)
- outlet_references = defaultdict(set)
- # We can't use a set here as we want to preserve order
- outlet_dataset_models: dict[DatasetModel, None] = {}
- input_dataset_models: dict[DatasetModel, None] = {}
- outlet_dataset_alias_models: set[DatasetAliasModel] = set()
- input_dataset_alias_models: set[DatasetAliasModel] = set()
-
- # here we go through dags and tasks to check for dataset references
- # if there are now None and previously there were some, we delete them
- # if there are now *any*, we add them to the above data structures, and
- # later we'll persist them to the database.
- for dag in dags:
- curr_orm_dag = existing_dags.get(dag.dag_id)
- if not (dataset_condition := dag.timetable.dataset_condition):
- if curr_orm_dag:
- if curr_orm_dag.schedule_dataset_references:
- curr_orm_dag.schedule_dataset_references = []
- if curr_orm_dag.schedule_dataset_alias_references:
- curr_orm_dag.schedule_dataset_alias_references = []
- else:
- for _, dataset in dataset_condition.iter_datasets():
- dag_references[dag.dag_id].add(("dataset", dataset.uri))
- input_dataset_models[DatasetModel.from_public(dataset)] = None
-
- for dataset_alias in dataset_condition.iter_dataset_aliases():
- dag_references[dag.dag_id].add(("dataset-alias", dataset_alias.name))
- input_dataset_alias_models.add(DatasetAliasModel.from_public(dataset_alias))
-
- curr_outlet_references = curr_orm_dag and curr_orm_dag.task_outlet_dataset_references
- for task in dag.tasks:
- dataset_outlets: list[Dataset] = []
- dataset_alias_outlets: list[DatasetAlias] = []
- for outlet in task.outlets:
- if isinstance(outlet, Dataset):
- dataset_outlets.append(outlet)
- elif isinstance(outlet, DatasetAlias):
- dataset_alias_outlets.append(outlet)
-
- if not dataset_outlets:
- if curr_outlet_references:
- this_task_outlet_refs = [
- x
- for x in curr_outlet_references
- if x.dag_id == dag.dag_id and x.task_id == task.task_id
- ]
- for ref in this_task_outlet_refs:
- curr_outlet_references.remove(ref)
-
- for d in dataset_outlets:
- outlet_dataset_models[DatasetModel.from_public(d)] = None
- outlet_references[(task.dag_id, task.task_id)].add(d.uri)
-
- for d_a in dataset_alias_outlets:
- outlet_dataset_alias_models.add(DatasetAliasModel.from_public(d_a))
-
- all_dataset_models = outlet_dataset_models
- all_dataset_models.update(input_dataset_models)
-
- # store datasets
- stored_dataset_models: dict[str, DatasetModel] = {}
- new_dataset_models: list[DatasetModel] = []
- for dataset in all_dataset_models:
- stored_dataset_model = session.scalar(
- select(DatasetModel).where(DatasetModel.uri == dataset.uri).limit(1)
- )
- if stored_dataset_model:
- # Some datasets may have been previously unreferenced, and therefore orphaned by the
- # scheduler. But if we're here, then we have found that dataset again in our DAGs, which
- # means that it is no longer an orphan, so set is_orphaned to False.
- stored_dataset_model.is_orphaned = expression.false()
- stored_dataset_models[stored_dataset_model.uri] = stored_dataset_model
- else:
- new_dataset_models.append(dataset)
- dataset_manager.create_datasets(dataset_models=new_dataset_models, session=session)
- stored_dataset_models.update(
- {dataset_model.uri: dataset_model for dataset_model in new_dataset_models}
- )
-
- del new_dataset_models
- del all_dataset_models
-
- # store dataset aliases
- all_datasets_alias_models = input_dataset_alias_models | outlet_dataset_alias_models
- stored_dataset_alias_models: dict[str, DatasetAliasModel] = {}
- new_dataset_alias_models: set[DatasetAliasModel] = set()
- if all_datasets_alias_models:
- all_dataset_alias_names = {
- dataset_alias_model.name for dataset_alias_model in all_datasets_alias_models
- }
-
- stored_dataset_alias_models = {
- dsa_m.name: dsa_m
- for dsa_m in session.scalars(
- select(DatasetAliasModel).where(DatasetAliasModel.name.in_(all_dataset_alias_names))
- ).fetchall()
- }
-
- if stored_dataset_alias_models:
- new_dataset_alias_models = {
- dataset_alias_model
- for dataset_alias_model in all_datasets_alias_models
- if dataset_alias_model.name not in stored_dataset_alias_models.keys()
- }
- else:
- new_dataset_alias_models = all_datasets_alias_models
-
- session.add_all(new_dataset_alias_models)
- session.flush()
- stored_dataset_alias_models.update(
- {
- dataset_alias_model.name: dataset_alias_model
- for dataset_alias_model in new_dataset_alias_models
- }
+ log.info("Sync %s DAGs", len(dags))
+ dags_by_ids = {dag.dag_id: dag for dag in dags}
+ del dags
+
+ orm_dags = collect_orm_dags(dags_by_ids, session=session)
+ orm_dags.update(
+ (dag_id, create_orm_dag(dag, session=session))
+ for dag_id, dag in dags_by_ids.items()
+ if dag_id not in orm_dags
)
- del new_dataset_alias_models
- del all_datasets_alias_models
+ update_orm_dags(dags_by_ids, orm_dags, processor_subdir=processor_subdir, session=session)
+ DagCode.bulk_sync_to_db((dag.fileloc for dag in dags_by_ids.values()), session=session)
Review Comment:
Performance should be roughly the same since `dags_by_ids` is in-memory. I do plan to further refactor this in a later PR, as mentioned in https://github.com/apache/airflow/pull/42245#discussion_r1766097408. This will be done in a similar fashion as DatasetCollection.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766176759
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
+ for dag in dags:
+ for _, dataset in dag.timetable.dataset_condition.iter_datasets():
+ yield dataset
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, Dataset):
+ yield obj
+
+
+def _find_all_dataset_aliases(dags: Iterable[DAG]) -> Iterator[DatasetAlias]:
+ for dag in dags:
+ for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases():
+ yield alias
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, DatasetAlias):
+ yield obj
+
+
+class DatasetCollection(NamedTuple):
+ """Datasets collected from DAGs."""
+
+ schedule_dataset_references: dict[str, list[Dataset]]
+ schedule_dataset_alias_references: dict[str, list[DatasetAlias]]
+ outlet_references: dict[str, list[tuple[str, Dataset]]]
+ datasets: dict[str, Dataset]
+ dataset_aliases: dict[str, DatasetAlias]
+
+ @classmethod
+ def collect(cls, dags: dict[str, DAG]) -> Self:
+ coll = cls(
+ schedule_dataset_references={
+ dag_id: [dataset for _, dataset in dag.timetable.dataset_condition.iter_datasets()]
+ for dag_id, dag in dags.items()
+ },
+ schedule_dataset_alias_references={
+ dag_id: [alias for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases()]
+ for dag_id, dag in dags.items()
+ },
+ outlet_references={
+ dag_id: [
+ (task_id, outlet)
+ for task_id, task in dag.task_dict.items()
+ for outlet in task.outlets
+ if isinstance(outlet, Dataset)
+ ]
+ for dag_id, dag in dags.items()
+ },
+ datasets={dataset.uri: dataset for dataset in _find_all_datasets(dags.values())},
+ dataset_aliases={alias.name: alias for alias in _find_all_dataset_aliases(dags.values())},
+ )
+ return coll
+
+ def add_datasets(self, *, session: Session) -> dict[str, DatasetModel]:
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #42245:
URL: https://github.com/apache/airflow/pull/42245#issuecomment-2359780802
Here we gooooooo
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766097408
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
Review Comment:
Should we reuse https://github.com/apache/airflow/pull/42245/files#diff-7bf8b670d67b3f14e4d5ee8e8fcfad752193dd2e0ad51c428812008cbf4c82cbR96 here?
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
Review Comment:
I'm thinking of whether we should refactor this and _find_all_dataset_aliases with `partial`, but I guess this would make it way more complicated with little gain?
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
+ for dag in dags:
+ for _, dataset in dag.timetable.dataset_condition.iter_datasets():
+ yield dataset
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, Dataset):
+ yield obj
+
+
+def _find_all_dataset_aliases(dags: Iterable[DAG]) -> Iterator[DatasetAlias]:
+ for dag in dags:
+ for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases():
+ yield alias
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, DatasetAlias):
+ yield obj
+
+
+class DatasetCollection(NamedTuple):
+ """Datasets collected from DAGs."""
+
+ schedule_dataset_references: dict[str, list[Dataset]]
+ schedule_dataset_alias_references: dict[str, list[DatasetAlias]]
+ outlet_references: dict[str, list[tuple[str, Dataset]]]
+ datasets: dict[str, Dataset]
+ dataset_aliases: dict[str, DatasetAlias]
+
+ @classmethod
+ def collect(cls, dags: dict[str, DAG]) -> Self:
+ coll = cls(
+ schedule_dataset_references={
+ dag_id: [dataset for _, dataset in dag.timetable.dataset_condition.iter_datasets()]
+ for dag_id, dag in dags.items()
+ },
+ schedule_dataset_alias_references={
+ dag_id: [alias for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases()]
+ for dag_id, dag in dags.items()
+ },
+ outlet_references={
+ dag_id: [
+ (task_id, outlet)
+ for task_id, task in dag.task_dict.items()
+ for outlet in task.outlets
+ if isinstance(outlet, Dataset)
+ ]
+ for dag_id, dag in dags.items()
+ },
+ datasets={dataset.uri: dataset for dataset in _find_all_datasets(dags.values())},
+ dataset_aliases={alias.name: alias for alias in _find_all_dataset_aliases(dags.values())},
+ )
+ return coll
+
+ def add_datasets(self, *, session: Session) -> dict[str, DatasetModel]:
Review Comment:
It looks like add datasets to this `DatasetCollection` named tuple. not sure whether `add_datasets_to_db` makes more sense š¤
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
Review Comment:
Why do we want to clean up the tags here?
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
+ for dag in dags:
+ for _, dataset in dag.timetable.dataset_condition.iter_datasets():
+ yield dataset
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, Dataset):
+ yield obj
+
+
+def _find_all_dataset_aliases(dags: Iterable[DAG]) -> Iterator[DatasetAlias]:
+ for dag in dags:
+ for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases():
+ yield alias
+ for task in dag.task_dict.values():
+ for obj in itertools.chain(task.inlets, task.outlets):
+ if isinstance(obj, DatasetAlias):
+ yield obj
+
+
+class DatasetCollection(NamedTuple):
+ """Datasets collected from DAGs."""
+
+ schedule_dataset_references: dict[str, list[Dataset]]
+ schedule_dataset_alias_references: dict[str, list[DatasetAlias]]
+ outlet_references: dict[str, list[tuple[str, Dataset]]]
+ datasets: dict[str, Dataset]
+ dataset_aliases: dict[str, DatasetAlias]
+
+ @classmethod
+ def collect(cls, dags: dict[str, DAG]) -> Self:
+ coll = cls(
+ schedule_dataset_references={
+ dag_id: [dataset for _, dataset in dag.timetable.dataset_condition.iter_datasets()]
+ for dag_id, dag in dags.items()
+ },
+ schedule_dataset_alias_references={
+ dag_id: [alias for _, alias in dag.timetable.dataset_condition.iter_dataset_aliases()]
+ for dag_id, dag in dags.items()
+ },
+ outlet_references={
+ dag_id: [
+ (task_id, outlet)
+ for task_id, task in dag.task_dict.items()
+ for outlet in task.outlets
+ if isinstance(outlet, Dataset)
+ ]
+ for dag_id, dag in dags.items()
+ },
+ datasets={dataset.uri: dataset for dataset in _find_all_datasets(dags.values())},
+ dataset_aliases={alias.name: alias for alias in _find_all_dataset_aliases(dags.values())},
+ )
+ return coll
+
+ def add_datasets(self, *, session: Session) -> dict[str, DatasetModel]:
Review Comment:
or `insert_datasets` maybe?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766115523
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
+
+ return existing_orm_dags
+
+
+def create_orm_dag(dag: DAG, session: Session) -> DagModel:
+ orm_dag = DagModel(dag_id=dag.dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag.dag_id)
+ session.add(orm_dag)
+ return orm_dag
+
+
+def _get_latest_runs_stmt(dag_ids: Collection[str]) -> Select:
+ """Build a select statement for retrieve the last automated run for each dag."""
+ if len(dag_ids) == 1: # Index optimized fast path to avoid more complicated & slower groupby queryplan.
+ (dag_id,) = dag_ids
+ last_automated_runs_subq = (
+ select(func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id == dag_id,
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .scalar_subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == dag_id,
+ DagRun.execution_date == last_automated_runs_subq,
+ )
+ else:
+ last_automated_runs_subq = (
+ select(DagRun.dag_id, func.max(DagRun.execution_date).label("max_execution_date"))
+ .where(
+ DagRun.dag_id.in_(dag_ids),
+ DagRun.run_type.in_((DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED)),
+ )
+ .group_by(DagRun.dag_id)
+ .subquery()
+ )
+ query = select(DagRun).where(
+ DagRun.dag_id == last_automated_runs_subq.c.dag_id,
+ DagRun.execution_date == last_automated_runs_subq.c.max_execution_date,
+ )
+ return query.options(
+ load_only(
+ DagRun.dag_id,
+ DagRun.execution_date,
+ DagRun.data_interval_start,
+ DagRun.data_interval_end,
+ )
+ )
+
+
+class _RunInfo(NamedTuple):
+ latest_runs: dict[str, DagRun]
+ num_active_runs: dict[str, int]
+
+ @classmethod
+ def calculate(cls, dags: dict[str, DAG], *, session: Session) -> Self:
+ # Skip these queries entirely if no DAGs can be scheduled to save time.
+ if not any(dag.timetable.can_be_scheduled for dag in dags.values()):
+ return cls({}, {})
+ return cls(
+ {run.dag_id: run for run in session.scalars(_get_latest_runs_stmt(dag_ids=dags))},
+ DagRun.active_runs_of_dags(dag_ids=dags, session=session),
+ )
+
+
+def update_orm_dags(
+ source_dags: dict[str, DAG],
+ target_dags: dict[str, DagModel],
+ *,
+ processor_subdir: str | None = None,
+ session: Session,
+) -> None:
+ """
+ Apply DAG attributes to DagModel objects.
+
+ Objects in ``target_dags`` are modified in-place.
+ """
+ run_info = _RunInfo.calculate(source_dags, session=session)
+
+ for dag_id, dm in sorted(target_dags.items()):
+ dag = source_dags[dag_id]
+ dm.fileloc = dag.fileloc
+ dm.owners = dag.owner
+ dm.is_active = True
+ dm.has_import_errors = False
+ dm.last_parsed_time = utcnow()
+ dm.default_view = dag.default_view
+ dm._dag_display_property_value = dag._dag_display_property_value
+ dm.description = dag.description
+ dm.max_active_tasks = dag.max_active_tasks
+ dm.max_active_runs = dag.max_active_runs
+ dm.max_consecutive_failed_dag_runs = dag.max_consecutive_failed_dag_runs
+ dm.has_task_concurrency_limits = any(
+ t.max_active_tis_per_dag is not None or t.max_active_tis_per_dagrun is not None for t in dag.tasks
+ )
+ dm.timetable_summary = dag.timetable.summary
+ dm.timetable_description = dag.timetable.description
+ dm.dataset_expression = dag.timetable.dataset_condition.as_expression()
+ dm.processor_subdir = processor_subdir
+
+ last_automated_run: DagRun | None = run_info.latest_runs.get(dag.dag_id)
+ if last_automated_run is None:
+ last_automated_data_interval = None
+ else:
+ last_automated_data_interval = dag.get_run_data_interval(last_automated_run)
+ if run_info.num_active_runs.get(dag.dag_id, 0) >= dm.max_active_runs:
+ dm.next_dagrun_create_after = None
+ else:
+ dm.calculate_dagrun_date_fields(dag, last_automated_data_interval)
+
+ if not dag.timetable.dataset_condition:
+ dm.schedule_dataset_references = []
+ dm.schedule_dataset_alias_references = []
+ # FIXME: STORE NEW REFERENCES.
+
+ dag_tags = set(dag.tags or ())
+ for orm_tag in (dm_tags := list(dm.tags or [])):
+ if orm_tag.name not in dag_tags:
+ session.delete(orm_tag)
+ dm.tags.remove(orm_tag)
+ orm_tag_names = {t.name for t in dm_tags}
+ for dag_tag in dag_tags:
+ if dag_tag not in orm_tag_names:
+ dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id)
+ dm.tags.append(dag_tag_orm)
+ session.add(dag_tag_orm)
+
+ dm_links = dm.dag_owner_links or []
+ for dm_link in dm_links:
+ if dm_link not in dag.owner_links:
+ session.delete(dm_link)
+ for owner_name, owner_link in dag.owner_links.items():
+ dag_owner_orm = DagOwnerAttributes(dag_id=dag.dag_id, owner=owner_name, link=owner_link)
+ session.add(dag_owner_orm)
+
+
+def _find_all_datasets(dags: Iterable[DAG]) -> Iterator[Dataset]:
Review Comment:
Yep, totally agree with it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor bulk_save_to_db [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #42245:
URL: https://github.com/apache/airflow/pull/42245#discussion_r1766115512
##########
airflow/dag_processing/collection.py:
##########
@@ -0,0 +1,408 @@
+#
+# 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.
+
+"""
+Utility code that write DAGs in bulk into the database.
+
+This should generally only be called by internal methods such as
+``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
+
+:meta private:
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+from typing import TYPE_CHECKING, NamedTuple
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import joinedload, load_only
+from sqlalchemy.sql import expression
+
+from airflow.datasets import Dataset, DatasetAlias
+from airflow.datasets.manager import dataset_manager
+from airflow.models.dag import DAG, DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagrun import DagRun
+from airflow.models.dataset import (
+ DagScheduleDatasetAliasReference,
+ DagScheduleDatasetReference,
+ DatasetAliasModel,
+ DatasetModel,
+ TaskOutletDatasetReference,
+)
+from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.timezone import utcnow
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+ from collections.abc import Collection, Iterable, Iterator
+
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql import Select
+
+ from airflow.typing_compat import Self
+
+log = logging.getLogger(__name__)
+
+
+def collect_orm_dags(dags: dict[str, DAG], *, session: Session) -> dict[str, DagModel]:
+ """
+ Collect DagModel objects from DAG objects.
+
+ An existing DagModel is fetched if there's a matching ID in the database.
+ Otherwise, a new DagModel is created and added to the session.
+ """
+ stmt = (
+ select(DagModel)
+ .options(joinedload(DagModel.tags, innerjoin=False))
+ .where(DagModel.dag_id.in_(dags))
+ .options(joinedload(DagModel.schedule_dataset_references))
+ .options(joinedload(DagModel.schedule_dataset_alias_references))
+ .options(joinedload(DagModel.task_outlet_dataset_references))
+ )
+ stmt = with_row_locks(stmt, of=DagModel, session=session)
+ existing_orm_dags = {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
+
+ for dag_id, dag in dags.items():
+ if dag_id in existing_orm_dags:
+ continue
+ orm_dag = DagModel(dag_id=dag_id)
+ if dag.is_paused_upon_creation is not None:
+ orm_dag.is_paused = dag.is_paused_upon_creation
+ orm_dag.tags = []
+ log.info("Creating ORM DAG for %s", dag_id)
+ session.add(orm_dag)
+ existing_orm_dags[dag_id] = orm_dag
Review Comment:
Likely, but Iād prefer doing this in a later PRā¦ this is basically just copy-pasted and split into parts.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org