Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/serialization/pydantic/dataset.py: 100%

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 datetime import datetime 

18from typing import List, Optional 

19 

20from airflow.utils.pydantic import BaseModel as BaseModelPydantic, ConfigDict 

21 

22 

23class DagScheduleDatasetReferencePydantic(BaseModelPydantic): 

24 """Serializable version of the DagScheduleDatasetReference ORM SqlAlchemyModel used by internal API.""" 

25 

26 dataset_id: int 

27 dag_id: str 

28 created_at: datetime 

29 updated_at: datetime 

30 

31 model_config = ConfigDict(from_attributes=True) 

32 

33 

34class TaskOutletDatasetReferencePydantic(BaseModelPydantic): 

35 """Serializable version of the TaskOutletDatasetReference ORM SqlAlchemyModel used by internal API.""" 

36 

37 dataset_id: int 

38 dag_id: str 

39 task_id: str 

40 created_at: datetime 

41 updated_at: datetime 

42 

43 model_config = ConfigDict(from_attributes=True) 

44 

45 

46class DatasetPydantic(BaseModelPydantic): 

47 """Serializable representation of the Dataset ORM SqlAlchemyModel used by internal API.""" 

48 

49 id: int 

50 uri: str 

51 extra: Optional[dict] 

52 created_at: datetime 

53 updated_at: datetime 

54 is_orphaned: bool 

55 

56 consuming_dags: List[DagScheduleDatasetReferencePydantic] 

57 producing_tasks: List[TaskOutletDatasetReferencePydantic] 

58 

59 model_config = ConfigDict(from_attributes=True) 

60 

61 

62class DatasetEventPydantic(BaseModelPydantic): 

63 """Serializable representation of the DatasetEvent ORM SqlAlchemyModel used by internal API.""" 

64 

65 id: int 

66 dataset_id: Optional[int] 

67 extra: dict 

68 source_task_id: Optional[str] 

69 source_dag_id: Optional[str] 

70 source_run_id: Optional[str] 

71 source_map_index: Optional[int] 

72 timestamp: datetime 

73 dataset: Optional[DatasetPydantic] 

74 

75 model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)