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

67 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-20 03:38 -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 

25 

26from sqlalchemy import Float, 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 

42class DOUBLE(Float): 

43 """The non-standard DOUBLE type.""" 

44 

45 __visit_name__ = "DOUBLE" 

46 

47 

48@compiles(TINYINT) 

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

50 """Return type name for TINYINT.""" 

51 return "TINYINT" 

52 

53 

54@compiles(DOUBLE) 

55def compile_double(type_: Any, compiler: Any, **kw: Any) -> str: 

56 """Return type name for double precision type.""" 

57 return "DOUBLE" 

58 

59 

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

61 

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

63 

64byte_map: _TypeMap = { 

65 MYSQL: mysql.TINYINT(), 

66 ORACLE: oracle.NUMBER(3), 

67 POSTGRES: postgresql.SMALLINT(), 

68} 

69 

70short_map: _TypeMap = { 

71 MYSQL: mysql.SMALLINT(), 

72 ORACLE: oracle.NUMBER(5), 

73 POSTGRES: postgresql.SMALLINT(), 

74} 

75 

76# Skip Oracle 

77int_map: _TypeMap = { 

78 MYSQL: mysql.INTEGER(), 

79 POSTGRES: postgresql.INTEGER(), 

80} 

81 

82long_map: _TypeMap = { 

83 MYSQL: mysql.BIGINT(), 

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

85 POSTGRES: postgresql.BIGINT(), 

86} 

87 

88float_map: _TypeMap = { 

89 MYSQL: mysql.FLOAT(), 

90 ORACLE: oracle.BINARY_FLOAT(), 

91 POSTGRES: postgresql.FLOAT(), 

92} 

93 

94double_map: _TypeMap = { 

95 MYSQL: mysql.DOUBLE(), 

96 ORACLE: oracle.BINARY_DOUBLE(), 

97 POSTGRES: postgresql.DOUBLE_PRECISION(), 

98} 

99 

100char_map: _TypeMap = { 

101 MYSQL: mysql.CHAR, 

102 ORACLE: oracle.CHAR, 

103 POSTGRES: postgresql.CHAR, 

104} 

105 

106string_map: _TypeMap = { 

107 MYSQL: mysql.VARCHAR, 

108 ORACLE: oracle.VARCHAR2, 

109 POSTGRES: postgresql.VARCHAR, 

110} 

111 

112unicode_map: _TypeMap = { 

113 MYSQL: mysql.NVARCHAR, 

114 ORACLE: oracle.NVARCHAR2, 

115 POSTGRES: postgresql.VARCHAR, 

116} 

117 

118text_map: _TypeMap = { 

119 MYSQL: mysql.LONGTEXT, 

120 ORACLE: oracle.CLOB, 

121 POSTGRES: postgresql.TEXT, 

122} 

123 

124binary_map: _TypeMap = { 

125 MYSQL: mysql.LONGBLOB, 

126 ORACLE: oracle.BLOB, 

127 POSTGRES: postgresql.BYTEA, 

128} 

129 

130 

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

132 """Return SQLAlchemy type for boolean.""" 

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

134 

135 

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

137 """Return SQLAlchemy type for byte.""" 

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

139 

140 

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

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

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

144 

145 

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

147 """Return SQLAlchemy type for integer.""" 

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

149 

150 

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

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

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

154 

155 

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

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

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

159 

160 

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

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

163 return _vary(DOUBLE(), double_map, kwargs) 

164 

165 

166def char(length: builtins.int, **kwargs: Any) -> types.TypeEngine: 

167 """Return SQLAlchemy type for character.""" 

168 return _vary(types.CHAR(length), char_map, kwargs, length) 

169 

170 

171def string(length: builtins.int, **kwargs: Any) -> types.TypeEngine: 

172 """Return SQLAlchemy type for string.""" 

173 return _vary(types.VARCHAR(length), string_map, kwargs, length) 

174 

175 

176def unicode(length: builtins.int, **kwargs: Any) -> types.TypeEngine: 

177 """Return SQLAlchemy type for unicode string.""" 

178 return _vary(types.NVARCHAR(length), unicode_map, kwargs, length) 

179 

180 

181def text(length: builtins.int, **kwargs: Any) -> types.TypeEngine: 

182 """Return SQLAlchemy type for text.""" 

183 return _vary(types.CLOB(length), text_map, kwargs, length) 

184 

185 

186def binary(length: builtins.int, **kwargs: Any) -> types.TypeEngine: 

187 """Return SQLAlchemy type for binary.""" 

188 return _vary(types.BLOB(length), binary_map, kwargs, length) 

189 

190 

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

192 """Return SQLAlchemy type for timestamp.""" 

193 return types.TIMESTAMP() 

194 

195 

196def _vary( 

197 type_: types.TypeEngine, 

198 variant_map: _TypeMap, 

199 overrides: _TypeMap, 

200 *args: Any, 

201) -> types.TypeEngine: 

202 variants: dict[str, types.TypeEngine | type[types.TypeEngine]] = dict(variant_map) 

203 variants.update(overrides) 

204 for dialect, variant in variants.items(): 

205 # If this is a class and not an instance, instantiate 

206 if isinstance(variant, type): 

207 variant = variant(*args) 

208 type_ = type_.with_variant(variant, dialect) 

209 return type_