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
« 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/>.
22import re
23from typing import Any
25from sqlalchemy import types
26from sqlalchemy.dialects import mysql, oracle, postgresql, sqlite
27from sqlalchemy.types import TypeEngine
29from ..datamodel import Column
31MYSQL = "mysql"
32ORACLE = "oracle"
33POSTGRES = "postgresql"
34SQLITE = "sqlite"
36TABLE_OPTS = {
37 "mysql:engine": "mysql_engine",
38 "mysql:charset": "mysql_charset",
39 "oracle:compress": "oracle_compress",
40}
42COLUMN_VARIANT_OVERRIDE = {
43 "mysql_datatype": "mysql",
44 "oracle_datatype": "oracle",
45 "postgresql_datatype": "postgresql",
46 "sqlite_datatype": "sqlite",
47}
49DIALECT_MODULES = {MYSQL: mysql, ORACLE: oracle, SQLITE: sqlite, POSTGRES: postgresql}
51_length_regex = re.compile(r"\((\d+)\)")
52"""A regular expression that is looking for numbers within parentheses."""
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]
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)
70def make_variant_dict(column_obj: Column) -> dict[str, TypeEngine[Any]]:
71 """Handle variant overrides for a `felis.datamodel.Column`.
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).
77 Parameters
78 ----------
79 column_obj : `felis.datamodel.Column`
80 The column object from which to build the variant dictionary.
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