Coverage for python/felis/db/sqltypes.py: 59%
66 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-01 15:16 -0700
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-01 15:16 -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 builtins
23from collections.abc import Mapping
24from typing import Any, Callable
26from sqlalchemy import SmallInteger, types
27from sqlalchemy.dialects import mysql, oracle, postgresql
28from sqlalchemy.ext.compiler import compiles
30MYSQL = "mysql"
31ORACLE = "oracle"
32POSTGRES = "postgresql"
33SQLITE = "sqlite"
36class TINYINT(SmallInteger):
37 """The non-standard TINYINT type."""
39 __visit_name__ = "TINYINT"
42@compiles(TINYINT)
43def compile_tinyint(type_: Any, compiler: Any, **kw: Any) -> str:
44 """Return type name for TINYINT."""
45 return "TINYINT"
48_TypeMap = Mapping[str, types.TypeEngine | type[types.TypeEngine]]
50boolean_map: _TypeMap = {MYSQL: mysql.BIT(1), ORACLE: oracle.NUMBER(1), POSTGRES: postgresql.BOOLEAN()}
52byte_map: _TypeMap = {
53 MYSQL: mysql.TINYINT(),
54 ORACLE: oracle.NUMBER(3),
55 POSTGRES: postgresql.SMALLINT(),
56}
58short_map: _TypeMap = {
59 MYSQL: mysql.SMALLINT(),
60 ORACLE: oracle.NUMBER(5),
61 POSTGRES: postgresql.SMALLINT(),
62}
64# Skip Oracle
65int_map: _TypeMap = {
66 MYSQL: mysql.INTEGER(),
67 POSTGRES: postgresql.INTEGER(),
68}
70long_map: _TypeMap = {
71 MYSQL: mysql.BIGINT(),
72 ORACLE: oracle.NUMBER(38, 0),
73 POSTGRES: postgresql.BIGINT(),
74}
76float_map: _TypeMap = {
77 MYSQL: mysql.FLOAT(),
78 ORACLE: oracle.BINARY_FLOAT(),
79 POSTGRES: postgresql.FLOAT(),
80}
82double_map: _TypeMap = {
83 MYSQL: mysql.DOUBLE(),
84 ORACLE: oracle.BINARY_DOUBLE(),
85 POSTGRES: postgresql.DOUBLE_PRECISION(),
86}
88char_map: _TypeMap = {
89 MYSQL: mysql.CHAR,
90 ORACLE: oracle.CHAR,
91 POSTGRES: postgresql.CHAR,
92}
94string_map: _TypeMap = {
95 MYSQL: mysql.VARCHAR,
96 ORACLE: oracle.VARCHAR2,
97 POSTGRES: postgresql.VARCHAR,
98}
100unicode_map: _TypeMap = {
101 MYSQL: mysql.NVARCHAR,
102 ORACLE: oracle.NVARCHAR2,
103 POSTGRES: postgresql.VARCHAR,
104}
106text_map: _TypeMap = {
107 MYSQL: mysql.LONGTEXT,
108 ORACLE: oracle.CLOB,
109 POSTGRES: postgresql.TEXT,
110}
112binary_map: _TypeMap = {
113 MYSQL: mysql.LONGBLOB,
114 ORACLE: oracle.BLOB,
115 POSTGRES: postgresql.BYTEA,
116}
119def boolean(**kwargs: Any) -> types.TypeEngine:
120 """Return SQLAlchemy type for boolean."""
121 return _vary(types.BOOLEAN(), boolean_map, kwargs)
124def byte(**kwargs: Any) -> types.TypeEngine:
125 """Return SQLAlchemy type for byte."""
126 return _vary(TINYINT(), byte_map, kwargs)
129def short(**kwargs: Any) -> types.TypeEngine:
130 """Return SQLAlchemy type for short integer."""
131 return _vary(types.SMALLINT(), short_map, kwargs)
134def int(**kwargs: Any) -> types.TypeEngine:
135 """Return SQLAlchemy type for integer."""
136 return _vary(types.INTEGER(), int_map, kwargs)
139def long(**kwargs: Any) -> types.TypeEngine:
140 """Return SQLAlchemy type for long integer."""
141 return _vary(types.BIGINT(), long_map, kwargs)
144def float(**kwargs: Any) -> types.TypeEngine:
145 """Return SQLAlchemy type for single precision float."""
146 return _vary(types.FLOAT(), float_map, kwargs)
149def double(**kwargs: Any) -> types.TypeEngine:
150 """Return SQLAlchemy type for double precision float."""
151 return _vary(types.DOUBLE(), double_map, kwargs)
154def char(length: builtins.int, **kwargs: Any) -> types.TypeEngine:
155 """Return SQLAlchemy type for character."""
156 return _vary(types.CHAR(length), char_map, kwargs, length)
159def string(length: builtins.int, **kwargs: Any) -> types.TypeEngine:
160 """Return SQLAlchemy type for string."""
161 return _vary(types.VARCHAR(length), string_map, kwargs, length)
164def unicode(length: builtins.int, **kwargs: Any) -> types.TypeEngine:
165 """Return SQLAlchemy type for unicode string."""
166 return _vary(types.NVARCHAR(length), unicode_map, kwargs, length)
169def text(**kwargs: Any) -> types.TypeEngine:
170 """Return SQLAlchemy type for text."""
171 return _vary(types.TEXT(), text_map, kwargs)
174def binary(length: builtins.int, **kwargs: Any) -> types.TypeEngine:
175 """Return SQLAlchemy type for binary."""
176 return _vary(types.BLOB(length), binary_map, kwargs, length)
179def timestamp(**kwargs: Any) -> types.TypeEngine:
180 """Return SQLAlchemy type for timestamp."""
181 return types.TIMESTAMP()
184def get_type_func(type_name: str) -> Callable:
185 """Return the function for the type with the given name."""
186 if type_name not in globals():
187 raise ValueError(f"Unknown type: {type_name}")
188 return globals()[type_name]
191def _vary(
192 type_: types.TypeEngine,
193 variant_map: _TypeMap,
194 overrides: _TypeMap,
195 *args: Any,
196) -> types.TypeEngine:
197 variants: dict[str, types.TypeEngine | type[types.TypeEngine]] = dict(variant_map)
198 variants.update(overrides)
199 for dialect, variant in variants.items():
200 # If this is a class and not an instance, instantiate
201 if callable(variant):
202 variant = variant(*args)
203 type_ = type_.with_variant(variant, dialect)
204 return type_