From a7a1efd507c393e8baec0dfd6b4642d1b663af97 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Thu, 19 Dec 2024 03:58:37 +0530 Subject: [PATCH 1/9] Add dry run for backfill --- .../core_api/datamodels/backfills.py | 13 ++++ .../core_api/openapi/v1-generated.yaml | 32 +++++++- .../core_api/routes/public/backfills.py | 75 +++++++++++++++---- airflow/ui/openapi-gen/queries/queries.ts | 2 +- .../ui/openapi-gen/requests/schemas.gen.ts | 36 +++++++++ .../ui/openapi-gen/requests/services.gen.ts | 2 +- airflow/ui/openapi-gen/requests/types.gen.ts | 19 ++++- .../core_api/routes/public/test_backfills.py | 1 + 8 files changed, 161 insertions(+), 19 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/backfills.py b/airflow/api_fastapi/core_api/datamodels/backfills.py index be04063907a9d..c5f6bd86ea242 100644 --- a/airflow/api_fastapi/core_api/datamodels/backfills.py +++ b/airflow/api_fastapi/core_api/datamodels/backfills.py @@ -33,6 +33,7 @@ class BackfillPostBody(BaseModel): dag_run_conf: dict = {} reprocess_behavior: ReprocessBehavior = ReprocessBehavior.NONE max_active_runs: int = 10 + dry_run: bool = False class BackfillResponse(BaseModel): @@ -56,3 +57,15 @@ class BackfillCollectionResponse(BaseModel): backfills: list[BackfillResponse] total_entries: int + + +class BackfillRunInfo(BaseModel): + """Data model for run information during a backfill operation.""" + + logical_date: datetime + + +class BackfillDryRunResponse(BaseModel): + """Serializer for responses in dry-run mode for backfill operations.""" + + run_info_list: list[BackfillRunInfo] diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 6e8beb4fc6af1..0f6b9dca06e82 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1143,7 +1143,10 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/BackfillResponse' + anyOf: + - $ref: '#/components/schemas/BackfillResponse' + - $ref: '#/components/schemas/BackfillDryRunResponse' + title: Response Create Backfill '401': content: application/json: @@ -6210,6 +6213,18 @@ components: - total_entries title: BackfillCollectionResponse description: Backfill Collection serializer for responses. + BackfillDryRunResponse: + properties: + run_info_list: + items: + $ref: '#/components/schemas/BackfillRunInfo' + type: array + title: Run Info List + type: object + required: + - run_info_list + title: BackfillDryRunResponse + description: Serializer for responses in dry-run mode for backfill operations. BackfillPostBody: properties: dag_id: @@ -6238,6 +6253,10 @@ components: type: integer title: Max Active Runs default: 10 + dry_run: + type: boolean + title: Dry Run + default: false type: object required: - dag_id @@ -6301,6 +6320,17 @@ components: - updated_at title: BackfillResponse description: Base serializer for Backfill. + BackfillRunInfo: + properties: + logical_date: + type: string + format: date-time + title: Logical Date + type: object + required: + - logical_date + title: BackfillRunInfo + description: Data model for run information during a backfill operation. BaseInfoResponse: properties: status: diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index 61d25597cdd15..4d8a0b1d33daf 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -19,7 +19,7 @@ from typing import Annotated from fastapi import Depends, HTTPException, status -from sqlalchemy import select, update +from sqlalchemy import desc, select, update from airflow.api_fastapi.common.db.common import ( AsyncSessionDep, @@ -30,8 +30,10 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.datamodels.backfills import ( BackfillCollectionResponse, + BackfillDryRunResponse, BackfillPostBody, BackfillResponse, + BackfillRunInfo, ) from airflow.api_fastapi.core_api.openapi.exceptions import ( create_openapi_http_exception_doc, @@ -41,9 +43,14 @@ AlreadyRunningBackfill, Backfill, BackfillDagRun, + BackfillDagRunExceptionReason, + ReprocessBehavior, _create_backfill, + _get_info_list, ) +from airflow.models.serialized_dag import SerializedDagModel from airflow.utils import timezone +from airflow.utils.sqlalchemy import nulls_first from airflow.utils.state import DagRunState backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") @@ -187,22 +194,62 @@ def cancel_backfill(backfill_id, session: SessionDep) -> BackfillResponse: ) def create_backfill( backfill_request: BackfillPostBody, -) -> BackfillResponse: + session: SessionDep, +) -> BackfillResponse | BackfillDryRunResponse: from_date = timezone.coerce_datetime(backfill_request.from_date) to_date = timezone.coerce_datetime(backfill_request.to_date) - try: - backfill_obj = _create_backfill( - dag_id=backfill_request.dag_id, + if not backfill_request.dry_run: + try: + backfill_obj = _create_backfill( + dag_id=backfill_request.dag_id, + from_date=from_date, + to_date=to_date, + max_active_runs=backfill_request.max_active_runs, + reverse=backfill_request.run_backwards, + dag_run_conf=backfill_request.dag_run_conf, + reprocess_behavior=backfill_request.reprocess_behavior, + ) + return BackfillResponse.model_validate(backfill_obj) + except AlreadyRunningBackfill: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"There is already a running backfill for dag {backfill_request.dag_id}", + ) + else: + serdag = session.scalar(SerializedDagModel.latest_item_select_object(backfill_request.dag_id)) + if not serdag: + raise HTTPException(status_code=404, detail=f"Could not find dag {backfill_request.dag_id}") + + info_list = _get_info_list( + dag=serdag.dag, from_date=from_date, to_date=to_date, - max_active_runs=backfill_request.max_active_runs, reverse=backfill_request.run_backwards, - dag_run_conf=backfill_request.dag_run_conf, - reprocess_behavior=backfill_request.reprocess_behavior, - ) - return BackfillResponse.model_validate(backfill_obj) - except AlreadyRunningBackfill: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"There is already a running backfill for dag {backfill_request.dag_id}", ) + backfill_response_item = [] + print(info_list) + for info in info_list: + print(info.logical_date) + dr = session.scalar( + select(DagRun) + .where(DagRun.logical_date == info.logical_date) + .order_by(nulls_first(desc(DagRun.start_date), session)) + .limit(1) + ) + + if dr: + non_create_reason = None + if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): + non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT + elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: + if dr.state != DagRunState.FAILED: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + if not non_create_reason: + backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + + else: + backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + + return BackfillDryRunResponse(run_info_list=backfill_response_item) diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 216c165ae330f..a41499117674a 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -2779,7 +2779,7 @@ export const useAssetServiceCreateAssetEvent = < * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns BackfillResponse Successful Response + * @returns unknown Successful Response * @throws ApiError */ export const useBackfillServiceCreateBackfill = < diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index cd9c2e5897dbf..60b735e335142 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -350,6 +350,23 @@ export const $BackfillCollectionResponse = { description: "Backfill Collection serializer for responses.", } as const; +export const $BackfillDryRunResponse = { + properties: { + run_info_list: { + items: { + $ref: "#/components/schemas/BackfillRunInfo", + }, + type: "array", + title: "Run Info List", + }, + }, + type: "object", + required: ["run_info_list"], + title: "BackfillDryRunResponse", + description: + "Serializer for responses in dry-run mode for backfill operations.", +} as const; + export const $BackfillPostBody = { properties: { dag_id: { @@ -385,6 +402,11 @@ export const $BackfillPostBody = { title: "Max Active Runs", default: 10, }, + dry_run: { + type: "boolean", + title: "Dry Run", + default: false, + }, }, type: "object", required: ["dag_id", "from_date", "to_date"], @@ -468,6 +490,20 @@ export const $BackfillResponse = { description: "Base serializer for Backfill.", } as const; +export const $BackfillRunInfo = { + properties: { + logical_date: { + type: "string", + format: "date-time", + title: "Logical Date", + }, + }, + type: "object", + required: ["logical_date"], + title: "BackfillRunInfo", + description: "Data model for run information during a backfill operation.", +} as const; + export const $BaseInfoResponse = { properties: { status: { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index c048ddb023053..b967f19258e32 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -837,7 +837,7 @@ export class BackfillService { * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns BackfillResponse Successful Response + * @returns unknown Successful Response * @throws ApiError */ public static createBackfill( diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 620ac86b815a7..5b652e9a46b6a 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -97,6 +97,13 @@ export type BackfillCollectionResponse = { total_entries: number; }; +/** + * Serializer for responses in dry-run mode for backfill operations. + */ +export type BackfillDryRunResponse = { + run_info_list: Array; +}; + /** * Object used for create backfill request. */ @@ -110,6 +117,7 @@ export type BackfillPostBody = { }; reprocess_behavior?: ReprocessBehavior; max_active_runs?: number; + dry_run?: boolean; }; /** @@ -131,6 +139,13 @@ export type BackfillResponse = { updated_at: string; }; +/** + * Data model for run information during a backfill operation. + */ +export type BackfillRunInfo = { + logical_date: string; +}; + /** * Base info serializer for responses. */ @@ -1497,7 +1512,7 @@ export type CreateBackfillData = { requestBody: BackfillPostBody; }; -export type CreateBackfillResponse = BackfillResponse; +export type CreateBackfillResponse = BackfillResponse | BackfillDryRunResponse; export type GetBackfillData = { backfillId: string; @@ -2629,7 +2644,7 @@ export type $OpenApiTs = { /** * Successful Response */ - 200: BackfillResponse; + 200: BackfillResponse | BackfillDryRunResponse; /** * Unauthorized */ diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index 1c64b10848f4b..a711c13ff7010 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -192,6 +192,7 @@ def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_cl "max_active_runs": max_active_runs, "run_backwards": False, "dag_run_conf": {"param1": "val1", "param2": True}, + "dry_run": False, } if repro_act is not None: data["reprocess_behavior"] = repro_act From d28893413149085cc8a9278378619c39fda5efc0 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Thu, 19 Dec 2024 03:58:37 +0530 Subject: [PATCH 2/9] Add dry run for backfill --- .../core_api/datamodels/backfills.py | 13 ++++ .../core_api/openapi/v1-generated.yaml | 32 +++++++- .../core_api/routes/public/backfills.py | 75 +++++++++++++++---- airflow/ui/openapi-gen/queries/queries.ts | 2 +- .../ui/openapi-gen/requests/schemas.gen.ts | 36 +++++++++ .../ui/openapi-gen/requests/services.gen.ts | 2 +- airflow/ui/openapi-gen/requests/types.gen.ts | 19 ++++- .../core_api/routes/public/test_backfills.py | 1 + 8 files changed, 161 insertions(+), 19 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/backfills.py b/airflow/api_fastapi/core_api/datamodels/backfills.py index be04063907a9d..c5f6bd86ea242 100644 --- a/airflow/api_fastapi/core_api/datamodels/backfills.py +++ b/airflow/api_fastapi/core_api/datamodels/backfills.py @@ -33,6 +33,7 @@ class BackfillPostBody(BaseModel): dag_run_conf: dict = {} reprocess_behavior: ReprocessBehavior = ReprocessBehavior.NONE max_active_runs: int = 10 + dry_run: bool = False class BackfillResponse(BaseModel): @@ -56,3 +57,15 @@ class BackfillCollectionResponse(BaseModel): backfills: list[BackfillResponse] total_entries: int + + +class BackfillRunInfo(BaseModel): + """Data model for run information during a backfill operation.""" + + logical_date: datetime + + +class BackfillDryRunResponse(BaseModel): + """Serializer for responses in dry-run mode for backfill operations.""" + + run_info_list: list[BackfillRunInfo] diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 980b045e34d06..bdc6ae842bc60 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1143,7 +1143,10 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/BackfillResponse' + anyOf: + - $ref: '#/components/schemas/BackfillResponse' + - $ref: '#/components/schemas/BackfillDryRunResponse' + title: Response Create Backfill '401': content: application/json: @@ -6210,6 +6213,18 @@ components: - total_entries title: BackfillCollectionResponse description: Backfill Collection serializer for responses. + BackfillDryRunResponse: + properties: + run_info_list: + items: + $ref: '#/components/schemas/BackfillRunInfo' + type: array + title: Run Info List + type: object + required: + - run_info_list + title: BackfillDryRunResponse + description: Serializer for responses in dry-run mode for backfill operations. BackfillPostBody: properties: dag_id: @@ -6238,6 +6253,10 @@ components: type: integer title: Max Active Runs default: 10 + dry_run: + type: boolean + title: Dry Run + default: false type: object required: - dag_id @@ -6301,6 +6320,17 @@ components: - updated_at title: BackfillResponse description: Base serializer for Backfill. + BackfillRunInfo: + properties: + logical_date: + type: string + format: date-time + title: Logical Date + type: object + required: + - logical_date + title: BackfillRunInfo + description: Data model for run information during a backfill operation. BaseInfoResponse: properties: status: diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index 61d25597cdd15..4d8a0b1d33daf 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -19,7 +19,7 @@ from typing import Annotated from fastapi import Depends, HTTPException, status -from sqlalchemy import select, update +from sqlalchemy import desc, select, update from airflow.api_fastapi.common.db.common import ( AsyncSessionDep, @@ -30,8 +30,10 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.datamodels.backfills import ( BackfillCollectionResponse, + BackfillDryRunResponse, BackfillPostBody, BackfillResponse, + BackfillRunInfo, ) from airflow.api_fastapi.core_api.openapi.exceptions import ( create_openapi_http_exception_doc, @@ -41,9 +43,14 @@ AlreadyRunningBackfill, Backfill, BackfillDagRun, + BackfillDagRunExceptionReason, + ReprocessBehavior, _create_backfill, + _get_info_list, ) +from airflow.models.serialized_dag import SerializedDagModel from airflow.utils import timezone +from airflow.utils.sqlalchemy import nulls_first from airflow.utils.state import DagRunState backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") @@ -187,22 +194,62 @@ def cancel_backfill(backfill_id, session: SessionDep) -> BackfillResponse: ) def create_backfill( backfill_request: BackfillPostBody, -) -> BackfillResponse: + session: SessionDep, +) -> BackfillResponse | BackfillDryRunResponse: from_date = timezone.coerce_datetime(backfill_request.from_date) to_date = timezone.coerce_datetime(backfill_request.to_date) - try: - backfill_obj = _create_backfill( - dag_id=backfill_request.dag_id, + if not backfill_request.dry_run: + try: + backfill_obj = _create_backfill( + dag_id=backfill_request.dag_id, + from_date=from_date, + to_date=to_date, + max_active_runs=backfill_request.max_active_runs, + reverse=backfill_request.run_backwards, + dag_run_conf=backfill_request.dag_run_conf, + reprocess_behavior=backfill_request.reprocess_behavior, + ) + return BackfillResponse.model_validate(backfill_obj) + except AlreadyRunningBackfill: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"There is already a running backfill for dag {backfill_request.dag_id}", + ) + else: + serdag = session.scalar(SerializedDagModel.latest_item_select_object(backfill_request.dag_id)) + if not serdag: + raise HTTPException(status_code=404, detail=f"Could not find dag {backfill_request.dag_id}") + + info_list = _get_info_list( + dag=serdag.dag, from_date=from_date, to_date=to_date, - max_active_runs=backfill_request.max_active_runs, reverse=backfill_request.run_backwards, - dag_run_conf=backfill_request.dag_run_conf, - reprocess_behavior=backfill_request.reprocess_behavior, - ) - return BackfillResponse.model_validate(backfill_obj) - except AlreadyRunningBackfill: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"There is already a running backfill for dag {backfill_request.dag_id}", ) + backfill_response_item = [] + print(info_list) + for info in info_list: + print(info.logical_date) + dr = session.scalar( + select(DagRun) + .where(DagRun.logical_date == info.logical_date) + .order_by(nulls_first(desc(DagRun.start_date), session)) + .limit(1) + ) + + if dr: + non_create_reason = None + if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): + non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT + elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: + if dr.state != DagRunState.FAILED: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + if not non_create_reason: + backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + + else: + backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + + return BackfillDryRunResponse(run_info_list=backfill_response_item) diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 216c165ae330f..a41499117674a 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -2779,7 +2779,7 @@ export const useAssetServiceCreateAssetEvent = < * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns BackfillResponse Successful Response + * @returns unknown Successful Response * @throws ApiError */ export const useBackfillServiceCreateBackfill = < diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index defc76ec7345f..f842c5c7c8352 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -350,6 +350,23 @@ export const $BackfillCollectionResponse = { description: "Backfill Collection serializer for responses.", } as const; +export const $BackfillDryRunResponse = { + properties: { + run_info_list: { + items: { + $ref: "#/components/schemas/BackfillRunInfo", + }, + type: "array", + title: "Run Info List", + }, + }, + type: "object", + required: ["run_info_list"], + title: "BackfillDryRunResponse", + description: + "Serializer for responses in dry-run mode for backfill operations.", +} as const; + export const $BackfillPostBody = { properties: { dag_id: { @@ -385,6 +402,11 @@ export const $BackfillPostBody = { title: "Max Active Runs", default: 10, }, + dry_run: { + type: "boolean", + title: "Dry Run", + default: false, + }, }, type: "object", required: ["dag_id", "from_date", "to_date"], @@ -468,6 +490,20 @@ export const $BackfillResponse = { description: "Base serializer for Backfill.", } as const; +export const $BackfillRunInfo = { + properties: { + logical_date: { + type: "string", + format: "date-time", + title: "Logical Date", + }, + }, + type: "object", + required: ["logical_date"], + title: "BackfillRunInfo", + description: "Data model for run information during a backfill operation.", +} as const; + export const $BaseInfoResponse = { properties: { status: { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index c048ddb023053..b967f19258e32 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -837,7 +837,7 @@ export class BackfillService { * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns BackfillResponse Successful Response + * @returns unknown Successful Response * @throws ApiError */ public static createBackfill( diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 9aeda695c3554..b1f17c39c86e7 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -97,6 +97,13 @@ export type BackfillCollectionResponse = { total_entries: number; }; +/** + * Serializer for responses in dry-run mode for backfill operations. + */ +export type BackfillDryRunResponse = { + run_info_list: Array; +}; + /** * Object used for create backfill request. */ @@ -110,6 +117,7 @@ export type BackfillPostBody = { }; reprocess_behavior?: ReprocessBehavior; max_active_runs?: number; + dry_run?: boolean; }; /** @@ -131,6 +139,13 @@ export type BackfillResponse = { updated_at: string; }; +/** + * Data model for run information during a backfill operation. + */ +export type BackfillRunInfo = { + logical_date: string; +}; + /** * Base info serializer for responses. */ @@ -1498,7 +1513,7 @@ export type CreateBackfillData = { requestBody: BackfillPostBody; }; -export type CreateBackfillResponse = BackfillResponse; +export type CreateBackfillResponse = BackfillResponse | BackfillDryRunResponse; export type GetBackfillData = { backfillId: string; @@ -2630,7 +2645,7 @@ export type $OpenApiTs = { /** * Successful Response */ - 200: BackfillResponse; + 200: BackfillResponse | BackfillDryRunResponse; /** * Unauthorized */ diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index 1c64b10848f4b..a711c13ff7010 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -192,6 +192,7 @@ def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_cl "max_active_runs": max_active_runs, "run_backwards": False, "dag_run_conf": {"param1": "val1", "param2": True}, + "dry_run": False, } if repro_act is not None: data["reprocess_behavior"] = repro_act From 5bd604e69dd741b1d1afdafdf5339d27b2e57828 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Thu, 19 Dec 2024 18:25:58 +0530 Subject: [PATCH 3/9] Add tests for backfill dry runs --- .../core_api/routes/public/backfills.py | 3 + .../core_api/routes/public/test_backfills.py | 87 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index 4d8a0b1d33daf..3bba1f115b760 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -243,6 +243,9 @@ def create_backfill( non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + elif backfill_request.reprocess_behavior is ReprocessBehavior.COMPLETED: + if dr.state == DagRunState.FAILED: + non_create_reason = BackfillDagRunExceptionReason.UNKNOWN elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: if dr.state != DagRunState.FAILED: non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index a711c13ff7010..f653489efedc3 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -216,6 +216,93 @@ def test_create_backfill(self, repro_act, repro_exp, session, dag_maker, test_cl } +class TestCreateBackfillDryRun(TestBackfillEndpoint): + @pytest.mark.parametrize( + "reprocess_behavior, expected_dates", + [ + ( + "none", + [ + {"logical_date": "2024-01-01T00:00:00Z"}, + {"logical_date": "2024-01-04T00:00:00Z"}, + {"logical_date": "2024-01-05T00:00:00Z"}, + ], + ), + ( + "failed", + [ + {"logical_date": "2024-01-01T00:00:00Z"}, + {"logical_date": "2024-01-03T00:00:00Z"}, # Reprocess failed + {"logical_date": "2024-01-04T00:00:00Z"}, + {"logical_date": "2024-01-05T00:00:00Z"}, + ], + ), + ( + "completed", + [ + {"logical_date": "2024-01-01T00:00:00Z"}, + {"logical_date": "2024-01-02T00:00:00Z"}, # Reprocess completed + {"logical_date": "2024-01-04T00:00:00Z"}, + {"logical_date": "2024-01-05T00:00:00Z"}, + ], + ), + ], + ) + def test_create_backfill_dry_run( + self, session, dag_maker, test_client, reprocess_behavior, expected_dates + ): + with dag_maker( + session=session, + dag_id="TEST_DAG_2", + schedule="0 0 * * *", + start_date=pendulum.parse("2024-01-01"), + ) as dag: + EmptyOperator(task_id="mytask") + + session.commit() + + existing_dagruns = [ + {"logical_date": pendulum.parse("2024-01-02"), "state": DagRunState.SUCCESS}, # Completed dag run + {"logical_date": pendulum.parse("2024-01-03"), "state": DagRunState.FAILED}, # Failed dag run + ] + for dagrun in existing_dagruns: + session.add( + DagRun( + dag_id=dag.dag_id, + run_id=f"manual__{dagrun['logical_date'].isoformat()}", + logical_date=dagrun["logical_date"], + state=dagrun["state"], + run_type="scheduled", + ) + ) + session.commit() + + from_date = pendulum.parse("2024-01-01") + from_date_iso = to_iso(from_date) + to_date = pendulum.parse("2024-01-05") + to_date_iso = to_iso(to_date) + + data = { + "dag_id": dag.dag_id, + "from_date": from_date_iso, + "to_date": to_date_iso, + "max_active_runs": 5, + "run_backwards": False, + "dag_run_conf": {"param1": "val1", "param2": True}, + "dry_run": True, + "reprocess_behavior": reprocess_behavior, + } + + response = test_client.post( + url="/public/backfills", + json=data, + ) + + assert response.status_code == 200 + response_json = response.json() + assert response_json["run_info_list"] == expected_dates + + class TestCancelBackfill(TestBackfillEndpoint): def test_cancel_backfill(self, session, test_client): (dag,) = self._create_dag_models() From 23b7b9b7e48f3f939bff31adaf90778b1c54e1d4 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Thu, 19 Dec 2024 19:14:28 +0530 Subject: [PATCH 4/9] Fix tests --- airflow/api_fastapi/core_api/routes/public/backfills.py | 3 --- tests/api_fastapi/core_api/routes/public/test_backfills.py | 3 ++- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index 3bba1f115b760..4d8a0b1d33daf 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -243,9 +243,6 @@ def create_backfill( non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - elif backfill_request.reprocess_behavior is ReprocessBehavior.COMPLETED: - if dr.state == DagRunState.FAILED: - non_create_reason = BackfillDagRunExceptionReason.UNKNOWN elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: if dr.state != DagRunState.FAILED: non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index f653489efedc3..c24b2b6cc8a8e 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -241,7 +241,8 @@ class TestCreateBackfillDryRun(TestBackfillEndpoint): "completed", [ {"logical_date": "2024-01-01T00:00:00Z"}, - {"logical_date": "2024-01-02T00:00:00Z"}, # Reprocess completed + {"logical_date": "2024-01-02T00:00:00Z"}, # Reprocess all + {"logical_date": "2024-01-03T00:00:00Z"}, {"logical_date": "2024-01-04T00:00:00Z"}, {"logical_date": "2024-01-05T00:00:00Z"}, ], From ab17ebf5cb5477eb4804add11b857df9e4c81288 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Mon, 23 Dec 2024 10:35:06 +0530 Subject: [PATCH 5/9] create a separate endpoint --- .../core_api/openapi/v1-generated.yaml | 54 ++++++++- .../core_api/routes/public/backfills.py | 112 ++++++++++-------- airflow/ui/openapi-gen/queries/common.ts | 3 + airflow/ui/openapi-gen/queries/queries.ts | 40 ++++++- .../ui/openapi-gen/requests/services.gen.ts | 29 ++++- airflow/ui/openapi-gen/requests/types.gen.ts | 41 ++++++- .../core_api/routes/public/test_backfills.py | 2 +- 7 files changed, 222 insertions(+), 59 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index bdc6ae842bc60..35e681381cb5f 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1143,10 +1143,7 @@ paths: content: application/json: schema: - anyOf: - - $ref: '#/components/schemas/BackfillResponse' - - $ref: '#/components/schemas/BackfillDryRunResponse' - title: Response Create Backfill + $ref: '#/components/schemas/BackfillResponse' '401': content: application/json: @@ -1368,6 +1365,55 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /public/backfills/dry_run: + post: + tags: + - Backfill + summary: Create Backfill Dry Run + operationId: create_backfill_dry_run + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/BackfillPostBody' + required: true + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/BackfillDryRunResponse' + '401': + description: Unauthorized + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '403': + description: Forbidden + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '404': + description: Not Found + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '409': + description: Conflict + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /public/connections/{connection_id}: delete: tags: diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index 4d8a0b1d33daf..be4c20957ed03 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -194,62 +194,74 @@ def cancel_backfill(backfill_id, session: SessionDep) -> BackfillResponse: ) def create_backfill( backfill_request: BackfillPostBody, - session: SessionDep, -) -> BackfillResponse | BackfillDryRunResponse: +) -> BackfillResponse: from_date = timezone.coerce_datetime(backfill_request.from_date) to_date = timezone.coerce_datetime(backfill_request.to_date) - if not backfill_request.dry_run: - try: - backfill_obj = _create_backfill( - dag_id=backfill_request.dag_id, - from_date=from_date, - to_date=to_date, - max_active_runs=backfill_request.max_active_runs, - reverse=backfill_request.run_backwards, - dag_run_conf=backfill_request.dag_run_conf, - reprocess_behavior=backfill_request.reprocess_behavior, - ) - return BackfillResponse.model_validate(backfill_obj) - except AlreadyRunningBackfill: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"There is already a running backfill for dag {backfill_request.dag_id}", - ) - else: - serdag = session.scalar(SerializedDagModel.latest_item_select_object(backfill_request.dag_id)) - if not serdag: - raise HTTPException(status_code=404, detail=f"Could not find dag {backfill_request.dag_id}") - - info_list = _get_info_list( - dag=serdag.dag, + try: + backfill_obj = _create_backfill( + dag_id=backfill_request.dag_id, from_date=from_date, to_date=to_date, + max_active_runs=backfill_request.max_active_runs, reverse=backfill_request.run_backwards, + dag_run_conf=backfill_request.dag_run_conf, + reprocess_behavior=backfill_request.reprocess_behavior, ) - backfill_response_item = [] - print(info_list) - for info in info_list: - print(info.logical_date) - dr = session.scalar( - select(DagRun) - .where(DagRun.logical_date == info.logical_date) - .order_by(nulls_first(desc(DagRun.start_date), session)) - .limit(1) - ) - - if dr: - non_create_reason = None - if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): - non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT - elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: - non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: - if dr.state != DagRunState.FAILED: - non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - if not non_create_reason: - backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + return BackfillResponse.model_validate(backfill_obj) + except AlreadyRunningBackfill: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"There is already a running backfill for dag {backfill_request.dag_id}", + ) + - else: +@backfills_router.post( + path="/dry_run", + responses=create_openapi_http_exception_doc( + [ + status.HTTP_404_NOT_FOUND, + status.HTTP_409_CONFLICT, + ] + ), +) +def create_backfill_dry_run( + backfill_request: BackfillPostBody, + session: SessionDep, +) -> BackfillDryRunResponse: + from_date = timezone.coerce_datetime(backfill_request.from_date) + to_date = timezone.coerce_datetime(backfill_request.to_date) + serdag = session.scalar(SerializedDagModel.latest_item_select_object(backfill_request.dag_id)) + if not serdag: + raise HTTPException(status_code=404, detail=f"Could not find dag {backfill_request.dag_id}") + + info_list = _get_info_list( + dag=serdag.dag, + from_date=from_date, + to_date=to_date, + reverse=backfill_request.run_backwards, + ) + backfill_response_item = [] + for info in info_list: + dr = session.scalar( + select(DagRun) + .where(DagRun.logical_date == info.logical_date) + .order_by(nulls_first(desc(DagRun.start_date), session)) + .limit(1) + ) + + if dr: + non_create_reason = None + if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): + non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT + elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: + if dr.state != DagRunState.FAILED: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + if not non_create_reason: backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) - return BackfillDryRunResponse(run_info_list=backfill_response_item) + else: + backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + + return BackfillDryRunResponse(run_info_list=backfill_response_item) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index ea18796932159..f06c2c91a14bb 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -1686,6 +1686,9 @@ export type AssetServiceCreateAssetEventMutationResult = Awaited< export type BackfillServiceCreateBackfillMutationResult = Awaited< ReturnType >; +export type BackfillServiceCreateBackfillDryRunMutationResult = Awaited< + ReturnType +>; export type ConnectionServicePostConnectionMutationResult = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index a41499117674a..b0bef485a23ee 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -2779,7 +2779,7 @@ export const useAssetServiceCreateAssetEvent = < * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns unknown Successful Response + * @returns BackfillResponse Successful Response * @throws ApiError */ export const useBackfillServiceCreateBackfill = < @@ -2813,6 +2813,44 @@ export const useBackfillServiceCreateBackfill = < }) as unknown as Promise, ...options, }); +/** + * Create Backfill Dry Run + * @param data The data for the request. + * @param data.requestBody + * @returns BackfillDryRunResponse Successful Response + * @throws ApiError + */ +export const useBackfillServiceCreateBackfillDryRun = < + TData = Common.BackfillServiceCreateBackfillDryRunMutationResult, + TError = unknown, + TContext = unknown, +>( + options?: Omit< + UseMutationOptions< + TData, + TError, + { + requestBody: BackfillPostBody; + }, + TContext + >, + "mutationFn" + >, +) => + useMutation< + TData, + TError, + { + requestBody: BackfillPostBody; + }, + TContext + >({ + mutationFn: ({ requestBody }) => + BackfillService.createBackfillDryRun({ + requestBody, + }) as unknown as Promise, + ...options, + }); /** * Post Connection * Create connection entry. diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index b967f19258e32..99967dfc0903e 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -54,6 +54,8 @@ import type { UnpauseBackfillResponse, CancelBackfillData, CancelBackfillResponse, + CreateBackfillDryRunData, + CreateBackfillDryRunResponse, DeleteConnectionData, DeleteConnectionResponse, GetConnectionData, @@ -837,7 +839,7 @@ export class BackfillService { * Create Backfill * @param data The data for the request. * @param data.requestBody - * @returns unknown Successful Response + * @returns BackfillResponse Successful Response * @throws ApiError */ public static createBackfill( @@ -960,6 +962,31 @@ export class BackfillService { }, }); } + + /** + * Create Backfill Dry Run + * @param data The data for the request. + * @param data.requestBody + * @returns BackfillDryRunResponse Successful Response + * @throws ApiError + */ + public static createBackfillDryRun( + data: CreateBackfillDryRunData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "POST", + url: "/public/backfills/dry_run", + body: data.requestBody, + mediaType: "application/json", + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", + 409: "Conflict", + 422: "Validation Error", + }, + }); + } } export class ConnectionService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index b1f17c39c86e7..f92d2fd23a56f 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1513,7 +1513,7 @@ export type CreateBackfillData = { requestBody: BackfillPostBody; }; -export type CreateBackfillResponse = BackfillResponse | BackfillDryRunResponse; +export type CreateBackfillResponse = BackfillResponse; export type GetBackfillData = { backfillId: string; @@ -1539,6 +1539,12 @@ export type CancelBackfillData = { export type CancelBackfillResponse = BackfillResponse; +export type CreateBackfillDryRunData = { + requestBody: BackfillPostBody; +}; + +export type CreateBackfillDryRunResponse = BackfillDryRunResponse; + export type DeleteConnectionData = { connectionId: string; }; @@ -2645,7 +2651,7 @@ export type $OpenApiTs = { /** * Successful Response */ - 200: BackfillResponse | BackfillDryRunResponse; + 200: BackfillResponse; /** * Unauthorized */ @@ -2789,6 +2795,37 @@ export type $OpenApiTs = { }; }; }; + "/public/backfills/dry_run": { + post: { + req: CreateBackfillDryRunData; + res: { + /** + * Successful Response + */ + 200: BackfillDryRunResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Conflict + */ + 409: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; "/public/connections/{connection_id}": { delete: { req: DeleteConnectionData; diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index c24b2b6cc8a8e..b4eae7010cf0f 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -295,7 +295,7 @@ def test_create_backfill_dry_run( } response = test_client.post( - url="/public/backfills", + url="/public/backfills/dry_run", json=data, ) From dcf0e2dbd44062a4dfdff07f8fe6abed7f881970 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Sun, 29 Dec 2024 14:24:03 +0530 Subject: [PATCH 6/9] Refactor _create_backfill function --- .../core_api/datamodels/backfills.py | 8 +- .../core_api/openapi/v1-generated.yaml | 56 ++--- .../core_api/routes/public/backfills.py | 68 ++---- airflow/models/backfill.py | 231 +++++++++++------- airflow/ui/openapi-gen/queries/queries.ts | 2 +- .../ui/openapi-gen/requests/schemas.gen.ts | 71 +++--- .../ui/openapi-gen/requests/services.gen.ts | 2 +- airflow/ui/openapi-gen/requests/types.gen.ts | 34 +-- .../core_api/routes/public/test_backfills.py | 2 +- 9 files changed, 254 insertions(+), 220 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/backfills.py b/airflow/api_fastapi/core_api/datamodels/backfills.py index c5f6bd86ea242..4926f28693e8e 100644 --- a/airflow/api_fastapi/core_api/datamodels/backfills.py +++ b/airflow/api_fastapi/core_api/datamodels/backfills.py @@ -33,7 +33,6 @@ class BackfillPostBody(BaseModel): dag_run_conf: dict = {} reprocess_behavior: ReprocessBehavior = ReprocessBehavior.NONE max_active_runs: int = 10 - dry_run: bool = False class BackfillResponse(BaseModel): @@ -59,13 +58,14 @@ class BackfillCollectionResponse(BaseModel): total_entries: int -class BackfillRunInfo(BaseModel): +class DryRunBackfillResponse(BaseModel): """Data model for run information during a backfill operation.""" logical_date: datetime -class BackfillDryRunResponse(BaseModel): +class DryRunBackfillCollectionResponse(BaseModel): """Serializer for responses in dry-run mode for backfill operations.""" - run_info_list: list[BackfillRunInfo] + backfills: list[DryRunBackfillResponse] + total_entries: int diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 41081fc97c603..a6560e7babba7 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1383,7 +1383,7 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/BackfillDryRunResponse' + $ref: '#/components/schemas/DryRunBackfillCollectionResponse' '401': description: Unauthorized content: @@ -6265,18 +6265,6 @@ components: - total_entries title: BackfillCollectionResponse description: Backfill Collection serializer for responses. - BackfillDryRunResponse: - properties: - run_info_list: - items: - $ref: '#/components/schemas/BackfillRunInfo' - type: array - title: Run Info List - type: object - required: - - run_info_list - title: BackfillDryRunResponse - description: Serializer for responses in dry-run mode for backfill operations. BackfillPostBody: properties: dag_id: @@ -6305,10 +6293,6 @@ components: type: integer title: Max Active Runs default: 10 - dry_run: - type: boolean - title: Dry Run - default: false type: object required: - dag_id @@ -6372,17 +6356,6 @@ components: - updated_at title: BackfillResponse description: Base serializer for Backfill. - BackfillRunInfo: - properties: - logical_date: - type: string - format: date-time - title: Logical Date - type: object - required: - - logical_date - title: BackfillRunInfo - description: Data model for run information during a backfill operation. BaseInfoResponse: properties: status: @@ -7800,6 +7773,33 @@ components: This is the set of allowable values for the ``warning_type`` field in the DagWarning model.' + DryRunBackfillCollectionResponse: + properties: + backfills: + items: + $ref: '#/components/schemas/DryRunBackfillResponse' + type: array + title: Backfills + total_entries: + type: integer + title: Total Entries + type: object + required: + - backfills + - total_entries + title: DryRunBackfillCollectionResponse + description: Serializer for responses in dry-run mode for backfill operations. + DryRunBackfillResponse: + properties: + logical_date: + type: string + format: date-time + title: Logical Date + type: object + required: + - logical_date + title: DryRunBackfillResponse + description: Data model for run information during a backfill operation. EdgeResponse: properties: is_setup_teardown: diff --git a/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow/api_fastapi/core_api/routes/public/backfills.py index be4c20957ed03..931ad167e9e10 100644 --- a/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -19,7 +19,7 @@ from typing import Annotated from fastapi import Depends, HTTPException, status -from sqlalchemy import desc, select, update +from sqlalchemy import select, update from airflow.api_fastapi.common.db.common import ( AsyncSessionDep, @@ -30,10 +30,10 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.datamodels.backfills import ( BackfillCollectionResponse, - BackfillDryRunResponse, BackfillPostBody, BackfillResponse, - BackfillRunInfo, + DryRunBackfillCollectionResponse, + DryRunBackfillResponse, ) from airflow.api_fastapi.core_api.openapi.exceptions import ( create_openapi_http_exception_doc, @@ -43,14 +43,9 @@ AlreadyRunningBackfill, Backfill, BackfillDagRun, - BackfillDagRunExceptionReason, - ReprocessBehavior, _create_backfill, - _get_info_list, ) -from airflow.models.serialized_dag import SerializedDagModel from airflow.utils import timezone -from airflow.utils.sqlalchemy import nulls_first from airflow.utils.state import DagRunState backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") @@ -225,43 +220,28 @@ def create_backfill( ), ) def create_backfill_dry_run( - backfill_request: BackfillPostBody, - session: SessionDep, -) -> BackfillDryRunResponse: - from_date = timezone.coerce_datetime(backfill_request.from_date) - to_date = timezone.coerce_datetime(backfill_request.to_date) - serdag = session.scalar(SerializedDagModel.latest_item_select_object(backfill_request.dag_id)) - if not serdag: - raise HTTPException(status_code=404, detail=f"Could not find dag {backfill_request.dag_id}") + body: BackfillPostBody, +) -> DryRunBackfillCollectionResponse: + from_date = timezone.coerce_datetime(body.from_date) + to_date = timezone.coerce_datetime(body.to_date) - info_list = _get_info_list( - dag=serdag.dag, - from_date=from_date, - to_date=to_date, - reverse=backfill_request.run_backwards, - ) - backfill_response_item = [] - for info in info_list: - dr = session.scalar( - select(DagRun) - .where(DagRun.logical_date == info.logical_date) - .order_by(nulls_first(desc(DagRun.start_date), session)) - .limit(1) + try: + backfills_dry_run = _create_backfill( + dag_id=body.dag_id, + from_date=from_date, + to_date=to_date, + max_active_runs=body.max_active_runs, + reverse=body.run_backwards, + dag_run_conf=body.dag_run_conf, + reprocess_behavior=body.reprocess_behavior, + dry_run=True, ) + backfills = [DryRunBackfillResponse(logical_date=logical_date) for logical_date in backfills_dry_run] - if dr: - non_create_reason = None - if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): - non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT - elif backfill_request.reprocess_behavior is ReprocessBehavior.NONE: - non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - elif backfill_request.reprocess_behavior is ReprocessBehavior.FAILED: - if dr.state != DagRunState.FAILED: - non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - if not non_create_reason: - backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) - - else: - backfill_response_item.append(BackfillRunInfo(logical_date=info.logical_date)) + return DryRunBackfillCollectionResponse(backfills=backfills, total_entries=len(backfills_dry_run)) - return BackfillDryRunResponse(run_info_list=backfill_response_item) + except AlreadyRunningBackfill: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="There is already a running backfill for the dag", + ) diff --git a/airflow/models/backfill.py b/airflow/models/backfill.py index 0e88fa15bb04f..32d1f99072687 100644 --- a/airflow/models/backfill.py +++ b/airflow/models/backfill.py @@ -25,7 +25,7 @@ import logging from enum import Enum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload from sqlalchemy import ( Boolean, @@ -33,7 +33,6 @@ ForeignKeyConstraint, Integer, UniqueConstraint, - desc, func, select, ) @@ -47,13 +46,15 @@ from airflow.settings import json from airflow.utils import timezone from airflow.utils.session import create_session -from airflow.utils.sqlalchemy import UtcDateTime, nulls_first, with_row_locks +from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType if TYPE_CHECKING: from datetime import datetime + from typing_extensions import Literal + log = logging.getLogger(__name__) @@ -158,72 +159,98 @@ def validate_sort_ordinal(self, key, val): def _create_backfill_dag_run( *, dag, - info, - reprocess_behavior: ReprocessBehavior, + dagrun_info_list, + reprocess_behavior: ReprocessBehavior | None = None, backfill_id, - dag_run_conf, - backfill_sort_ordinal, + dag_run_conf: dict | None, session, -): + dry_run, +) -> list[datetime]: from airflow.models import DagRun - with session.begin_nested() as nested: - dr = session.scalar( - with_row_locks( - select(DagRun) - .where(DagRun.logical_date == info.logical_date) - .order_by(nulls_first(desc(DagRun.start_date), session=session)) - .limit(1), - session=session, - ) - ) - if dr: - non_create_reason = None - if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): - non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT - elif reprocess_behavior is ReprocessBehavior.NONE: - non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - elif reprocess_behavior is ReprocessBehavior.FAILED: - if dr.state != DagRunState.FAILED: + backfill_sort_ordinal = 0 + logical_dates = [] + dagrun_infos = list(dagrun_info_list) + + if reprocess_behavior is None: + reprocess_behavior = ReprocessBehavior.NONE + if dag_run_conf is None: + dag_run_conf = {} + + existing_dag_runs = { + dr.logical_date: dr + for dr in session.scalars( + select(DagRun) + .where(DagRun.dag_id == dag.dag_id) + .where(DagRun.logical_date.in_([info.logical_date for info in dagrun_infos])) + ).all() + } + + for info in dagrun_infos: + backfill_sort_ordinal += 1 + dr = existing_dag_runs.get(info.logical_date) + non_create_reason = None + + with session.begin_nested() as nested: + if dr: + if dr.state not in (DagRunState.SUCCESS, DagRunState.FAILED): + non_create_reason = BackfillDagRunExceptionReason.IN_FLIGHT + elif reprocess_behavior is ReprocessBehavior.NONE: non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS - if non_create_reason: - # rolling back here restores to start of this nested tran - # which releases the lock on the latest dag run, since we - # are not creating a new one - nested.rollback() + elif reprocess_behavior is ReprocessBehavior.FAILED: + if dr.state != DagRunState.FAILED: + non_create_reason = BackfillDagRunExceptionReason.ALREADY_EXISTS + + if non_create_reason: + if not dry_run: + nested.rollback() + session.add( + BackfillDagRun( + backfill_id=backfill_id, + dag_run_id=None, + logical_date=dr.logical_date, + exception_reason=non_create_reason, + sort_ordinal=backfill_sort_ordinal, + ) + ) + else: + logical_dates.append(dr.logical_date) + else: + logical_dates.append(info.logical_date) + + if not non_create_reason and not dry_run: + dag_version = DagVersion.get_latest_version(dag.dag_id, session=session) + dr = dag.create_dagrun( + triggered_by=DagRunTriggeredByType.BACKFILL, + logical_date=info.logical_date, + data_interval=info.data_interval, + start_date=timezone.utcnow(), + state=DagRunState.QUEUED, + external_trigger=False, + conf=dag_run_conf, + run_type=DagRunType.BACKFILL_JOB, + creating_job_id=None, + session=session, + backfill_id=backfill_id, + dag_version=dag_version, + ) session.add( BackfillDagRun( backfill_id=backfill_id, - dag_run_id=None, - logical_date=info.logical_date, - exception_reason=non_create_reason, + dag_run_id=dr.id, sort_ordinal=backfill_sort_ordinal, + logical_date=info.logical_date, ) ) - return - dag_version = DagVersion.get_latest_version(dag.dag_id, session=session) - dr = dag.create_dagrun( - triggered_by=DagRunTriggeredByType.BACKFILL, - logical_date=info.logical_date, - data_interval=info.data_interval, - start_date=timezone.utcnow(), - state=DagRunState.QUEUED, - external_trigger=False, - conf=dag_run_conf, - run_type=DagRunType.BACKFILL_JOB, - creating_job_id=None, - session=session, - backfill_id=backfill_id, - dag_version=dag_version, - ) - session.add( - BackfillDagRun( - backfill_id=backfill_id, - dag_run_id=dr.id, - sort_ordinal=backfill_sort_ordinal, - logical_date=info.logical_date, - ) - ) + + log.info( + "created backfill dag run dag_id=%s backfill_id=%s, info=%s", + dag.dag_id, + backfill_id, + info, + ) + + return logical_dates def _get_info_list( @@ -241,6 +268,34 @@ def _get_info_list( return dagrun_info_list +@overload +def _create_backfill( + *, + dag_id: str, + from_date: datetime, + to_date: datetime, + max_active_runs: int, + reverse: bool, + dag_run_conf: dict | None, + reprocess_behavior: ReprocessBehavior | None = ..., + dry_run: Literal[True], +) -> list[datetime]: ... + + +@overload +def _create_backfill( + *, + dag_id: str, + from_date: datetime, + to_date: datetime, + max_active_runs: int, + reverse: bool, + dag_run_conf: dict | None, + reprocess_behavior: ReprocessBehavior | None = ..., + dry_run: Literal[False] | None = ..., +) -> Backfill | None: ... + + def _create_backfill( *, dag_id: str, @@ -250,7 +305,8 @@ def _create_backfill( reverse: bool, dag_run_conf: dict | None, reprocess_behavior: ReprocessBehavior | None = None, -) -> Backfill | None: + dry_run: bool | None = False, +) -> Backfill | list[datetime] | None: from airflow.models import DagModel from airflow.models.serialized_dag import SerializedDagModel @@ -284,18 +340,23 @@ def _create_backfill( "You must set reprocess behavior to reprocess completed or " "reprocess failed" ) - br = Backfill( - dag_id=dag_id, - from_date=from_date, - to_date=to_date, - max_active_runs=max_active_runs, - dag_run_conf=dag_run_conf, - reprocess_behavior=reprocess_behavior, - ) - session.add(br) - session.commit() - backfill_sort_ordinal = 0 + backfill_id = None + + if not dry_run: + br = Backfill( + dag_id=dag_id, + from_date=from_date, + to_date=to_date, + max_active_runs=max_active_runs, + dag_run_conf=dag_run_conf, + reprocess_behavior=reprocess_behavior, + ) + + session.add(br) + session.commit() + + backfill_id = br.id dagrun_info_list = _get_info_list( from_date=from_date, @@ -316,21 +377,15 @@ def _create_backfill( ) if not dag_model: raise RuntimeError(f"Dag {dag_id} not found") - for info in dagrun_info_list: - backfill_sort_ordinal += 1 - _create_backfill_dag_run( - dag=dag, - info=info, - backfill_id=br.id, - dag_run_conf=br.dag_run_conf, - reprocess_behavior=br.reprocess_behavior, - backfill_sort_ordinal=backfill_sort_ordinal, - session=session, - ) - log.info( - "created backfill dag run dag_id=%s backfill_id=%s, info=%s", - dag.dag_id, - br.id, - info, - ) - return br + + backfill_response = _create_backfill_dag_run( + dag=dag, + dagrun_info_list=dagrun_info_list, + backfill_id=backfill_id, + dag_run_conf=dag_run_conf, + reprocess_behavior=reprocess_behavior, + session=session, + dry_run=dry_run, + ) + + return br if not dry_run else backfill_response diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index b0bef485a23ee..9de4e5e9275d1 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -2817,7 +2817,7 @@ export const useBackfillServiceCreateBackfill = < * Create Backfill Dry Run * @param data The data for the request. * @param data.requestBody - * @returns BackfillDryRunResponse Successful Response + * @returns DryRunBackfillCollectionResponse Successful Response * @throws ApiError */ export const useBackfillServiceCreateBackfillDryRun = < diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 995bb3e71c4bb..a1c232465b3ac 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -350,23 +350,6 @@ export const $BackfillCollectionResponse = { description: "Backfill Collection serializer for responses.", } as const; -export const $BackfillDryRunResponse = { - properties: { - run_info_list: { - items: { - $ref: "#/components/schemas/BackfillRunInfo", - }, - type: "array", - title: "Run Info List", - }, - }, - type: "object", - required: ["run_info_list"], - title: "BackfillDryRunResponse", - description: - "Serializer for responses in dry-run mode for backfill operations.", -} as const; - export const $BackfillPostBody = { properties: { dag_id: { @@ -402,11 +385,6 @@ export const $BackfillPostBody = { title: "Max Active Runs", default: 10, }, - dry_run: { - type: "boolean", - title: "Dry Run", - default: false, - }, }, type: "object", required: ["dag_id", "from_date", "to_date"], @@ -490,20 +468,6 @@ export const $BackfillResponse = { description: "Base serializer for Backfill.", } as const; -export const $BackfillRunInfo = { - properties: { - logical_date: { - type: "string", - format: "date-time", - title: "Logical Date", - }, - }, - type: "object", - required: ["logical_date"], - title: "BackfillRunInfo", - description: "Data model for run information during a backfill operation.", -} as const; - export const $BaseInfoResponse = { properties: { status: { @@ -2666,6 +2630,41 @@ This is the set of allowable values for the \`\`warning_type\`\` field in the DagWarning model.`, } as const; +export const $DryRunBackfillCollectionResponse = { + properties: { + backfills: { + items: { + $ref: "#/components/schemas/DryRunBackfillResponse", + }, + type: "array", + title: "Backfills", + }, + total_entries: { + type: "integer", + title: "Total Entries", + }, + }, + type: "object", + required: ["backfills", "total_entries"], + title: "DryRunBackfillCollectionResponse", + description: + "Serializer for responses in dry-run mode for backfill operations.", +} as const; + +export const $DryRunBackfillResponse = { + properties: { + logical_date: { + type: "string", + format: "date-time", + title: "Logical Date", + }, + }, + type: "object", + required: ["logical_date"], + title: "DryRunBackfillResponse", + description: "Data model for run information during a backfill operation.", +} as const; + export const $EdgeResponse = { properties: { is_setup_teardown: { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index be8a0cda9660d..f25fb39a9639e 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -967,7 +967,7 @@ export class BackfillService { * Create Backfill Dry Run * @param data The data for the request. * @param data.requestBody - * @returns BackfillDryRunResponse Successful Response + * @returns DryRunBackfillCollectionResponse Successful Response * @throws ApiError */ public static createBackfillDryRun( diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 6494219e2479e..97016b99ace07 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -97,13 +97,6 @@ export type BackfillCollectionResponse = { total_entries: number; }; -/** - * Serializer for responses in dry-run mode for backfill operations. - */ -export type BackfillDryRunResponse = { - run_info_list: Array; -}; - /** * Object used for create backfill request. */ @@ -117,7 +110,6 @@ export type BackfillPostBody = { }; reprocess_behavior?: ReprocessBehavior; max_active_runs?: number; - dry_run?: boolean; }; /** @@ -139,13 +131,6 @@ export type BackfillResponse = { updated_at: string; }; -/** - * Data model for run information during a backfill operation. - */ -export type BackfillRunInfo = { - logical_date: string; -}; - /** * Base info serializer for responses. */ @@ -651,6 +636,21 @@ export type DagTagResponse = { */ export type DagWarningType = "asset conflict" | "non-existent pool"; +/** + * Serializer for responses in dry-run mode for backfill operations. + */ +export type DryRunBackfillCollectionResponse = { + backfills: Array; + total_entries: number; +}; + +/** + * Data model for run information during a backfill operation. + */ +export type DryRunBackfillResponse = { + logical_date: string; +}; + /** * Edge serializer for responses. */ @@ -1543,7 +1543,7 @@ export type CreateBackfillDryRunData = { requestBody: BackfillPostBody; }; -export type CreateBackfillDryRunResponse = BackfillDryRunResponse; +export type CreateBackfillDryRunResponse = DryRunBackfillCollectionResponse; export type DeleteConnectionData = { connectionId: string; @@ -2802,7 +2802,7 @@ export type $OpenApiTs = { /** * Successful Response */ - 200: BackfillDryRunResponse; + 200: DryRunBackfillCollectionResponse; /** * Unauthorized */ diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index b4eae7010cf0f..e98540c16ff98 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -301,7 +301,7 @@ def test_create_backfill_dry_run( assert response.status_code == 200 response_json = response.json() - assert response_json["run_info_list"] == expected_dates + assert response_json["backfills"] == expected_dates class TestCancelBackfill(TestBackfillEndpoint): From 828907c32b17e9c861488d500996c01062d2e6c0 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Sun, 29 Dec 2024 17:40:37 +0530 Subject: [PATCH 7/9] Fix static checks --- airflow/models/backfill.py | 28 +++++++++++++++++-- .../core_api/routes/public/test_backfills.py | 1 - 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/airflow/models/backfill.py b/airflow/models/backfill.py index 32d1f99072687..01b23086354ef 100644 --- a/airflow/models/backfill.py +++ b/airflow/models/backfill.py @@ -177,15 +177,39 @@ def _create_backfill_dag_run( if dag_run_conf is None: dag_run_conf = {} + dag_run_ranked = ( + select( + DagRun.logical_date, + DagRun.start_date, + DagRun.run_id, + DagRun.dag_id, + func.row_number() + .over( + partition_by=DagRun.logical_date, + order_by=(DagRun.start_date.desc().nullsfirst()), + ) + .label("row_number"), + ) + .where(DagRun.dag_id == dag.dag_id) + .where(DagRun.logical_date.in_([info.logical_date for info in dagrun_infos])) + .subquery() + ) + existing_dag_runs = { dr.logical_date: dr for dr in session.scalars( select(DagRun) - .where(DagRun.dag_id == dag.dag_id) - .where(DagRun.logical_date.in_([info.logical_date for info in dagrun_infos])) + .join( + dag_run_ranked, + (DagRun.logical_date == dag_run_ranked.c.logical_date) + & (DagRun.dag_id == dag_run_ranked.c.dag_id), + ) + .where(dag_run_ranked.c.row_number == 1) ).all() } + print(existing_dag_runs) + for info in dagrun_infos: backfill_sort_ordinal += 1 dr = existing_dag_runs.get(info.logical_date) diff --git a/tests/api_fastapi/core_api/routes/public/test_backfills.py b/tests/api_fastapi/core_api/routes/public/test_backfills.py index e98540c16ff98..4f0ae7918e41d 100644 --- a/tests/api_fastapi/core_api/routes/public/test_backfills.py +++ b/tests/api_fastapi/core_api/routes/public/test_backfills.py @@ -290,7 +290,6 @@ def test_create_backfill_dry_run( "max_active_runs": 5, "run_backwards": False, "dag_run_conf": {"param1": "val1", "param2": True}, - "dry_run": True, "reprocess_behavior": reprocess_behavior, } From b130be48530672261fb324ad0d2ab6d6567d7bf3 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Sun, 29 Dec 2024 19:28:00 +0530 Subject: [PATCH 8/9] Fix static checks --- airflow/models/backfill.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/airflow/models/backfill.py b/airflow/models/backfill.py index 01b23086354ef..cf55fc300cf2c 100644 --- a/airflow/models/backfill.py +++ b/airflow/models/backfill.py @@ -181,7 +181,6 @@ def _create_backfill_dag_run( select( DagRun.logical_date, DagRun.start_date, - DagRun.run_id, DagRun.dag_id, func.row_number() .over( @@ -202,6 +201,10 @@ def _create_backfill_dag_run( .join( dag_run_ranked, (DagRun.logical_date == dag_run_ranked.c.logical_date) + & ( + (DagRun.start_date == dag_run_ranked.c.start_date) + | ((DagRun.start_date.is_(None)) & (dag_run_ranked.c.start_date.is_(None))) + ) & (DagRun.dag_id == dag_run_ranked.c.dag_id), ) .where(dag_run_ranked.c.row_number == 1) From f4b7602a2094a03330f2b99ec54776c30aa69282 Mon Sep 17 00:00:00 2001 From: Sneha Prabhu Date: Sun, 29 Dec 2024 23:55:02 +0530 Subject: [PATCH 9/9] Fix mysql static checks --- airflow/models/backfill.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/airflow/models/backfill.py b/airflow/models/backfill.py index cf55fc300cf2c..c9158bbad240d 100644 --- a/airflow/models/backfill.py +++ b/airflow/models/backfill.py @@ -27,15 +27,7 @@ from enum import Enum from typing import TYPE_CHECKING, overload -from sqlalchemy import ( - Boolean, - Column, - ForeignKeyConstraint, - Integer, - UniqueConstraint, - func, - select, -) +from sqlalchemy import Boolean, Column, ForeignKeyConstraint, Integer, UniqueConstraint, case, func, select from sqlalchemy.orm import relationship, validates from sqlalchemy_jsonfield import JSONField @@ -185,7 +177,7 @@ def _create_backfill_dag_run( func.row_number() .over( partition_by=DagRun.logical_date, - order_by=(DagRun.start_date.desc().nullsfirst()), + order_by=(case([(DagRun.start_date.is_(None), 0)], else_=1), DagRun.start_date.desc()), ) .label("row_number"), )