Coverage for python/felis/db/sqltypes.py: 59%

66 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 builtins 

23from collections.abc import Mapping 

24from typing import Any, Callable 

25 

26from sqlalchemy import SmallInteger, types 

27from sqlalchemy.dialects import mysql, oracle, postgresql 

28from sqlalchemy.ext.compiler import compiles 

29 

30MYSQL = "mysql" 

31ORACLE = "oracle" 

32POSTGRES = "postgresql" 

33SQLITE = "sqlite" 

34 

35 

36class TINYINT(SmallInteger): 

37 """The non-standard TINYINT type.""" 

38 

39 __visit_name__ = "TINYINT" 

40 

41 

42@compiles(TINYINT) 

43def compile_tinyint(type_: Any, compiler: Any, **kw: Any) -> str: 

44 """Return type name for TINYINT.""" 

45 return "TINYINT" 

46 

47 

48_TypeMap = Mapping[str, types.TypeEngine | type[types.TypeEngine]] 

49 

50boolean_map: _TypeMap = {MYSQL: mysql.BOOLEAN, ORACLE: oracle.NUMBER(1), POSTGRES: postgresql.BOOLEAN()} 

51 

52byte_map: _TypeMap = { 

53 MYSQL: mysql.TINYINT(), 

54 ORACLE: oracle.NUMBER(3), 

55 POSTGRES: postgresql.SMALLINT(), 

56} 

57 

58short_map: _TypeMap = { 

59 MYSQL: mysql.SMALLINT(), 

60 ORACLE: oracle.NUMBER(5), 

61 POSTGRES: postgresql.SMALLINT(), 

62} 

63 

64# Skip Oracle 

65int_map: _TypeMap = { 

66 MYSQL: mysql.INTEGER(), 

67 POSTGRES: postgresql.INTEGER(), 

68} 

69 

70long_map: _TypeMap = { 

71 MYSQL: mysql.BIGINT(), 

72 ORACLE: oracle.NUMBER(38, 0), 

73 POSTGRES: postgresql.BIGINT(), 

74} 

75 

76float_map: _TypeMap = { 

77 MYSQL: mysql.FLOAT(), 

78 ORACLE: oracle.BINARY_FLOAT(), 

79 POSTGRES: postgresql.FLOAT(), 

80} 

81 

82double_map: _TypeMap = { 

83 MYSQL: mysql.DOUBLE(), 

84 ORACLE: oracle.BINARY_DOUBLE(), 

85 POSTGRES: postgresql.DOUBLE_PRECISION(), 

86} 

87 

88char_map: _TypeMap = { 

89 MYSQL: mysql.CHAR, 

90 ORACLE: oracle.CHAR, 

91 POSTGRES: postgresql.CHAR, 

92} 

93 

94string_map: _TypeMap = { 

95 MYSQL: mysql.VARCHAR, 

96 ORACLE: oracle.VARCHAR2, 

97 POSTGRES: postgresql.VARCHAR, 

98} 

99 

100unicode_map: _TypeMap = { 

101 MYSQL: mysql.NVARCHAR, 

102 ORACLE: oracle.NVARCHAR2, 

103 POSTGRES: postgresql.VARCHAR, 

104} 

105 

106text_map: _TypeMap = { 

107 MYSQL: mysql.LONGTEXT, 

108 ORACLE: oracle.CLOB, 

109 POSTGRES: postgresql.TEXT, 

110} 

111 

112binary_map: _TypeMap = { 

113 MYSQL: mysql.LONGBLOB, 

114 ORACLE: oracle.BLOB, 

115 POSTGRES: postgresql.BYTEA, 

116} 

117 

118 

119def boolean(**kwargs: Any) -> types.TypeEngine: 

120 """Return SQLAlchemy type for boolean.""" 

121 return _vary(types.BOOLEAN(), boolean_map, kwargs) 

122 

123 

124def byte(**kwargs: Any) -> types.TypeEngine: 

125 """Return SQLAlchemy type for byte.""" 

126 return _vary(TINYINT(), byte_map, kwargs) 

127 

128 

129def short(**kwargs: Any) -> types.TypeEngine: 

130 """Return SQLAlchemy type for short integer.""" 

131 return _vary(types.SMALLINT(), short_map, kwargs) 

132 

133 

134def int(**kwargs: Any) -> types.TypeEngine: 

135 """Return SQLAlchemy type for integer.""" 

136 return _vary(types.INTEGER(), int_map, kwargs) 

137 

138 

139def long(**kwargs: Any) -> types.TypeEngine: 

140 """Return SQLAlchemy type for long integer.""" 

141 return _vary(types.BIGINT(), long_map, kwargs) 

142 

143 

144def float(**kwargs: Any) -> types.TypeEngine: 

145 """Return SQLAlchemy type for single precision float.""" 

146 return _vary(types.FLOAT(), float_map, kwargs) 

147 

148 

149def double(**kwargs: Any) -> types.TypeEngine: 

150 """Return SQLAlchemy type for double precision float.""" 

151 return _vary(types.DOUBLE(), double_map, kwargs) 

152 

153 

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) 

157 

158 

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) 

162 

163 

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) 

167 

168 

169def text(**kwargs: Any) -> types.TypeEngine: 

170 """Return SQLAlchemy type for text.""" 

171 return _vary(types.TEXT(), text_map, kwargs) 

172 

173 

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) 

177 

178 

179def timestamp(**kwargs: Any) -> types.TypeEngine: 

180 """Return SQLAlchemy type for timestamp.""" 

181 return types.TIMESTAMP() 

182 

183 

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] 

189 

190 

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_