Coverage for python/lsst/dax/apdb/schema_model.py: 67%

168 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 02:52 -0700

1# This file is part of dax_apdb. 

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 

22from __future__ import annotations 

23 

24__all__ = [ 

25 "CheckConstraint", 

26 "Column", 

27 "Constraint", 

28 "ExtraDataTypes", 

29 "ForeignKeyConstraint", 

30 "Index", 

31 "Schema", 

32 "Table", 

33 "UniqueConstraint", 

34] 

35 

36import dataclasses 

37from collections.abc import Iterable, Mapping, MutableMapping 

38from enum import Enum 

39from typing import Any 

40 

41import felis.datamodel 

42 

43_Mapping = Mapping[str, Any] 

44 

45 

46class ExtraDataTypes(Enum): 

47 """Additional column data types that we need in dax_apdb.""" 

48 

49 UUID = "uuid" 

50 

51 

52DataTypes = felis.datamodel.DataType | ExtraDataTypes 

53 

54 

55def _strip_keys(map: _Mapping, keys: Iterable[str]) -> _Mapping: 

56 """Return a copy of a dictionary with some keys removed.""" 

57 keys = set(keys) 

58 return {key: value for key, value in map.items() if key not in keys} 

59 

60 

61def _make_iterable(obj: str | Iterable[str]) -> Iterable[str]: 

62 """Make an iterable out of string or list of strings.""" 

63 if isinstance(obj, str): 

64 yield obj 

65 else: 

66 yield from obj 

67 

68 

69@dataclasses.dataclass 

70class Column: 

71 """Column representation in schema.""" 

72 

73 name: str 

74 """Column name.""" 

75 

76 id: str 

77 """Felis ID for this column.""" 

78 

79 datatype: DataTypes 

80 """Column type, one of the enums defined in DataType.""" 

81 

82 length: int | None = None 

83 """Optional length for string/binary columns""" 

84 

85 nullable: bool = True 

86 """True for nullable columns.""" 

87 

88 value: Any = None 

89 """Default value for column, can be `None`.""" 

90 

91 autoincrement: bool | None = None 

92 """Unspecified value results in `None`.""" 

93 

94 description: str | None = None 

95 """Column description.""" 

96 

97 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict) 

98 """Additional annotations for this column.""" 

99 

100 table: Table | None = None 

101 """Table which defines this column, usually not `None`.""" 

102 

103 @classmethod 

104 def from_felis(cls, dm_column: felis.datamodel.Column) -> Column: 

105 """Convert Felis column definition into instance of this class. 

106 

107 Parameters 

108 ---------- 

109 dm_column : `felis.datamodel.Column` 

110 Felis column definition. 

111 

112 Returns 

113 ------- 

114 column : `Column` 

115 Converted column definition. 

116 """ 

117 column = cls( 

118 name=dm_column.name, 

119 id=dm_column.id, 

120 datatype=dm_column.datatype, 

121 length=dm_column.length, 

122 value=dm_column.value, 

123 description=dm_column.description, 

124 nullable=dm_column.nullable if dm_column.nullable is not None else True, 

125 autoincrement=dm_column.autoincrement, 

126 annotations=_strip_keys( 

127 dict(dm_column), 

128 ["name", "id", "datatype", "length", "nullable", "value", "autoincrement", "description"], 

129 ), 

130 ) 

131 return column 

132 

133 def clone(self) -> Column: 

134 """Make a clone of self.""" 

135 return dataclasses.replace(self, table=None) 

136 

137 

138@dataclasses.dataclass 

139class Index: 

140 """Index representation.""" 

141 

142 name: str 

143 """index name, can be empty.""" 

144 

145 id: str 

146 """Felis ID for this index.""" 

147 

148 columns: list[Column] = dataclasses.field(default_factory=list) 

149 """List of columns in index, one of the ``columns`` or ``expressions`` 

150 must be non-empty. 

151 """ 

152 

153 expressions: list[str] = dataclasses.field(default_factory=list) 

154 """List of expressions in index, one of the ``columns`` or ``expressions`` 

155 must be non-empty. 

156 """ 

157 

158 description: str | None = None 

159 """Index description.""" 

160 

161 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict) 

162 """Additional annotations for this index.""" 

163 

164 @classmethod 

165 def from_felis(cls, dm_index: felis.datamodel.Index, columns: Mapping[str, Column]) -> Index: 

166 """Convert Felis index definition into instance of this class. 

167 

168 Parameters 

169 ---------- 

170 dm_index : `felis.datamodel.Index` 

171 Felis index definition. 

172 columns : `~collections.abc.Mapping` [`str`, `Column`] 

173 Mapping of column ID to `Column` instance. 

174 

175 Returns 

176 ------- 

177 index : `Index` 

178 Converted index definition. 

179 """ 

180 return cls( 

181 name=dm_index.name, 

182 id=dm_index.id, 

183 columns=[columns[c] for c in (dm_index.columns or [])], 

184 expressions=dm_index.expressions or [], 

185 description=dm_index.description, 

186 annotations=_strip_keys(dict(dm_index), ["name", "id", "columns", "expressions", "description"]), 

187 ) 

188 

189 

190@dataclasses.dataclass 

191class Constraint: 

192 """Constraint description, this is a base class, actual constraints will be 

193 instances of one of the subclasses. 

194 """ 

195 

196 name: str | None 

197 """Constraint name.""" 

198 

199 id: str 

200 """Felis ID for this constraint.""" 

201 

202 deferrable: bool = False 

203 """If `True` then this constraint will be declared as deferrable.""" 

204 

205 initially: str | None = None 

206 """Value for ``INITIALLY`` clause, only used of ``deferrable`` is True.""" 

207 

208 description: str | None = None 

209 """Constraint description.""" 

210 

211 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict) 

212 """Additional annotations for this constraint.""" 

213 

214 @classmethod 

215 def from_felis(cls, dm_constr: felis.datamodel.Constraint, columns: Mapping[str, Column]) -> Constraint: 

216 """Convert Felis constraint definition into instance of this class. 

217 

218 Parameters 

219 ---------- 

220 dm_const : `felis.datamodel.Constraint` 

221 Felis constraint definition. 

222 columns : `~collections.abc.Mapping` [`str`, `Column`] 

223 Mapping of column ID to `Column` instance. 

224 

225 Returns 

226 ------- 

227 constraint : `Constraint` 

228 Converted constraint definition. 

229 """ 

230 if isinstance(dm_constr, felis.datamodel.UniqueConstraint): 

231 return UniqueConstraint( 

232 name=dm_constr.name, 

233 id=dm_constr.id, 

234 columns=[columns[c] for c in dm_constr.columns], 

235 deferrable=dm_constr.deferrable, 

236 initially=dm_constr.initially, 

237 description=dm_constr.description, 

238 annotations=_strip_keys( 

239 dict(dm_constr), 

240 ["name", "type", "id", "columns", "deferrable", "initially", "description"], 

241 ), 

242 ) 

243 elif isinstance(dm_constr, felis.datamodel.ForeignKeyConstraint): 

244 return ForeignKeyConstraint( 

245 name=dm_constr.name, 

246 id=dm_constr.id, 

247 columns=[columns[c] for c in dm_constr.columns], 

248 referenced_columns=[columns[c] for c in dm_constr.referenced_columns], 

249 deferrable=dm_constr.deferrable, 

250 initially=dm_constr.initially, 

251 description=dm_constr.description, 

252 annotations=_strip_keys( 

253 dict(dm_constr), 

254 [ 

255 "name", 

256 "id", 

257 "type", 

258 "columns", 

259 "deferrable", 

260 "initially", 

261 "referenced_columns", 

262 "description", 

263 ], 

264 ), 

265 ) 

266 elif isinstance(dm_constr, felis.datamodel.CheckConstraint): 

267 return CheckConstraint( 

268 name=dm_constr.name, 

269 id=dm_constr.id, 

270 expression=dm_constr.expression, 

271 deferrable=dm_constr.deferrable, 

272 initially=dm_constr.initially, 

273 description=dm_constr.description, 

274 annotations=_strip_keys( 

275 dict(dm_constr), 

276 ["name", "id", "type", "expression", "deferrable", "initially", "description"], 

277 ), 

278 ) 

279 else: 

280 raise TypeError(f"Unexpected constraint type: {dm_constr.type}") 

281 

282 

283@dataclasses.dataclass 

284class UniqueConstraint(Constraint): 

285 """Description of unique constraint.""" 

286 

287 columns: list[Column] = dataclasses.field(default_factory=list) 

288 """List of columns in this constraint, all columns belong to the same table 

289 as the constraint itself. 

290 """ 

291 

292 

293@dataclasses.dataclass 

294class ForeignKeyConstraint(Constraint): 

295 """Description of foreign key constraint.""" 

296 

297 columns: list[Column] = dataclasses.field(default_factory=list) 

298 """List of columns in this constraint, all columns belong to the same table 

299 as the constraint itself. 

300 """ 

301 

302 referenced_columns: list[Column] = dataclasses.field(default_factory=list) 

303 """List of referenced columns, the number of columns must be the same as in 

304 ``Constraint.columns`` list. All columns must belong to the same table, 

305 which is different from the table of this constraint. 

306 """ 

307 

308 onupdate: str | None = None 

309 """What to do when parent table columns are updated. Typical values are 

310 CASCADE, DELETE and RESTRICT. 

311 """ 

312 

313 ondelete: str | None = None 

314 """What to do when parent table columns are deleted. Typical values are 

315 CASCADE, DELETE and RESTRICT. 

316 """ 

317 

318 @property 

319 def referenced_table(self) -> Table: 

320 """Table referenced by this constraint.""" 

321 assert len(self.referenced_columns) > 0, "column list cannot be empty" 

322 ref_table = self.referenced_columns[0].table 

323 assert ref_table is not None, "foreign key column must have table defined" 

324 return ref_table 

325 

326 

327@dataclasses.dataclass 

328class CheckConstraint(Constraint): 

329 """Description of check constraint.""" 

330 

331 expression: str = "" 

332 """Expression on one or more columns on the table, must be non-empty.""" 

333 

334 

335@dataclasses.dataclass 

336class Table: 

337 """Description of a single table schema.""" 

338 

339 name: str 

340 """Table name.""" 

341 

342 id: str 

343 """Felis ID for this table.""" 

344 

345 columns: list[Column] 

346 """List of Column instances.""" 

347 

348 primary_key: list[Column] 

349 """List of Column that constitute a primary key, may be empty.""" 

350 

351 constraints: list[Constraint] 

352 """List of Constraint instances, can be empty.""" 

353 

354 indexes: list[Index] 

355 """List of Index instances, can be empty.""" 

356 

357 description: str | None = None 

358 """Table description.""" 

359 

360 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict) 

361 """Additional annotations for this table.""" 

362 

363 def __post_init__(self) -> None: 

364 """Update all columns to point to this table.""" 

365 for column in self.columns: 

366 column.table = self 

367 

368 @classmethod 

369 def from_felis(cls, dm_table: felis.datamodel.Table, columns: Mapping[str, Column]) -> Table: 

370 """Convert Felis table definition into instance of this class. 

371 

372 Parameters 

373 ---------- 

374 dm_table : `felis.datamodel.Table` 

375 Felis table definition. 

376 columns : `~collections.abc.Mapping` [`str`, `Column`] 

377 Mapping of column ID to `Column` instance. 

378 

379 Returns 

380 ------- 

381 table : `Table` 

382 Converted table definition. 

383 """ 

384 table_columns = [columns[c.id] for c in dm_table.columns] 

385 if dm_table.primary_key: 

386 pk_columns = [columns[c] for c in _make_iterable(dm_table.primary_key)] 

387 else: 

388 pk_columns = [] 

389 constraints = [Constraint.from_felis(constr, columns) for constr in dm_table.constraints] 

390 indices = [Index.from_felis(dm_idx, columns) for dm_idx in dm_table.indexes] 

391 table = cls( 

392 name=dm_table.name, 

393 id=dm_table.id, 

394 columns=table_columns, 

395 primary_key=pk_columns, 

396 constraints=constraints, 

397 indexes=indices, 

398 description=dm_table.description, 

399 annotations=_strip_keys( 

400 dict(dm_table), 

401 ["name", "id", "columns", "primaryKey", "constraints", "indexes", "description"], 

402 ), 

403 ) 

404 return table 

405 

406 

407@dataclasses.dataclass 

408class Schema: 

409 """Complete schema description, collection of tables.""" 

410 

411 name: str 

412 """Schema name.""" 

413 

414 id: str 

415 """Felis ID for this schema.""" 

416 

417 tables: list[Table] 

418 """Collection of table definitions.""" 

419 

420 version: felis.datamodel.SchemaVersion | None = None 

421 """Schema version description.""" 

422 

423 description: str | None = None 

424 """Schema description.""" 

425 

426 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict) 

427 """Additional annotations for this table.""" 

428 

429 @classmethod 

430 def from_felis(cls, dm_schema: felis.datamodel.Schema) -> Schema: 

431 """Convert felis schema definition to instance of this class. 

432 

433 Parameters 

434 ---------- 

435 dm_schema : `felis.datamodel.Schema` 

436 Felis schema definition. 

437 

438 Returns 

439 ------- 

440 schema : `Schema` 

441 Converted schema definition. 

442 """ 

443 # Convert all columns first. 

444 columns: MutableMapping[str, Column] = {} 

445 for dm_table in dm_schema.tables: 

446 for dm_column in dm_table.columns: 

447 column = Column.from_felis(dm_column) 

448 columns[column.id] = column 

449 

450 tables = [Table.from_felis(dm_table, columns) for dm_table in dm_schema.tables] 

451 

452 version: felis.datamodel.SchemaVersion | None 

453 if isinstance(dm_schema.version, str): 

454 version = felis.datamodel.SchemaVersion(current=dm_schema.version) 

455 else: 

456 version = dm_schema.version 

457 

458 schema = cls( 

459 name=dm_schema.name, 

460 id=dm_schema.id, 

461 tables=tables, 

462 version=version, 

463 description=dm_schema.description, 

464 annotations=_strip_keys(dict(dm_schema), ["name", "id", "tables", "description"]), 

465 ) 

466 return schema