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

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, mapped_column, registry 

24 

25from airflow.configuration import conf 

26 

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

28 

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

30# below mean, see: 

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

32naming_convention = { 

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

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

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

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

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

38} 

39 

40 

41def _get_schema(): 

42 if not SQL_ALCHEMY_SCHEMA or SQL_ALCHEMY_SCHEMA.isspace(): 

43 return None 

44 return SQL_ALCHEMY_SCHEMA 

45 

46 

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

48mapper_registry = registry(metadata=metadata) 

49_sentinel = object() 

50 

51if TYPE_CHECKING: 

52 Base = Any 

53else: 

54 Base = mapper_registry.generate_base() 

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

56 # unblock the development of SQLA2 support. 

57 Base.__allow_unmapped__ = True 

58 

59ID_LEN = 250 

60 

61 

62def get_id_collation_args(): 

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

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

65 if collation: 

66 return {"collation": collation} 

67 # Automatically use utf8mb3_bin collation for mysql 

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

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

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

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

72 # 

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

74 # 

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

76 # parameters, so we use the connection 

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

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

79 return {"collation": "utf8mb3_bin"} 

80 return {} 

81 

82 

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

84 

85 

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

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

88 

89 

90class TaskInstanceDependencies(Base): 

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

92 

93 __abstract__ = True 

94 

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

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

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

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