Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/services/ui/calendar.py: 15%
140 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 14:22 +0000
1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
17from __future__ import annotations
19import collections
20from collections.abc import Sequence
21from datetime import datetime
22from typing import Literal, cast
24import sqlalchemy as sa
25import structlog
26from sqlalchemy.engine import Row
27from sqlalchemy.orm import InstrumentedAttribute, Session
29from airflow._shared.timezones import timezone
30from airflow.api_fastapi.common.parameters import RangeFilter
31from airflow.api_fastapi.core_api.datamodels.ui.calendar import (
32 CalendarTimeRangeCollectionResponse,
33 CalendarTimeRangeResponse,
34)
35from airflow.models.dagrun import DagRun
36from airflow.serialization.definitions.dag import SerializedDAG
37from airflow.timetables._cron import CronMixin
38from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction
39from airflow.utils.sqlalchemy import get_dialect_name
41log = structlog.get_logger(logger_name=__name__)
44class CalendarService:
45 """Service class for calendar-related operations."""
47 MAX_PLANNED_RUNS: int = 2000
49 def get_calendar_data(
50 self,
51 dag_id: str,
52 session: Session,
53 dag: SerializedDAG,
54 logical_date: RangeFilter,
55 partition_date: RangeFilter,
56 granularity: Literal["hourly", "daily"] = "daily",
57 ) -> CalendarTimeRangeCollectionResponse:
58 """
59 Get calendar data for a Dag including historical and planned runs.
61 Args:
62 dag_id: The Dag ID
63 session: Database session
64 dag: The Dag object
65 logical_date: Date range filter for logical_date
66 partition_date: Date range filter for partition_date
67 granularity: Time granularity ("hourly" or "daily")
69 Returns:
70 List of calendar time range results
71 """
72 date_filter = partition_date if partition_date.is_active() else logical_date
73 historical_data, raw_dag_states = self._get_historical_dag_runs(
74 dag_id,
75 session,
76 date_filter,
77 granularity,
78 )
80 planned_data = self._get_planned_dag_runs(dag, raw_dag_states, date_filter, granularity)
82 all_data = historical_data + planned_data
83 return CalendarTimeRangeCollectionResponse(
84 total_entries=len(all_data),
85 dag_runs=all_data,
86 )
88 def _get_historical_dag_runs(
89 self,
90 dag_id: str,
91 session: Session,
92 date_filter: RangeFilter,
93 granularity: Literal["hourly", "daily"],
94 ) -> tuple[list[CalendarTimeRangeResponse], Sequence[Row]]:
95 """Get historical Dag runs from the database."""
96 dialect = get_dialect_name(session)
98 effective_date = sa.func.coalesce(DagRun.partition_date, DagRun.logical_date)
99 time_expression = self._get_time_truncation_expression(effective_date, granularity, dialect)
101 select_stmt = (
102 sa.select(
103 time_expression.label("datetime"),
104 DagRun.state,
105 sa.func.max(DagRun.data_interval_start).label("data_interval_start"),
106 sa.func.max(DagRun.data_interval_end).label("data_interval_end"),
107 sa.func.max(DagRun.run_after).label("run_after"),
108 sa.func.max(DagRun.partition_date).label("partition_date"),
109 sa.func.count("*").label("count"),
110 )
111 .where(DagRun.dag_id == dag_id)
112 .group_by(time_expression, DagRun.state)
113 .order_by(time_expression.asc())
114 )
116 select_stmt = date_filter.to_orm(select_stmt)
117 dag_states = session.execute(select_stmt).all()
119 calendar_results = [
120 CalendarTimeRangeResponse(
121 # ds.datetime in sqlite and mysql is a string, in postgresql it is a datetime
122 date=ds.datetime,
123 state=ds.state,
124 count=int(ds._mapping["count"]),
125 )
126 for ds in dag_states
127 ]
129 return calendar_results, dag_states
131 def _get_planned_dag_runs(
132 self,
133 dag: SerializedDAG,
134 raw_dag_states: Sequence[Row],
135 date_filter: RangeFilter,
136 granularity: Literal["hourly", "daily"],
137 ) -> list[CalendarTimeRangeResponse]:
138 """Get planned Dag runs based on the Dag's timetable."""
139 if not self._should_calculate_planned_runs(dag, raw_dag_states):
140 return []
142 last_state = raw_dag_states[-1]
144 if dag.timetable.partitioned:
145 last_run_after = timezone.coerce_datetime(last_state.run_after)
146 last_partition_date = timezone.coerce_datetime(last_state.partition_date)
147 if not last_run_after or not last_partition_date:
148 return []
149 year = last_partition_date.year
150 last_info = DagRunInfo(
151 run_after=last_run_after,
152 partition_date=last_partition_date,
153 partition_key=None,
154 data_interval=None,
155 )
156 else:
157 last_data_interval = self._get_last_data_interval(raw_dag_states)
158 if not last_data_interval:
159 return []
160 year = last_data_interval.end.year
161 if isinstance(dag.timetable, CronMixin):
162 return self._calculate_cron_planned_runs(
163 dag, last_data_interval, year, date_filter, granularity
164 )
165 last_info = DagRunInfo(
166 run_after=last_data_interval.end,
167 data_interval=last_data_interval,
168 partition_date=None,
169 partition_key=None,
170 )
172 restriction = TimeRestriction(
173 timezone.coerce_datetime(dag.start_date) if dag.start_date else None,
174 timezone.coerce_datetime(dag.end_date) if dag.end_date else None,
175 False,
176 )
178 return self._calculate_timetable_planned_runs(
179 dag, last_info, year, restriction, date_filter, granularity
180 )
182 def _should_calculate_planned_runs(self, dag: SerializedDAG, raw_dag_states: Sequence[Row]) -> bool:
183 """Check if we should calculate planned runs."""
184 if not raw_dag_states or not dag.timetable.periodic:
185 return False
186 last = raw_dag_states[-1]
187 return bool(last.data_interval_start and last.data_interval_end) or bool(last.partition_date)
189 def _get_last_data_interval(self, raw_dag_states: Sequence[Row]) -> DataInterval | None:
190 """Extract the last data interval from raw database results."""
191 if not raw_dag_states:
192 return None
194 last_state = raw_dag_states[-1]
195 if not (last_state.data_interval_start and last_state.data_interval_end):
196 return None
198 return DataInterval(
199 timezone.coerce_datetime(last_state.data_interval_start),
200 timezone.coerce_datetime(last_state.data_interval_end),
201 )
203 def _calculate_cron_planned_runs(
204 self,
205 dag: SerializedDAG,
206 last_data_interval: DataInterval,
207 year: int,
208 date_filter: RangeFilter,
209 granularity: Literal["hourly", "daily"],
210 ) -> list[CalendarTimeRangeResponse]:
211 """Calculate planned runs for cron-based timetables."""
212 dates: dict[datetime, int] = collections.Counter()
214 cron_timetable = cast("CronMixin", dag.timetable)
215 dt = last_data_interval.end
217 # Step with CronMixin._get_next so planned instants match the scheduler exactly,
218 # including its DST gap/fold handling. Cap the iteration like
219 # _calculate_timetable_planned_runs does; a high-frequency expression (e.g.
220 # "* * * * *", or a seconds-resolution cron) would otherwise take hundreds of
221 # thousands of steps before hitting the year boundary.
222 for _ in range(self.MAX_PLANNED_RUNS):
223 dt = cron_timetable._get_next(dt)
224 if dt.year != year:
225 break
226 if dag.end_date and dt > dag.end_date:
227 break
228 if not self._is_date_in_range(dt, date_filter):
229 continue
231 dates[self._truncate_datetime_for_granularity(dt, granularity)] += 1
233 return [
234 CalendarTimeRangeResponse(date=dt, state="planned", count=count) for dt, count in dates.items()
235 ]
237 def _calculate_timetable_planned_runs(
238 self,
239 dag: SerializedDAG,
240 last_info: DagRunInfo,
241 year: int,
242 restriction: TimeRestriction,
243 date_filter: RangeFilter,
244 granularity: Literal["hourly", "daily"],
245 ) -> list[CalendarTimeRangeResponse]:
246 """Calculate planned runs for generic timetables."""
247 dates: dict[datetime, int] = collections.Counter()
248 prev_run_after = last_info.run_after
249 total_planned = 0
251 while total_planned < self.MAX_PLANNED_RUNS:
252 curr_info = dag.timetable.next_dagrun_info_v2(
253 last_dagrun_info=last_info,
254 restriction=restriction,
255 )
257 if curr_info is None:
258 break
259 if curr_info.run_after <= prev_run_after:
260 break
262 effective_date = curr_info.partition_date or curr_info.logical_date
263 if not effective_date:
264 break
265 if effective_date.year != year:
266 break
268 if not self._is_date_in_range(effective_date, date_filter):
269 last_info = curr_info
270 prev_run_after = curr_info.run_after
271 total_planned += 1
272 continue
274 dt = self._truncate_datetime_for_granularity(effective_date, granularity)
275 dates[dt] += 1
276 last_info = curr_info
277 prev_run_after = curr_info.run_after
278 total_planned += 1
280 return [
281 CalendarTimeRangeResponse(date=dt, state="planned", count=count) for dt, count in dates.items()
282 ]
284 def _get_time_truncation_expression(
285 self,
286 column: InstrumentedAttribute[datetime | None] | sa.sql.elements.ColumnElement,
287 granularity: Literal["hourly", "daily"],
288 dialect: str | None,
289 ) -> sa.sql.elements.ColumnElement:
290 """
291 Get database-specific time truncation expression for SQLAlchemy.
293 We want to return always timestamp for both hourly and daily truncation.
294 Unfortunately different databases have different functions for truncating datetime, so we need to handle
295 them separately.
297 Args:
298 column: The datetime column to truncate
299 granularity: Either "hourly" or "daily"
300 dialect: Database dialect ("postgresql", "mysql", "sqlite")
302 Returns:
303 SQLAlchemy expression for time truncation
305 Raises:
306 ValueError: If the dialect is not supported
307 """
308 if granularity == "hourly":
309 if dialect == "postgresql":
310 expression = sa.func.date_trunc("hour", column)
311 elif dialect == "mysql":
312 expression = sa.func.date_format(column, "%Y-%m-%dT%H:00:00Z")
313 elif dialect == "sqlite":
314 expression = sa.func.strftime("%Y-%m-%dT%H:00:00Z", column)
315 else:
316 raise ValueError(f"Unsupported dialect: {dialect}")
317 else:
318 if dialect == "postgresql":
319 expression = sa.func.timezone("UTC", sa.func.cast(sa.func.cast(column, sa.Date), sa.DateTime))
320 elif dialect == "mysql":
321 expression = sa.func.date_format(column, "%Y-%m-%dT%00:00:00Z")
322 elif dialect == "sqlite":
323 expression = sa.func.strftime("%Y-%m-%dT00:00:00Z", column)
324 else:
325 raise ValueError(f"Unsupported dialect: {dialect}")
326 return expression
328 def _truncate_datetime_for_granularity(
329 self,
330 dt: datetime,
331 granularity: Literal["hourly", "daily"],
332 ) -> datetime:
333 """
334 Truncate datetime based on granularity for planned tasks grouping.
336 Args:
337 dt: The datetime to truncate
338 granularity: Either "hourly" or "daily"
340 Returns:
341 Truncated datetime
342 """
343 if granularity == "hourly":
344 return dt.replace(minute=0, second=0, microsecond=0)
345 return dt.replace(hour=0, minute=0, second=0, microsecond=0)
347 def _is_date_in_range(self, dt: datetime, date_filter: RangeFilter) -> bool:
348 """Check if a date is within the specified range filter."""
349 if not date_filter.value:
350 return True
352 if date_filter.value.lower_bound_gte and dt < date_filter.value.lower_bound_gte:
353 return False
354 if date_filter.value.lower_bound_gt and dt <= date_filter.value.lower_bound_gt:
355 return False
356 if date_filter.value.upper_bound_lte and dt > date_filter.value.upper_bound_lte:
357 return False
358 if date_filter.value.upper_bound_lt and dt >= date_filter.value.upper_bound_lt:
359 return False
361 return True