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

1"""Map Felis types to SQLAlchemy types.""" 

2 

3# This file is part of felis. 

4# 

5# Developed for the LSST Data Management System. 

6# This product includes software developed by the LSST Project 

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

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

9# for details of code ownership. 

10# 

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

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

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

14# (at your option) any later version. 

15# 

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

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

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

19# GNU General Public License for more details. 

20# 

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

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

23 

24from __future__ import annotations 

25 

26import builtins 

27from collections.abc import Callable, Mapping 

28from typing import Any 

29 

30from sqlalchemy import SmallInteger, types 

31from sqlalchemy.dialects import mysql, postgresql 

32from sqlalchemy.ext.compiler import compiles 

33 

34__all__ = [ 

35 "binary", 

36 "boolean", 

37 "byte", 

38 "char", 

39 "double", 

40 "float", 

41 "get_type_func", 

42 "int", 

43 "long", 

44 "short", 

45 "string", 

46 "text", 

47 "timestamp", 

48 "unicode", 

49] 

50 

51MYSQL = "mysql" 

52POSTGRES = "postgresql" 

53SQLITE = "sqlite" 

54 

55 

56class TINYINT(SmallInteger): 

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

58 

59 __visit_name__ = "TINYINT" 

60 

61 

62@compiles(TINYINT) 

63def compile_tinyint(type_: Any, compiler: Any, **kwargs: Any) -> str: 

64 """Compile the non-standard ``TINYINT`` type to SQL. 

65 

66 Parameters 

67 ---------- 

68 type_ 

69 The type object. 

70 compiler 

71 The compiler object. 

72 **kwargs 

73 Additional keyword arguments. 

74 

75 Returns 

76 ------- 

77 `str` 

78 The compiled SQL for TINYINT. 

79 

80 Notes 

81 ----- 

82 This function returns the SQL for the the TINYINT type. The function 

83 signature and parameters are defined by SQLAlchemy. 

84 """ 

85 return "TINYINT" 

86 

87 

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

89 

90boolean_map: _TypeMap = {MYSQL: mysql.BOOLEAN, POSTGRES: postgresql.BOOLEAN()} 

91 

92byte_map: _TypeMap = { 

93 MYSQL: mysql.TINYINT(), 

94 POSTGRES: postgresql.SMALLINT(), 

95} 

96 

97short_map: _TypeMap = { 

98 MYSQL: mysql.SMALLINT(), 

99 POSTGRES: postgresql.SMALLINT(), 

100} 

101 

102int_map: _TypeMap = { 

103 MYSQL: mysql.INTEGER(), 

104 POSTGRES: postgresql.INTEGER(), 

105} 

106 

107long_map: _TypeMap = { 

108 MYSQL: mysql.BIGINT(), 

109 POSTGRES: postgresql.BIGINT(), 

110} 

111 

112float_map: _TypeMap = { 

113 MYSQL: mysql.FLOAT(), 

114 POSTGRES: postgresql.FLOAT(), 

115} 

116 

117double_map: _TypeMap = { 

118 MYSQL: mysql.DOUBLE(), 

119 POSTGRES: postgresql.DOUBLE_PRECISION(), 

120} 

121 

122char_map: _TypeMap = { 

123 MYSQL: mysql.CHAR, 

124 POSTGRES: postgresql.CHAR, 

125} 

126 

127string_map: _TypeMap = { 

128 MYSQL: mysql.VARCHAR, 

129 POSTGRES: postgresql.VARCHAR, 

130} 

131 

132unicode_map: _TypeMap = { 

133 MYSQL: mysql.NVARCHAR, 

134 POSTGRES: postgresql.VARCHAR, 

135} 

136 

137text_map: _TypeMap = { 

138 MYSQL: mysql.LONGTEXT, 

139 POSTGRES: postgresql.TEXT, 

140} 

141 

142binary_map: _TypeMap = { 

143 MYSQL: mysql.LONGBLOB, 

144 POSTGRES: postgresql.BYTEA, 

145} 

146 

147timestamp_map: _TypeMap = { 

148 MYSQL: mysql.DATETIME(timezone=False), 

149 POSTGRES: postgresql.TIMESTAMP(timezone=False), 

150} 

151 

152 

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

154 """Get the SQL type for Felis `~felis.types.Boolean` with variants. 

155 

156 Parameters 

157 ---------- 

158 **kwargs 

159 Additional keyword arguments to pass to the type object. 

160 

161 Returns 

162 ------- 

163 `~sqlalchemy.types.TypeEngine` 

164 The SQL type for a Felis boolean. 

165 """ 

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

167 

168 

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

170 """Get the SQL type for Felis `~felis.types.Byte` with variants. 

171 

172 Parameters 

173 ---------- 

174 **kwargs 

175 Additional keyword arguments to pass to the type object. 

176 

177 Returns 

178 ------- 

179 `~sqlalchemy.types.TypeEngine` 

180 The SQL type for a Felis byte. 

181 """ 

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

183 

184 

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

186 """Get the SQL type for Felis `~felis.types.Short` with variants. 

187 

188 Parameters 

189 ---------- 

190 **kwargs 

191 Additional keyword arguments to pass to the type object. 

192 

193 Returns 

194 ------- 

195 `~sqlalchemy.types.TypeEngine` 

196 The SQL type for a Felis short. 

197 """ 

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

199 

200 

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

202 """Get the SQL type for Felis `~felis.types.Int` with variants. 

203 

204 Parameters 

205 ---------- 

206 **kwargs 

207 Additional keyword arguments to pass to the type object. 

208 

209 Returns 

210 ------- 

211 `~sqlalchemy.types.TypeEngine` 

212 The SQL type for a Felis int. 

213 """ 

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

215 

216 

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

218 """Get the SQL type for Felis `~felis.types.Long` with variants. 

219 

220 Parameters 

221 ---------- 

222 **kwargs 

223 Additional keyword arguments to pass to the type object. 

224 

225 Returns 

226 ------- 

227 `~sqlalchemy.types.TypeEngine` 

228 The SQL type for a Felis long. 

229 """ 

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

231 

232 

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

234 """Get the SQL type for Felis `~felis.types.Float` with variants. 

235 

236 Parameters 

237 ---------- 

238 **kwargs 

239 Additional keyword arguments to pass to the type object. 

240 

241 Returns 

242 ------- 

243 `~sqlalchemy.types.TypeEngine` 

244 The SQL type for a Felis float. 

245 """ 

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

247 

248 

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

250 """Get the SQL type for Felis `~felis.types.Double` with variants. 

251 

252 Parameters 

253 ---------- 

254 **kwargs 

255 Additional keyword arguments to pass to the type object. 

256 

257 Returns 

258 ------- 

259 `~sqlalchemy.types.TypeEngine` 

260 The SQL type for a Felis double. 

261 """ 

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

263 

264 

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

266 """Get the SQL type for Felis `~felis.types.Char` with variants. 

267 

268 Parameters 

269 ---------- 

270 length 

271 The length of the character field. 

272 **kwargs 

273 Additional keyword arguments to pass to the type object. 

274 

275 Returns 

276 ------- 

277 `~sqlalchemy.types.TypeEngine` 

278 The SQL type for a Felis char. 

279 """ 

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

281 

282 

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

284 """Get the SQL type for Felis `~felis.types.String` with variants. 

285 

286 Parameters 

287 ---------- 

288 length 

289 The length of the string field. 

290 **kwargs 

291 Additional keyword arguments to pass to the type object. 

292 

293 Returns 

294 ------- 

295 `~sqlalchemy.types.TypeEngine` 

296 The SQL type for a Felis string. 

297 """ 

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

299 

300 

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

302 """Get the SQL type for Felis `~felis.types.Unicode` with variants. 

303 

304 Parameters 

305 ---------- 

306 length 

307 The length of the unicode string field. 

308 **kwargs 

309 Additional keyword arguments to pass to the type object. 

310 

311 Returns 

312 ------- 

313 `~sqlalchemy.types.TypeEngine` 

314 The SQL type for a Felis unicode string. 

315 """ 

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

317 

318 

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

320 """Get the SQL type for Felis `~felis.types.Text` with variants. 

321 

322 Parameters 

323 ---------- 

324 **kwargs 

325 Additional keyword arguments to pass to the type object. 

326 

327 Returns 

328 ------- 

329 `~sqlalchemy.types.TypeEngine` 

330 The SQL type for Felis text. 

331 """ 

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

333 

334 

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

336 """Get the SQL type for Felis `~felis.types.Binary` with variants. 

337 

338 Parameters 

339 ---------- 

340 length 

341 The length of the binary field. 

342 **kwargs 

343 Additional keyword arguments to pass to the type object. 

344 

345 Returns 

346 ------- 

347 `~sqlalchemy.types.TypeEngine` 

348 The SQL type for Felis binary. 

349 """ 

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

351 

352 

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

354 """Get the SQL type for Felis `~felis.types.Timestamp` with variants. 

355 

356 Parameters 

357 ---------- 

358 **kwargs 

359 Additional keyword arguments to pass to the type object. 

360 

361 Returns 

362 ------- 

363 `~sqlalchemy.types.TypeEngine` 

364 The SQL type for a Felis timestamp. 

365 """ 

366 return _vary(types.TIMESTAMP(timezone=False), timestamp_map, kwargs) 

367 

368 

369def get_type_func(type_name: str) -> Callable: 

370 """Find the function which creates a specific SQL type by its Felis type 

371 name. 

372 

373 Parameters 

374 ---------- 

375 type_name 

376 The name of the type function to get. 

377 

378 Returns 

379 ------- 

380 `Callable` 

381 The function for the type. 

382 

383 Raises 

384 ------ 

385 ValueError 

386 Raised if the type name is not recognized. 

387 

388 Notes 

389 ----- 

390 This maps the type name to the function that creates the SQL type. This is 

391 the main way to get the type functions from the type names. 

392 """ 

393 if type_name not in globals(): 

394 raise ValueError(f"Unknown type: {type_name}") 

395 return globals()[type_name] 

396 

397 

398def _vary( 

399 type_: types.TypeEngine, 

400 variant_map: _TypeMap, 

401 overrides: _TypeMap, 

402 *args: Any, 

403) -> types.TypeEngine: 

404 """Add datatype variants and overrides to a SQLAlchemy type. 

405 

406 Parameters 

407 ---------- 

408 type_ 

409 The base SQLAlchemy type object. This is essentially a default 

410 SQLAlchemy ``TypeEngine`` object, which will apply if there is no 

411 variant or type override from the schema. 

412 variant_map 

413 The dictionary of dialects to types. Each key is a string representing 

414 a dialect name, and each value is either an instance of 

415 ``TypeEngine`` representing the variant type object or a callable 

416 reference to its class type that will be instantiated later. 

417 overrides 

418 The dictionary of dialects to types to override the defaults. Each key 

419 is a string representing a dialect name and type with a similar 

420 structure as the `variant_map`. 

421 args 

422 The extra arguments to pass to the type object. 

423 

424 Notes 

425 ----- 

426 This function is intended for internal use only. It builds a SQLAlchemy 

427 ``TypeEngine`` that includes variants and overrides defined by Felis. 

428 """ 

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

430 variants.update(overrides) 

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

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

433 if callable(variant): 

434 variant = variant(*args) 

435 type_ = type_.with_variant(variant, dialect) 

436 return type_