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

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 Column, Integer, MetaData, String, text 

23from sqlalchemy.orm import 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 

56ID_LEN = 250 

57 

58 

59def get_id_collation_args(): 

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

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

62 if collation: 

63 return {"collation": collation} 

64 else: 

65 # Automatically use utf8mb3_bin collation for mysql 

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

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

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

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

70 # 

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

72 # 

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

74 # parameters, so we use the connection 

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

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

77 return {"collation": "utf8mb3_bin"} 

78 return {} 

79 

80 

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

82 

83 

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

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

86 

87 

88class TaskInstanceDependencies(Base): 

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

90 

91 __abstract__ = True 

92 

93 task_id = Column(StringID(), nullable=False) 

94 dag_id = Column(StringID(), nullable=False) 

95 run_id = Column(StringID(), nullable=False) 

96 map_index = Column(Integer, nullable=False, server_default=text("-1"))