Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/models/base.py: 88%

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 Any 

21 

22from sqlalchemy import MetaData, String 

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 

50Base: Any = mapper_registry.generate_base() 

51 

52ID_LEN = 250 

53 

54 

55def get_id_collation_args(): 

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

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

58 if collation: 

59 return {"collation": collation} 

60 else: 

61 # Automatically use utf8mb3_bin collation for mysql 

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

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

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

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

66 # 

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

68 # 

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

70 # parameters, so we use the connection 

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

72 if conn.startswith("mysql") or conn.startswith("mariadb"): 

73 return {"collation": "utf8mb3_bin"} 

74 return {} 

75 

76 

77COLLATION_ARGS = get_id_collation_args() 

78 

79 

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

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