Coverage for python / lsst / dax / apdb / schema_model.py: 66%
174 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-14 23:46 +0000
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/>.
22from __future__ import annotations
24__all__ = [
25 "CheckConstraint",
26 "Column",
27 "Constraint",
28 "ExtraDataTypes",
29 "ForeignKeyConstraint",
30 "Index",
31 "Schema",
32 "Table",
33 "UniqueConstraint",
34]
36import dataclasses
37from collections.abc import Iterable, Mapping, MutableMapping
38from enum import Enum
39from typing import Any
41import felis.datamodel
43_Mapping = Mapping[str, Any]
46class ExtraDataTypes(Enum):
47 """Additional column data types that we need in dax_apdb."""
49 UUID = "uuid"
52DataTypes = felis.datamodel.DataType | ExtraDataTypes
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}
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
69_data_type_size: Mapping[DataTypes, int] = {
70 felis.datamodel.DataType.boolean: 1,
71 felis.datamodel.DataType.byte: 1,
72 felis.datamodel.DataType.short: 2,
73 felis.datamodel.DataType.int: 4,
74 felis.datamodel.DataType.long: 8,
75 felis.datamodel.DataType.float: 4,
76 felis.datamodel.DataType.double: 8,
77 felis.datamodel.DataType.char: 1,
78 felis.datamodel.DataType.string: 2, # approximation, depends on character set
79 felis.datamodel.DataType.unicode: 2, # approximation, depends on character set
80 felis.datamodel.DataType.text: 2, # approximation, depends on character set
81 felis.datamodel.DataType.binary: 1,
82 felis.datamodel.DataType.timestamp: 8, # May be different depending on backend
83 ExtraDataTypes.UUID: 16,
84}
87@dataclasses.dataclass
88class Column:
89 """Column representation in schema."""
91 name: str
92 """Column name."""
94 id: str
95 """Felis ID for this column."""
97 datatype: DataTypes
98 """Column type, one of the enums defined in DataType."""
100 length: int | None = None
101 """Optional length for string/binary columns"""
103 nullable: bool = True
104 """True for nullable columns."""
106 value: Any = None
107 """Default value for column, can be `None`."""
109 autoincrement: bool | None = None
110 """Unspecified value results in `None`."""
112 description: str | None = None
113 """Column description."""
115 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict)
116 """Additional annotations for this column."""
118 table: Table | None = None
119 """Table which defines this column, usually not `None`."""
121 @classmethod
122 def from_felis(cls, dm_column: felis.datamodel.Column) -> Column:
123 """Convert Felis column definition into instance of this class.
125 Parameters
126 ----------
127 dm_column : `felis.datamodel.Column`
128 Felis column definition.
130 Returns
131 -------
132 column : `Column`
133 Converted column definition.
134 """
135 column = cls(
136 name=dm_column.name,
137 id=dm_column.id,
138 datatype=dm_column.datatype,
139 length=dm_column.length,
140 value=dm_column.value,
141 description=dm_column.description,
142 nullable=dm_column.nullable if dm_column.nullable is not None else True,
143 autoincrement=dm_column.autoincrement,
144 annotations=_strip_keys(
145 dict(dm_column),
146 ["name", "id", "datatype", "length", "nullable", "value", "autoincrement", "description"],
147 ),
148 )
149 return column
151 def clone(self) -> Column:
152 """Make a clone of self."""
153 return dataclasses.replace(self, table=None)
155 def size(self) -> int:
156 """Return size in bytes of this column.
158 Returns
159 -------
160 size : `int`
161 Size in bytes for this column, typically represents in-memory size
162 of the corresponding data type. May or may not be the same as
163 storage size or wire-level protocol size.
164 """
165 size = _data_type_size[self.datatype]
166 if self.length is not None:
167 size *= self.length
168 return size
171@dataclasses.dataclass
172class Index:
173 """Index representation."""
175 name: str
176 """index name, can be empty."""
178 id: str
179 """Felis ID for this index."""
181 columns: list[Column] = dataclasses.field(default_factory=list)
182 """List of columns in index, one of the ``columns`` or ``expressions``
183 must be non-empty.
184 """
186 expressions: list[str] = dataclasses.field(default_factory=list)
187 """List of expressions in index, one of the ``columns`` or ``expressions``
188 must be non-empty.
189 """
191 description: str | None = None
192 """Index description."""
194 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict)
195 """Additional annotations for this index."""
197 @classmethod
198 def from_felis(cls, dm_index: felis.datamodel.Index, columns: Mapping[str, Column]) -> Index:
199 """Convert Felis index definition into instance of this class.
201 Parameters
202 ----------
203 dm_index : `felis.datamodel.Index`
204 Felis index definition.
205 columns : `~collections.abc.Mapping` [`str`, `Column`]
206 Mapping of column ID to `Column` instance.
208 Returns
209 -------
210 index : `Index`
211 Converted index definition.
212 """
213 return cls(
214 name=dm_index.name,
215 id=dm_index.id,
216 columns=[columns[c] for c in (dm_index.columns or [])],
217 expressions=dm_index.expressions or [],
218 description=dm_index.description,
219 annotations=_strip_keys(dict(dm_index), ["name", "id", "columns", "expressions", "description"]),
220 )
223@dataclasses.dataclass
224class Constraint:
225 """Constraint description, this is a base class, actual constraints will be
226 instances of one of the subclasses.
227 """
229 name: str | None
230 """Constraint name."""
232 id: str
233 """Felis ID for this constraint."""
235 deferrable: bool = False
236 """If `True` then this constraint will be declared as deferrable."""
238 initially: str | None = None
239 """Value for ``INITIALLY`` clause, only used of ``deferrable`` is True."""
241 description: str | None = None
242 """Constraint description."""
244 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict)
245 """Additional annotations for this constraint."""
247 @classmethod
248 def from_felis(cls, dm_constr: felis.datamodel.Constraint, columns: Mapping[str, Column]) -> Constraint:
249 """Convert Felis constraint definition into instance of this class.
251 Parameters
252 ----------
253 dm_const : `felis.datamodel.Constraint`
254 Felis constraint definition.
255 columns : `~collections.abc.Mapping` [`str`, `Column`]
256 Mapping of column ID to `Column` instance.
258 Returns
259 -------
260 constraint : `Constraint`
261 Converted constraint definition.
262 """
263 if isinstance(dm_constr, felis.datamodel.UniqueConstraint):
264 return UniqueConstraint(
265 name=dm_constr.name,
266 id=dm_constr.id,
267 columns=[columns[c] for c in dm_constr.columns],
268 deferrable=dm_constr.deferrable,
269 initially=dm_constr.initially,
270 description=dm_constr.description,
271 annotations=_strip_keys(
272 dict(dm_constr),
273 ["name", "type", "id", "columns", "deferrable", "initially", "description"],
274 ),
275 )
276 elif isinstance(dm_constr, felis.datamodel.ForeignKeyConstraint):
277 return ForeignKeyConstraint(
278 name=dm_constr.name,
279 id=dm_constr.id,
280 columns=[columns[c] for c in dm_constr.columns],
281 referenced_columns=[columns[c] for c in dm_constr.referenced_columns],
282 deferrable=dm_constr.deferrable,
283 initially=dm_constr.initially,
284 description=dm_constr.description,
285 annotations=_strip_keys(
286 dict(dm_constr),
287 [
288 "name",
289 "id",
290 "type",
291 "columns",
292 "deferrable",
293 "initially",
294 "referenced_columns",
295 "description",
296 ],
297 ),
298 )
299 elif isinstance(dm_constr, felis.datamodel.CheckConstraint):
300 return CheckConstraint(
301 name=dm_constr.name,
302 id=dm_constr.id,
303 expression=dm_constr.expression,
304 deferrable=dm_constr.deferrable,
305 initially=dm_constr.initially,
306 description=dm_constr.description,
307 annotations=_strip_keys(
308 dict(dm_constr),
309 ["name", "id", "type", "expression", "deferrable", "initially", "description"],
310 ),
311 )
312 else:
313 raise TypeError(f"Unexpected constraint type: {dm_constr}")
316@dataclasses.dataclass
317class UniqueConstraint(Constraint):
318 """Description of unique constraint."""
320 columns: list[Column] = dataclasses.field(default_factory=list)
321 """List of columns in this constraint, all columns belong to the same table
322 as the constraint itself.
323 """
326@dataclasses.dataclass
327class ForeignKeyConstraint(Constraint):
328 """Description of foreign key constraint."""
330 columns: list[Column] = dataclasses.field(default_factory=list)
331 """List of columns in this constraint, all columns belong to the same table
332 as the constraint itself.
333 """
335 referenced_columns: list[Column] = dataclasses.field(default_factory=list)
336 """List of referenced columns, the number of columns must be the same as in
337 ``Constraint.columns`` list. All columns must belong to the same table,
338 which is different from the table of this constraint.
339 """
341 onupdate: str | None = None
342 """What to do when parent table columns are updated. Typical values are
343 CASCADE, DELETE and RESTRICT.
344 """
346 ondelete: str | None = None
347 """What to do when parent table columns are deleted. Typical values are
348 CASCADE, DELETE and RESTRICT.
349 """
351 @property
352 def referenced_table(self) -> Table:
353 """Table referenced by this constraint."""
354 assert len(self.referenced_columns) > 0, "column list cannot be empty"
355 ref_table = self.referenced_columns[0].table
356 assert ref_table is not None, "foreign key column must have table defined"
357 return ref_table
360@dataclasses.dataclass
361class CheckConstraint(Constraint):
362 """Description of check constraint."""
364 expression: str = ""
365 """Expression on one or more columns on the table, must be non-empty."""
368@dataclasses.dataclass
369class Table:
370 """Description of a single table schema."""
372 name: str
373 """Table name."""
375 id: str
376 """Felis ID for this table."""
378 columns: list[Column]
379 """List of Column instances."""
381 primary_key: list[Column]
382 """List of Column that constitute a primary key, may be empty."""
384 constraints: list[Constraint]
385 """List of Constraint instances, can be empty."""
387 indexes: list[Index]
388 """List of Index instances, can be empty."""
390 description: str | None = None
391 """Table description."""
393 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict)
394 """Additional annotations for this table."""
396 def __post_init__(self) -> None:
397 """Update all columns to point to this table."""
398 for column in self.columns:
399 column.table = self
401 @classmethod
402 def from_felis(cls, dm_table: felis.datamodel.Table, columns: Mapping[str, Column]) -> Table:
403 """Convert Felis table definition into instance of this class.
405 Parameters
406 ----------
407 dm_table : `felis.datamodel.Table`
408 Felis table definition.
409 columns : `~collections.abc.Mapping` [`str`, `Column`]
410 Mapping of column ID to `Column` instance.
412 Returns
413 -------
414 table : `Table`
415 Converted table definition.
416 """
417 table_columns = [columns[c.id] for c in dm_table.columns]
418 if dm_table.primary_key:
419 pk_columns = [columns[c] for c in _make_iterable(dm_table.primary_key)]
420 else:
421 pk_columns = []
422 constraints = [Constraint.from_felis(constr, columns) for constr in dm_table.constraints]
423 indices = [Index.from_felis(dm_idx, columns) for dm_idx in dm_table.indexes]
424 table = cls(
425 name=dm_table.name,
426 id=dm_table.id,
427 columns=table_columns,
428 primary_key=pk_columns,
429 constraints=constraints,
430 indexes=indices,
431 description=dm_table.description,
432 annotations=_strip_keys(
433 dict(dm_table),
434 ["name", "id", "columns", "primaryKey", "constraints", "indexes", "description"],
435 ),
436 )
437 return table
440@dataclasses.dataclass
441class Schema:
442 """Complete schema description, collection of tables."""
444 name: str
445 """Schema name."""
447 id: str
448 """Felis ID for this schema."""
450 tables: list[Table]
451 """Collection of table definitions."""
453 version: felis.datamodel.SchemaVersion | None = None
454 """Schema version description."""
456 description: str | None = None
457 """Schema description."""
459 annotations: Mapping[str, Any] = dataclasses.field(default_factory=dict)
460 """Additional annotations for this table."""
462 @classmethod
463 def from_felis(cls, dm_schema: felis.datamodel.Schema) -> Schema:
464 """Convert felis schema definition to instance of this class.
466 Parameters
467 ----------
468 dm_schema : `felis.datamodel.Schema`
469 Felis schema definition.
471 Returns
472 -------
473 schema : `Schema`
474 Converted schema definition.
475 """
476 # Convert all columns first.
477 columns: MutableMapping[str, Column] = {}
478 for dm_table in dm_schema.tables:
479 for dm_column in dm_table.columns:
480 column = Column.from_felis(dm_column)
481 columns[column.id] = column
483 tables = [Table.from_felis(dm_table, columns) for dm_table in dm_schema.tables]
485 version: felis.datamodel.SchemaVersion | None
486 if isinstance(dm_schema.version, str):
487 version = felis.datamodel.SchemaVersion(current=dm_schema.version)
488 else:
489 version = dm_schema.version
491 schema = cls(
492 name=dm_schema.name,
493 id=dm_schema.id,
494 tables=tables,
495 version=version,
496 description=dm_schema.description,
497 annotations=_strip_keys(dict(dm_schema), ["name", "id", "tables", "description"]),
498 )
499 return schema