Coverage for python/lsst/dax/apdb/sql/apdbSqlSchema.py: 19%
169 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:35 -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# (http://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 <http://www.gnu.org/licenses/>.
22"""Module responsible for APDB schema operations."""
24from __future__ import annotations
26__all__ = ["ApdbSqlSchema", "ExtraTables"]
28import enum
29import itertools
30import logging
31from collections.abc import Mapping
33import felis.datamodel
34import sqlalchemy
36from .. import schema_model
37from ..apdbSchema import ApdbSchema, ApdbTables
38from .modelToSql import ModelToSql
40_LOG = logging.getLogger(__name__)
43class InconsistentSchemaError(RuntimeError):
44 """Exception raised when schema state is inconsistent."""
47@enum.unique
48class ExtraTables(enum.Enum):
49 """Names of the tables used for tracking insert IDs."""
51 ApdbReplicaChunks = "ApdbReplicaChunks"
52 """Name of the table for replica chunks records."""
54 DiaObjectChunks = "DiaObjectChunks"
55 """Name of the table for DIAObject chunk data."""
57 DiaSourceChunks = "DiaSourceChunks"
58 """Name of the table for DIASource chunk data."""
60 DiaForcedSourceChunks = "DiaForcedSourceChunks"
61 """Name of the table for DIAForcedSource chunk data."""
63 def table_name(self, prefix: str = "") -> str:
64 """Return full table name."""
65 return prefix + self.value
67 @classmethod
68 def replica_chunk_tables(cls) -> Mapping[ExtraTables, ApdbTables]:
69 """Return mapping of tables used for replica chunk storage to their
70 corresponding regular tables.
71 """
72 return {
73 cls.DiaObjectChunks: ApdbTables.DiaObject,
74 cls.DiaSourceChunks: ApdbTables.DiaSource,
75 cls.DiaForcedSourceChunks: ApdbTables.DiaForcedSource,
76 }
79class ApdbSqlSchema(ApdbSchema):
80 """Class for management of APDB schema.
82 Attributes
83 ----------
84 objects : `sqlalchemy.Table`
85 DiaObject table instance
86 objects_last : `sqlalchemy.Table`
87 DiaObjectLast table instance, may be None
88 sources : `sqlalchemy.Table`
89 DiaSource table instance
90 forcedSources : `sqlalchemy.Table`
91 DiaForcedSource table instance
92 replication_enabled : `bool`
93 If true then schema has tables for replication chunks.
95 Parameters
96 ----------
97 engine : `sqlalchemy.engine.Engine`
98 SQLAlchemy engine instance
99 dia_object_index : `str`
100 Indexing mode for DiaObject table, see `ApdbSqlConfig.dia_object_index`
101 for details.
102 htm_index_column : `str`
103 Name of a HTM index column for DiaObject and DiaSource tables.
104 schema_file : `str`
105 Name of the YAML schema file.
106 schema_name : `str`, optional
107 Name of the schema in YAML files.
108 prefix : `str`, optional
109 Prefix to add to all schema elements.
110 namespace : `str`, optional
111 Namespace (or schema name) to use for all APDB tables.
112 enable_replica : `bool`, optional
113 If `True` then use additional tables for replica chunks.
114 """
116 pixel_id_tables = (ApdbTables.DiaObject, ApdbTables.DiaObjectLast, ApdbTables.DiaSource)
117 """Tables that need pixelId column for spatial indexing."""
119 def __init__(
120 self,
121 engine: sqlalchemy.engine.Engine,
122 dia_object_index: str,
123 htm_index_column: str,
124 schema_file: str,
125 schema_name: str = "ApdbSchema",
126 prefix: str = "",
127 namespace: str | None = None,
128 enable_replica: bool = False,
129 ):
130 super().__init__(schema_file, schema_name)
132 self._engine = engine
133 self._dia_object_index = dia_object_index
134 self._htm_index_column = htm_index_column
135 self._prefix = prefix
136 self._enable_replica = enable_replica
138 self._metadata = sqlalchemy.schema.MetaData(schema=namespace)
140 # Add pixelId column and index to tables that need it
141 for table in self.pixel_id_tables:
142 tableDef = self.tableSchemas.get(table)
143 if not tableDef:
144 continue
145 column = schema_model.Column(
146 id=f"#{htm_index_column}",
147 name=htm_index_column,
148 datatype=felis.datamodel.DataType.long,
149 nullable=False,
150 value=None,
151 description="Pixelization index column.",
152 table=tableDef,
153 )
154 tableDef.columns.append(column)
156 # Adjust index if needed
157 if table == ApdbTables.DiaObject and self._dia_object_index == "pix_id_iov":
158 tableDef.primary_key.insert(0, column)
160 if table is ApdbTables.DiaObjectLast:
161 # use it as a leading PK column
162 tableDef.primary_key.insert(0, column)
163 else:
164 # make a regular index
165 name = f"IDX_{tableDef.name}_{htm_index_column}"
166 index = schema_model.Index(id=f"#{name}", name=name, columns=[column])
167 tableDef.indexes.append(index)
169 # generate schema for all tables, must be called last
170 apdb_tables = self._make_apdb_tables()
171 extra_tables = self._make_extra_tables(apdb_tables)
173 converter = ModelToSql(metadata=self._metadata, prefix=self._prefix)
174 id_to_table = converter.make_tables(itertools.chain(apdb_tables.values(), extra_tables.values()))
176 self._apdb_tables = {
177 apdb_enum: id_to_table[table_model.id] for apdb_enum, table_model in apdb_tables.items()
178 }
179 self._extra_tables = {
180 extra_enum: id_to_table[table_model.id] for extra_enum, table_model in extra_tables.items()
181 }
183 self._has_replica_chunks: bool | None = None
184 self._metadata_check: bool | None = None
186 def empty(self) -> bool:
187 """Return True if database schema is empty.
189 Returns
190 -------
191 empty : `bool`
192 `True` if none of the required APDB tables exist in the database,
193 `False` if all required tables exist.
195 Raises
196 ------
197 InconsistentSchemaError
198 Raised when some of the required tables exist but not all.
199 """
200 inspector = sqlalchemy.inspect(self._engine)
201 table_names = set(inspector.get_table_names(self._metadata.schema))
203 existing_tables = []
204 missing_tables = []
205 for table_enum in self._apdb_tables:
206 table_name = table_enum.table_name(self._prefix)
207 if table_name in table_names:
208 existing_tables.append(table_name)
209 else:
210 missing_tables.append(table_name)
212 if not missing_tables:
213 return False
214 elif not existing_tables:
215 return True
216 else:
217 raise InconsistentSchemaError(
218 f"Only some required APDB tables exist: {existing_tables}, missing tables: {missing_tables}"
219 )
221 def makeSchema(self, drop: bool = False) -> None:
222 """Create or re-create all tables.
224 Parameters
225 ----------
226 drop : `bool`, optional
227 If True then drop tables before creating new ones.
228 """
229 # Create namespace if it does not exist yet, for now this only makes
230 # sense for postgres.
231 if self._metadata.schema:
232 dialect = self._engine.dialect
233 quoted_schema = dialect.preparer(dialect).quote_schema(self._metadata.schema)
234 create_schema = sqlalchemy.schema.DDL(
235 "CREATE SCHEMA IF NOT EXISTS %(schema)s", context={"schema": quoted_schema}
236 ).execute_if(dialect="postgresql")
237 sqlalchemy.event.listen(self._metadata, "before_create", create_schema)
239 # create all tables (optionally drop first)
240 if drop:
241 _LOG.info("dropping all tables")
242 self._metadata.drop_all(self._engine)
243 _LOG.info("creating all tables")
244 self._metadata.create_all(self._engine)
246 # Reset possibly cached value.
247 self._has_replica_chunks = None
248 self._metadata_check = None
250 def get_table(self, table_enum: ApdbTables | ExtraTables) -> sqlalchemy.schema.Table:
251 """Return SQLAlchemy table instance for a specified table type/enum.
253 Parameters
254 ----------
255 table_enum : `ApdbTables` or `ExtraTables`
256 Type of table to return.
258 Returns
259 -------
260 table : `sqlalchemy.schema.Table`
261 Table instance.
263 Raises
264 ------
265 ValueError
266 Raised if ``table_enum`` is not valid for this database.
267 """
268 try:
269 if isinstance(table_enum, ApdbTables):
270 if table_enum is ApdbTables.metadata:
271 # There may be cases when schema is configured with the
272 # metadata table but database is still missing it. Check
273 # that table actually exists in the database. Note that
274 # this may interact with `makeSchema`.
275 if self._metadata_check is None:
276 inspector = sqlalchemy.inspect(self._engine)
277 table_name = table_enum.table_name(self._prefix)
278 self._metadata_check = inspector.has_table(table_name, schema=self._metadata.schema)
279 if not self._metadata_check:
280 # this will be caught below
281 raise LookupError("metadata table is missing")
282 return self._apdb_tables[table_enum]
283 else:
284 return self._extra_tables[table_enum]
285 except LookupError:
286 raise ValueError(f"Table type {table_enum} does not exist in the schema") from None
288 def get_apdb_columns(self, table_enum: ApdbTables | ExtraTables) -> list[sqlalchemy.schema.Column]:
289 """Return list of columns defined for a table in APDB schema.
291 Returned list excludes columns that are implementation-specific, e.g.
292 ``pixelId`` column is not include in the returned list.
294 Parameters
295 ----------
296 table_enum : `ApdbTables` or `ExtraTables`
297 Type of table.
299 Returns
300 -------
301 table : `list` [`sqlalchemy.schema.Column`]
302 Table instance.
304 Raises
305 ------
306 ValueError
307 Raised if ``table_enum`` is not valid for this database.
308 """
309 table = self.get_table(table_enum)
310 exclude_columns = set()
311 if table_enum in self.pixel_id_tables:
312 exclude_columns.add(self._htm_index_column)
313 return [column for column in table.columns if column.name not in exclude_columns]
315 @property
316 def replication_enabled(self) -> bool:
317 """True if replication is enabled (`bool`)."""
318 return self._enable_replica
320 def _make_apdb_tables(self, mysql_engine: str = "InnoDB") -> Mapping[ApdbTables, schema_model.Table]:
321 """Generate schema for regular tables.
323 Parameters
324 ----------
325 mysql_engine : `str`, optional
326 MySQL engine type to use for new tables.
327 """
328 tables: dict[ApdbTables, schema_model.Table] = {}
329 for table_enum in ApdbTables:
330 if table_enum is ApdbTables.DiaObjectLast and self._dia_object_index != "last_object_table":
331 continue
332 if table_enum is ApdbTables.metadata and table_enum not in self.tableSchemas:
333 # Schema does not define metadata.
334 continue
335 if table_enum is ApdbTables.SSSource:
336 # We do not support SSSource table yet.
337 continue
338 table = self.tableSchemas[table_enum]
339 tables[table_enum] = table
341 return tables
343 def _make_extra_tables(
344 self, apdb_tables: Mapping[ApdbTables, schema_model.Table]
345 ) -> Mapping[ExtraTables, schema_model.Table]:
346 """Generate schema for insert ID tables."""
347 if not self._enable_replica:
348 return {}
350 tables = {}
351 column_defs: list[schema_model.Column] = [
352 schema_model.Column(
353 name="apdb_replica_chunk",
354 id="#ApdbReplicaChunks.apdb_replica_chunk",
355 datatype=felis.datamodel.DataType.long,
356 ),
357 schema_model.Column(
358 name="last_update_time",
359 id="#ApdbReplicaChunks.last_update_time",
360 datatype=felis.datamodel.DataType.timestamp,
361 nullable=False,
362 ),
363 schema_model.Column(
364 name="unique_id",
365 id="#ApdbReplicaChunks.unique_id",
366 datatype=schema_model.ExtraDataTypes.UUID,
367 nullable=False,
368 ),
369 ]
370 parent_table = schema_model.Table(
371 name=ExtraTables.ApdbReplicaChunks.table_name(self._prefix),
372 id="#ApdbReplicaChunks",
373 columns=column_defs,
374 primary_key=[column_defs[0]],
375 constraints=[],
376 indexes=[],
377 )
378 tables[ExtraTables.ApdbReplicaChunks] = parent_table
380 for table_enum, apdb_enum in ExtraTables.replica_chunk_tables().items():
381 apdb_table = apdb_tables[apdb_enum]
382 table_name = table_enum.table_name(self._prefix)
384 columns = self._replicaChunkColumns(table_enum, apdb_enum)
385 column_map = {column.name: column for column in columns}
386 # PK is the same as for original table
387 pk_columns = [column_map[column.name] for column in apdb_table.primary_key]
389 indices = self._replicaChunkIndices(table_enum, column_map)
390 constraints = self._replicaChunkConstraints(table_enum, apdb_table, parent_table, column_map)
391 table = schema_model.Table(
392 name=table_name,
393 id=f"#{table_name}",
394 columns=columns,
395 primary_key=pk_columns,
396 indexes=indices,
397 constraints=constraints,
398 )
399 tables[table_enum] = table
401 return tables
403 def _replicaChunkColumns(
404 self, table_enum: ExtraTables, apdb_enum: ApdbTables
405 ) -> list[schema_model.Column]:
406 """Return list of columns for replica chunks tables."""
407 table_name = table_enum.table_name()
408 column_defs: list[schema_model.Column] = [
409 schema_model.Column(
410 name="apdb_replica_chunk",
411 id=f"#{table_name}.apdb_replica_chunk",
412 datatype=felis.datamodel.DataType.long,
413 nullable=False,
414 )
415 ]
416 if table_enum in ExtraTables.replica_chunk_tables():
417 table_model = self.tableSchemas[apdb_enum]
418 column_defs += [column.clone() for column in table_model.primary_key]
419 else:
420 assert False, "Above branches have to cover all enum values"
421 return column_defs
423 def _replicaChunkIndices(
424 self,
425 table_enum: ExtraTables,
426 column_map: Mapping[str, schema_model.Column],
427 ) -> list[schema_model.Index]:
428 """Return set of indices for replica chunk table."""
429 index_defs: list[schema_model.Index] = []
430 if table_enum in ExtraTables.replica_chunk_tables():
431 # Non-unique index on replica chunk column.
432 name = self._prefix + table_enum.name + "_apdb_replica_chunk_idx"
433 column = column_map["apdb_replica_chunk"]
434 index_defs.append(schema_model.Index(name=name, id=f"#{name}", columns=[column]))
435 return index_defs
437 def _replicaChunkConstraints(
438 self,
439 table_enum: ExtraTables,
440 apdb_table: schema_model.Table,
441 parent_table: schema_model.Table,
442 column_map: Mapping[str, schema_model.Column],
443 ) -> list[schema_model.Constraint]:
444 """Return set of constraints for replica chunk table."""
445 constraints: list[schema_model.Constraint] = []
446 replica_chunk_tables = ExtraTables.replica_chunk_tables()
447 if table_enum in replica_chunk_tables:
448 # Foreign key to original table
449 name = f"{table_enum.table_name()}_fk_{apdb_table.name}"
450 other_columns = apdb_table.primary_key
451 this_columns = [column_map[column.name] for column in apdb_table.primary_key]
452 constraints.append(
453 schema_model.ForeignKeyConstraint(
454 name=name,
455 id=f"#{name}",
456 columns=this_columns,
457 referenced_columns=other_columns,
458 onupdate="CASCADE",
459 ondelete="CASCADE",
460 )
461 )
463 # Foreign key to parent chunk ID table
464 name = f"{table_enum.table_name()}_fk_{parent_table.name}"
465 other_columns = parent_table.primary_key
466 this_columns = [column_map[column.name] for column in parent_table.primary_key]
467 constraints.append(
468 schema_model.ForeignKeyConstraint(
469 name=name,
470 id=f"#{name}",
471 columns=this_columns,
472 referenced_columns=other_columns,
473 onupdate="CASCADE",
474 ondelete="CASCADE",
475 )
476 )
477 return constraints