Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/base.py: 89%

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 

18 

19from abc import ABC, abstractmethod 

20from typing import TYPE_CHECKING, Generic, TypeVar, Union, get_args, get_origin 

21 

22from pydantic import BaseModel as PydanticBaseModel, ConfigDict, create_model 

23 

24if TYPE_CHECKING: 24 ↛ 25line 24 didn't jump to line 25 because the condition on line 24 was never true

25 from sqlalchemy.sql import Select 

26 

27T = TypeVar("T") 

28 

29 

30class BaseModel(PydanticBaseModel): 

31 """ 

32 Base pydantic model for REST API. 

33 

34 :meta private: 

35 """ 

36 

37 model_config = ConfigDict(from_attributes=True, populate_by_name=True) 

38 

39 

40class StrictBaseModel(BaseModel): 

41 """ 

42 StrictBaseModel is a base Pydantic model for REST API that does not allow any extra fields. 

43 

44 Use this class for models that should not have any extra fields in the payload. 

45 

46 :meta private: 

47 """ 

48 

49 model_config = ConfigDict(from_attributes=True, populate_by_name=True, extra="forbid") 

50 

51 

52def make_partial_model(model: type[PydanticBaseModel]) -> type[PydanticBaseModel]: 

53 """Create a version of a Pydantic model where all fields are Optional with default=None.""" 

54 field_overrides: dict = {} 

55 for field_name, field_info in model.model_fields.items(): 

56 ann = field_info.annotation 

57 origin = get_origin(ann) 

58 if not (origin is Union and type(None) in get_args(ann)): 58 ↛ 60line 58 didn't jump to line 60 because the condition on line 58 was always true

59 ann = ann | None # type: ignore[operator, assignment] 

60 new_info = field_info._copy() 

61 new_info.default = None 

62 new_info._attributes_set["default"] = None 

63 field_overrides[field_name] = (ann, new_info) 

64 

65 return create_model( 

66 f"{model.__name__}Partial", 

67 __base__=model, 

68 **field_overrides, 

69 ) 

70 

71 

72class OrmClause(Generic[T], ABC): 

73 """ 

74 Base class for filtering clauses with paginated_select. 

75 

76 The subclasses should implement the `to_orm` method and set the `value` attribute. 

77 """ 

78 

79 def __init__(self, value: T | None = None): 

80 self.value = value 

81 

82 @abstractmethod 

83 def to_orm(self, select: Select) -> Select: 

84 pass