Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/models/base.py: 87%

1# 

2# Licensed to the Apache Software Foundation (ASF) under one 

3# or more contributor license agreements. See the NOTICE file 

4# distributed with this work for additional information 

5# regarding copyright ownership. The ASF licenses this file 

6# to you under the Apache License, Version 2.0 (the 

7# "License"); you may not use this file except in compliance 

8# with the License. You may obtain a copy of the License at 

9# 

10# http://www.apache.org/licenses/LICENSE-2.0 

11# 

12# Unless required by applicable law or agreed to in writing, 

13# software distributed under the License is distributed on an 

14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

15# KIND, either express or implied. See the License for the 

16# specific language governing permissions and limitations 

17# under the License. 

18from __future__ import annotations 

19 

20from typing import TYPE_CHECKING, Any 

21 

22from sqlalchemy import Integer, MetaData, String, text 

23from sqlalchemy.orm import Mapped, registry 

24 

25from airflow.configuration import conf 

26from airflow.utils.sqlalchemy import is_sqlalchemy_v1, mapped_column 

27 

28SQL_ALCHEMY_SCHEMA = conf.get("database", "SQL_ALCHEMY_SCHEMA") 

29 

30# For more information about what the tokens in the naming convention 

31# below mean, see: 

32# https://docs.sqlalchemy.org/en/14/core/metadata.html#sqlalchemy.schema.MetaData.params.naming_convention 

33naming_convention = { 

34 "ix": "idx_%(column_0_N_label)s", 

35 "uq": "%(table_name)s_%(column_0_N_name)s_uq", 

36 "ck": "ck_%(table_name)s_%(constraint_name)s", 

37 "fk": "%(table_name)s_%(column_0_name)s_fkey", 

38 "pk": "%(table_name)s_pkey", 

39} 

40 

41 

42def _get_schema(): 

43 if not SQL_ALCHEMY_SCHEMA or SQL_ALCHEMY_SCHEMA.isspace(): 

44 return None 

45 return SQL_ALCHEMY_SCHEMA 

46 

47 

48metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention) 

49mapper_registry = registry(metadata=metadata) 

50_sentinel = object() 

51 

52if TYPE_CHECKING: 

53 Base = Any 

54else: 

55 Base = mapper_registry.generate_base() 

56 # TEMPORARY workaround to allow using unmapped (v1.4) models in SQLAlchemy 2.0. It is intended only to 

57 # unblock the development of SQLA2 support. 

58 if not is_sqlalchemy_v1(): 

59 Base.__allow_unmapped__ = True 

60 

61ID_LEN = 250 

62 

63 

64def get_id_collation_args(): 

65 """Get SQLAlchemy args to use for COLLATION.""" 

66 collation = conf.get("database", "sql_engine_collation_for_ids", fallback=None) 

67 if collation: 

68 return {"collation": collation} 

69 # Automatically use utf8mb3_bin collation for mysql 

70 # This is backwards-compatible. All our IDS are ASCII anyway so even if 

71 # we migrate from previously installed database with different collation and we end up mixture of 

72 # COLLATIONS, it's not a problem whatsoever (and we keep it small enough so that our indexes 

73 # for MYSQL will not exceed the maximum index size. 

74 # 

75 # See https://github.com/apache/airflow/pull/17603#issuecomment-901121618. 

76 # 

77 # We cannot use session/dialect as at this point we are trying to determine the right connection 

78 # parameters, so we use the connection 

79 conn = conf.get("database", "sql_alchemy_conn", fallback="") 

80 if conn.startswith(("mysql", "mariadb")): 

81 return {"collation": "utf8mb3_bin"} 

82 return {} 

83 

84 

85COLLATION_ARGS: dict[str, Any] = get_id_collation_args() 

86 

87 

88def StringID(*, length=ID_LEN, **kwargs) -> String: 

89 return String(length=length, **kwargs, **COLLATION_ARGS) 

90 

91 

92class TaskInstanceDependencies(Base): 

93 """Base class for depending models linked to TaskInstance.""" 

94 

95 __abstract__ = True 

96 

97 task_id: Mapped[str] = mapped_column(StringID(), nullable=False) 

98 dag_id: Mapped[str] = mapped_column(StringID(), nullable=False) 

99 run_id: Mapped[str] = mapped_column(StringID(), nullable=False) 

100 map_index: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("-1"))