Coverage for python/felis/db/_variants.py: 40%

33 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-25 10:20 -0700

1# This file is part of felis. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22import re 

23from typing import Any 

24 

25from sqlalchemy import types 

26from sqlalchemy.dialects import mysql, oracle, postgresql, sqlite 

27from sqlalchemy.types import TypeEngine 

28 

29from ..datamodel import Column 

30 

31MYSQL = "mysql" 

32ORACLE = "oracle" 

33POSTGRES = "postgresql" 

34SQLITE = "sqlite" 

35 

36TABLE_OPTS = { 

37 "mysql:engine": "mysql_engine", 

38 "mysql:charset": "mysql_charset", 

39 "oracle:compress": "oracle_compress", 

40} 

41 

42COLUMN_VARIANT_OVERRIDE = { 

43 "mysql_datatype": "mysql", 

44 "oracle_datatype": "oracle", 

45 "postgresql_datatype": "postgresql", 

46 "sqlite_datatype": "sqlite", 

47} 

48 

49DIALECT_MODULES = {MYSQL: mysql, ORACLE: oracle, SQLITE: sqlite, POSTGRES: postgresql} 

50 

51_length_regex = re.compile(r"\((\d+)\)") 

52"""A regular expression that is looking for numbers within parentheses.""" 

53 

54 

55def process_variant_override(dialect_name: str, variant_override_str: str) -> types.TypeEngine: 

56 """Return variant type for given dialect.""" 

57 dialect = DIALECT_MODULES[dialect_name] 

58 variant_type_name = variant_override_str.split("(")[0] 

59 

60 # Process Variant Type 

61 if variant_type_name not in dir(dialect): 

62 raise ValueError(f"Type {variant_type_name} not found in dialect {dialect_name}") 

63 variant_type = getattr(dialect, variant_type_name) 

64 length_params = [] 

65 if match := _length_regex.search(variant_override_str): 

66 length_params.extend([int(i) for i in match.group(1).split(",")]) 

67 return variant_type(*length_params) 

68 

69 

70def make_variant_dict(column_obj: Column) -> dict[str, TypeEngine[Any]]: 

71 """Handle variant overrides for a `felis.datamodel.Column`. 

72 

73 This function will return a dictionary of `str` to 

74 `sqlalchemy.types.TypeEngine` containing variant datatype information 

75 (e.g., for mysql, postgresql, etc). 

76 

77 Parameters 

78 ---------- 

79 column_obj : `felis.datamodel.Column` 

80 The column object from which to build the variant dictionary. 

81 

82 Returns 

83 ------- 

84 variant_dict : `dict` 

85 The dictionary of `str` to `sqlalchemy.types.TypeEngine` containing 

86 variant datatype information (e.g., for mysql, postgresql, etc). 

87 """ 

88 variant_dict = {} 

89 for field_name, value in iter(column_obj): 

90 if field_name in COLUMN_VARIANT_OVERRIDE and value is not None: 

91 dialect = COLUMN_VARIANT_OVERRIDE[field_name] 

92 variant: TypeEngine = process_variant_override(dialect, value) 

93 variant_dict[dialect] = variant 

94 return variant_dict