Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/datamodels/ui/common.py: 81%

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. 

17 

18from __future__ import annotations 

19 

20import inspect 

21from datetime import datetime 

22from typing import Generic, Literal, TypeVar 

23 

24from pydantic import computed_field, field_validator 

25 

26from airflow._shared.timezones import timezone 

27from airflow.api_fastapi.core_api.base import BaseModel 

28from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse 

29from airflow.utils.state import DagRunState 

30from airflow.utils.types import DagRunType 

31 

32 

33class BaseEdgeResponse(BaseModel): 

34 """Base Edge serializer for responses.""" 

35 

36 source_id: str 

37 target_id: str 

38 

39 

40class BaseNodeResponse(BaseModel): 

41 """Base Node serializer for responses.""" 

42 

43 id: str 

44 label: str 

45 type: Literal[ 

46 "join", 

47 "task", 

48 "asset-condition", 

49 "asset", 

50 "asset-alias", 

51 "asset-name-ref", 

52 "asset-uri-ref", 

53 "dag", 

54 "sensor", 

55 "trigger", 

56 ] 

57 team: str | None = None 

58 

59 

60E = TypeVar("E", bound=BaseEdgeResponse) 

61N = TypeVar("N", bound=BaseNodeResponse) 

62 

63 

64class GridNodeResponse(BaseModel): 

65 """Base Node serializer for responses.""" 

66 

67 id: str 

68 label: str 

69 children: list[GridNodeResponse] | None = None 

70 is_mapped: bool | None 

71 setup_teardown_type: Literal["setup", "teardown"] | None = None 

72 doc_md: str | None = None 

73 

74 @field_validator("doc_md", mode="before") 

75 @classmethod 

76 def get_doc_md(cls, doc_md: str | None) -> str | None: 

77 """Clean indentation in doc md.""" 

78 if doc_md is None: 

79 return None 

80 return inspect.cleandoc(doc_md) 

81 

82 

83class GridRunsResponse(BaseModel): 

84 """Base Node serializer for responses.""" 

85 

86 dag_id: str 

87 run_id: str 

88 queued_at: datetime | None 

89 start_date: datetime | None 

90 end_date: datetime | None 

91 run_after: datetime 

92 state: DagRunState | None 

93 run_type: DagRunType 

94 dag_versions: list[DagVersionResponse] = [] 

95 has_missed_deadline: bool 

96 has_note: bool 

97 

98 @computed_field 

99 def duration(self) -> float: 

100 if self.start_date: 

101 end_date = self.end_date or timezone.utcnow() 

102 return (end_date - self.start_date).total_seconds() 

103 return 0 

104 

105 

106class BaseGraphResponse(BaseModel, Generic[E, N]): 

107 """Base Graph serializer for responses.""" 

108 

109 edges: list[E] 

110 nodes: list[N]