Lee-W commented on code in PR #63631:
URL: https://github.com/apache/airflow/pull/63631#discussion_r2939427062


##########
airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_calendar.py:
##########
@@ -181,3 +183,120 @@ def test_hourly_calendar(self, setup_dag_runs, 
test_client, query_params, result
         body = response.json()
 
         assert body == result
+
+
+class TestPartitionedCalendar:
+    """Calendar tests for partitioned DAGs (AIP-76) which use partition_date 
instead of logical_date."""

Review Comment:
   ```suggestion
       """Calendar tests for partitioned Dags (AIP-76) which use partition_date 
instead of logical_date."""
   ```



##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py:
##########
@@ -95,21 +93,25 @@ def _get_historical_dag_runs(
         """Get historical DAG runs from the database."""
         dialect = get_dialect_name(session)
 
-        time_expression = 
self._get_time_truncation_expression(DagRun.logical_date, granularity, dialect)
+        effective_date = sa.func.coalesce(DagRun.partition_date, 
DagRun.logical_date)
+        time_expression = self._get_time_truncation_expression(effective_date, 
granularity, dialect)
 
         select_stmt = (
             sa.select(
                 time_expression.label("datetime"),
                 DagRun.state,
                 
sa.func.max(DagRun.data_interval_start).label("data_interval_start"),
                 
sa.func.max(DagRun.data_interval_end).label("data_interval_end"),
+                sa.func.max(DagRun.run_after).label("run_after"),
+                sa.func.max(DagRun.partition_date).label("partition_date"),
                 sa.func.count("*").label("count"),
             )
             .where(DagRun.dag_id == dag_id)
             .group_by(time_expression, DagRun.state)
             .order_by(time_expression.asc())
         )
 
+        logical_date.attribute = effective_date

Review Comment:
   but it could be partition_date as well?



##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py:
##########
@@ -136,31 +138,53 @@ def _get_planned_dag_runs(
         if not self._should_calculate_planned_runs(dag, raw_dag_states):
             return []
 
-        last_data_interval = self._get_last_data_interval(raw_dag_states)
-        if not last_data_interval:
-            return []
+        last_state = raw_dag_states[-1]
+        is_partitioned = dag.timetable.partitioned
+
+        if is_partitioned:
+            last_run_after = timezone.coerce_datetime(last_state.run_after)

Review Comment:
   ```suggestion
           if dag.timetable.partitioned:
               last_run_after = timezone.coerce_datetime(last_state.run_after)
   ```



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to