Coverage for python/lsst/dax/apdb/sql/apdbSqlSchema.py: 83%
201 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 09:06 +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# (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 ApdbUpdateRecordChunks = "ApdbUpdateRecordChunks"
64 """Name of the table for ApdbUpdateRecord chunk data."""
66 def table_name(self, prefix: str = "") -> str:
67 """Return full table name."""
68 return prefix + self.value
70 @classmethod
71 def replica_chunk_tables(cls) -> Mapping[ExtraTables, ApdbTables]:
72 """Return mapping of tables used for replica chunk storage to their
73 corresponding regular tables.
74 """
75 return {
76 cls.DiaObjectChunks: ApdbTables.DiaObject,
77 cls.DiaSourceChunks: ApdbTables.DiaSource,
78 cls.DiaForcedSourceChunks: ApdbTables.DiaForcedSource,
79 }
82class ApdbSqlSchema:
83 """Class for management of APDB schema.
85 Attributes
86 ----------
87 objects : `sqlalchemy.Table`
88 DiaObject table instance
89 objects_last : `sqlalchemy.Table`
90 DiaObjectLast table instance, may be None
91 sources : `sqlalchemy.Table`
92 DiaSource table instance
93 forcedSources : `sqlalchemy.Table`
94 DiaForcedSource table instance
95 replication_enabled : `bool`
96 If true then schema has tables for replication chunks.
98 Parameters
99 ----------
100 table_schema : `ApdbSchema`
101 Basic description of table schema.
102 engine : `sqlalchemy.engine.Engine`
103 SQLAlchemy engine instance
104 dia_object_index : `str`
105 Indexing mode for DiaObject table, see `ApdbSqlConfig.dia_object_index`
106 for details.
107 htm_index_column : `str`
108 Name of a HTM index column for DiaObject and DiaSource tables.
109 prefix : `str`, optional
110 Prefix to add to all schema elements.
111 namespace : `str`, optional
112 Namespace (or schema name) to use for all APDB tables.
113 enable_replica : `bool`, optional
114 If `True` then use additional tables for replica chunks.
115 """
117 pixel_id_tables = (ApdbTables.DiaObject, ApdbTables.DiaObjectLast, ApdbTables.DiaSource)
118 """Tables that need pixelId column for spatial indexing."""
120 def __init__(
121 self,
122 table_schema: ApdbSchema,
123 engine: sqlalchemy.engine.Engine,
124 dia_object_index: str,
125 htm_index_column: str,
126 prefix: str = "",
127 namespace: str | None = None,
128 enable_replica: bool = False,
129 ):
130 self._table_schema = table_schema
131 self._engine = engine
132 self._dia_object_index = dia_object_index
133 self._htm_index_column = htm_index_column
134 self._prefix = prefix
135 self._enable_replica = enable_replica
137 self._metadata = sqlalchemy.schema.MetaData(schema=namespace)
139 # Add pixelId column and index to tables that need it
140 for table in self.pixel_id_tables:
141 tableDef = table_schema.tableSchemas.get(table)
142 if not tableDef: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 continue
144 column = schema_model.Column(
145 id=f"#{htm_index_column}",
146 name=htm_index_column,
147 datatype=felis.datamodel.DataType.long,
148 nullable=False,
149 value=None,
150 description="Pixelization index column.",
151 table=tableDef,
152 )
153 tableDef.columns.append(column)
155 # Adjust index if needed
156 if table == ApdbTables.DiaObject and self._dia_object_index == "pix_id_iov":
157 tableDef.primary_key.insert(0, column)
159 if table is ApdbTables.DiaObjectLast:
160 # use it as a leading PK column
161 tableDef.primary_key.insert(0, column)
162 else:
163 # make a regular index
164 name = f"IDX_{tableDef.name}_{htm_index_column}"
165 index = schema_model.Index(id=f"#{name}", name=name, columns=[column])
166 tableDef.indexes.append(index)
168 # generate schema for all tables, must be called last
169 apdb_tables = self._make_apdb_tables()
170 extra_tables = self._make_extra_tables(apdb_tables)
172 converter = ModelToSql(metadata=self._metadata, prefix=self._prefix)
173 id_to_table = converter.make_tables(itertools.chain(apdb_tables.values(), extra_tables.values()))
175 self._apdb_tables = {
176 apdb_enum: id_to_table[table_model.id] for apdb_enum, table_model in apdb_tables.items()
177 }
178 self._extra_tables = {
179 extra_enum: id_to_table[table_model.id] for extra_enum, table_model in extra_tables.items()
180 }
182 self._has_replica_chunks: bool | None = None
183 self._metadata_check: bool | None = None
185 @property
186 def tableSchemas(self) -> Mapping[ApdbTables, schema_model.Table]:
187 return self._table_schema.tableSchemas
189 def empty(self) -> bool:
190 """Return True if database schema is empty.
192 Returns
193 -------
194 empty : `bool`
195 `True` if none of the required APDB tables exist in the database,
196 `False` if all required tables exist.
198 Raises
199 ------
200 InconsistentSchemaError
201 Raised when some of the required tables exist but not all.
202 """
203 inspector = sqlalchemy.inspect(self._engine)
204 table_names = set(inspector.get_table_names(self._metadata.schema))
206 existing_tables = []
207 missing_tables = []
208 for table_enum in self._apdb_tables:
209 table_name = table_enum.table_name(self._prefix)
210 if table_name in table_names:
211 existing_tables.append(table_name)
212 else:
213 missing_tables.append(table_name)
215 if not missing_tables:
216 return False
217 elif not existing_tables:
218 return True
219 else:
220 raise InconsistentSchemaError(
221 f"Only some required APDB tables exist: {existing_tables}, missing tables: {missing_tables}"
222 )
224 def makeSchema(self, drop: bool = False) -> None:
225 """Create or re-create all tables.
227 Parameters
228 ----------
229 drop : `bool`, optional
230 If True then drop tables before creating new ones.
231 """
232 # Create namespace if it does not exist yet, for now this only makes
233 # sense for postgres.
234 if self._metadata.schema: 234 ↛ 235line 234 didn't jump to line 235 because the condition on line 234 was never true
235 dialect = self._engine.dialect
236 quoted_schema = dialect.preparer(dialect).quote_schema(self._metadata.schema)
237 create_schema = sqlalchemy.schema.DDL(
238 "CREATE SCHEMA IF NOT EXISTS %(schema)s", context={"schema": quoted_schema}
239 ).execute_if(dialect="postgresql")
240 sqlalchemy.event.listen(self._metadata, "before_create", create_schema)
242 # create all tables (optionally drop first)
243 if drop:
244 _LOG.info("dropping all tables")
245 self._metadata.drop_all(self._engine)
247 # If we are dropping existing tables then we also want to drop
248 # `alembic_version` table if it exists.
249 metadata = sqlalchemy.schema.MetaData(schema=self._metadata.schema)
250 try:
251 table = sqlalchemy.schema.Table("alembic_version", metadata, autoload_with=self._engine)
252 except sqlalchemy.exc.NoSuchTableError:
253 # Nothing to delete.
254 pass
255 else:
256 table.drop(self._engine)
258 _LOG.info("creating all tables")
259 self._metadata.create_all(self._engine)
261 # Reset possibly cached value.
262 self._has_replica_chunks = None
263 self._metadata_check = None
265 def get_table(self, table_enum: ApdbTables | ExtraTables) -> sqlalchemy.schema.Table:
266 """Return SQLAlchemy table instance for a specified table type/enum.
268 Parameters
269 ----------
270 table_enum : `ApdbTables` or `ExtraTables`
271 Type of table to return.
273 Returns
274 -------
275 table : `sqlalchemy.schema.Table`
276 Table instance.
278 Raises
279 ------
280 ValueError
281 Raised if ``table_enum`` is not valid for this database.
282 """
283 try:
284 if isinstance(table_enum, ApdbTables):
285 if table_enum is ApdbTables.metadata:
286 # There may be cases when schema is configured with the
287 # metadata table but database is still missing it. Check
288 # that table actually exists in the database. Note that
289 # this may interact with `makeSchema`.
290 if self._metadata_check is None:
291 inspector = sqlalchemy.inspect(self._engine)
292 table_name = table_enum.table_name(self._prefix)
293 self._metadata_check = inspector.has_table(table_name, schema=self._metadata.schema)
294 if not self._metadata_check:
295 # this will be caught below
296 raise LookupError("metadata table is missing")
297 return self._apdb_tables[table_enum]
298 else:
299 return self._extra_tables[table_enum]
300 except LookupError:
301 raise ValueError(f"Table type {table_enum} does not exist in the schema") from None
303 def check_column(self, table_enum: ApdbTables | ExtraTables, column: str) -> bool:
304 """Check for the existence of the column in a given table, checking is
305 done against database, not APDB schema.
307 Parameters
308 ----------
309 table_enum : `ApdbTables` or `ExtraTables`
310 Table to check for a column.
311 column : `str`
312 Name of the column to check.
314 Returns
315 -------
316 exists : `bool`
317 True if column exists, False otherwise.
318 """
319 inspector = sqlalchemy.inspect(self._engine)
320 table_name = table_enum.table_name(self._prefix)
321 columns = inspector.get_columns(table_name, schema=self._metadata.schema)
322 for col in columns: 322 ↛ 325line 322 didn't jump to line 325 because the loop on line 322 didn't complete
323 if col["name"] == column:
324 return True
325 return False
327 def get_apdb_columns(self, table_enum: ApdbTables | ExtraTables) -> list[sqlalchemy.schema.Column]:
328 """Return list of columns defined for a table in APDB schema.
330 Returned list excludes columns that are implementation-specific, e.g.
331 ``pixelId`` column is not include in the returned list.
333 Parameters
334 ----------
335 table_enum : `ApdbTables` or `ExtraTables`
336 Type of table.
338 Returns
339 -------
340 table : `list` [`sqlalchemy.schema.Column`]
341 Table instance.
343 Raises
344 ------
345 ValueError
346 Raised if ``table_enum`` is not valid for this database.
347 """
348 table = self.get_table(table_enum)
349 exclude_columns = set()
350 if table_enum in self.pixel_id_tables:
351 exclude_columns.add(self._htm_index_column)
352 return [column for column in table.columns if column.name not in exclude_columns]
354 @property
355 def replication_enabled(self) -> bool:
356 """True if replication is enabled (`bool`)."""
357 return self._enable_replica
359 def _make_apdb_tables(self, mysql_engine: str = "InnoDB") -> Mapping[ApdbTables, schema_model.Table]:
360 """Generate schema for regular tables.
362 Parameters
363 ----------
364 mysql_engine : `str`, optional
365 MySQL engine type to use for new tables.
366 """
367 tables: dict[ApdbTables, schema_model.Table] = {}
368 for table_enum in ApdbTables:
369 if table_enum is ApdbTables.DiaObjectLast and self._dia_object_index != "last_object_table":
370 continue
371 if table_enum is ApdbTables.metadata and table_enum not in self.tableSchemas:
372 # Schema does not define metadata.
373 continue
374 if table_enum in (ApdbTables.SSObject, ApdbTables.SSSource):
375 # SSObject/SSSource do not exist in APDB, but are defined in
376 # ApdbTables. The reason is that AP wants to have schema of
377 # these tables for alert-related business.
378 continue
379 table = self.tableSchemas[table_enum]
381 if table_enum is ApdbTables.DiaObjectLast:
382 # In the past DiaObjectLast table did not have validityStart.
383 validity_start_column = "validityStartMjdTai"
384 try:
385 if not self.check_column(ApdbTables.DiaObjectLast, validity_start_column): 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true
386 for column in table.columns:
387 if column.name == validity_start_column:
388 table.columns.remove(column)
389 break
390 except sqlalchemy.exc.NoSuchTableError:
391 # Table does not exist yet, will be created later.
392 pass
394 tables[table_enum] = table
396 return tables
398 def _make_extra_tables(
399 self, apdb_tables: Mapping[ApdbTables, schema_model.Table]
400 ) -> Mapping[ExtraTables, schema_model.Table]:
401 """Generate schema for insert ID tables."""
402 if not self._enable_replica:
403 return {}
405 tables = {}
406 column_defs: list[schema_model.Column] = [
407 schema_model.Column(
408 name="apdb_replica_chunk",
409 id="#ApdbReplicaChunks.apdb_replica_chunk",
410 datatype=felis.datamodel.DataType.long,
411 ),
412 schema_model.Column(
413 name="last_update_time",
414 id="#ApdbReplicaChunks.last_update_time",
415 datatype=felis.datamodel.DataType.timestamp,
416 nullable=False,
417 ),
418 schema_model.Column(
419 name="unique_id",
420 id="#ApdbReplicaChunks.unique_id",
421 datatype=schema_model.ExtraDataTypes.UUID,
422 nullable=False,
423 ),
424 ]
425 parent_table = schema_model.Table(
426 name=ExtraTables.ApdbReplicaChunks.table_name(self._prefix),
427 id="#ApdbReplicaChunks",
428 columns=column_defs,
429 primary_key=[column_defs[0]],
430 constraints=[],
431 indexes=[],
432 )
433 tables[ExtraTables.ApdbReplicaChunks] = parent_table
435 for table_enum, apdb_enum in ExtraTables.replica_chunk_tables().items():
436 apdb_table = apdb_tables[apdb_enum]
437 table_name = table_enum.table_name(self._prefix)
439 columns = self._replicaChunkColumns(table_enum, apdb_enum)
440 column_map = {column.name: column for column in columns}
441 # PK is the same as for original table
442 pk_columns = [column_map[column.name] for column in apdb_table.primary_key]
444 indices = self._replicaChunkIndices(table_enum, column_map)
445 constraints = self._replicaChunkConstraints(table_enum, apdb_table, parent_table, column_map)
446 table = schema_model.Table(
447 name=table_name,
448 id=f"#{table_name}",
449 columns=columns,
450 primary_key=pk_columns,
451 indexes=indices,
452 constraints=constraints,
453 )
454 tables[table_enum] = table
456 # ApdbUpdateRecordChunks table.
457 table_name = ExtraTables.ApdbUpdateRecordChunks.table_name(self._prefix)
458 columns = [
459 schema_model.Column(
460 name="apdb_replica_chunk",
461 id=f"#{table_name}.apdb_replica_chunk",
462 datatype=felis.datamodel.DataType.long,
463 nullable=False,
464 ),
465 schema_model.Column(
466 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_time_ns",
467 name="update_time_ns",
468 datatype=felis.datamodel.DataType.long,
469 nullable=False,
470 ),
471 schema_model.Column(
472 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_order",
473 name="update_order",
474 datatype=felis.datamodel.DataType.int,
475 nullable=False,
476 ),
477 schema_model.Column(
478 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_unique_id",
479 name="update_unique_id",
480 datatype=schema_model.ExtraDataTypes.UUID,
481 nullable=False,
482 ),
483 schema_model.Column(
484 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_payload",
485 name="update_payload",
486 datatype=felis.datamodel.DataType.string,
487 nullable=False,
488 ),
489 ]
490 tables[ExtraTables.ApdbUpdateRecordChunks] = schema_model.Table(
491 name=table_name,
492 id=f"#{table_name}",
493 columns=columns,
494 primary_key=columns[:4],
495 indexes=[],
496 constraints=[],
497 )
499 return tables
501 def _replicaChunkColumns(
502 self, table_enum: ExtraTables, apdb_enum: ApdbTables
503 ) -> list[schema_model.Column]:
504 """Return list of columns for replica chunks tables."""
505 table_name = table_enum.table_name()
506 column_defs: list[schema_model.Column] = [
507 schema_model.Column(
508 name="apdb_replica_chunk",
509 id=f"#{table_name}.apdb_replica_chunk",
510 datatype=felis.datamodel.DataType.long,
511 nullable=False,
512 )
513 ]
514 if table_enum in ExtraTables.replica_chunk_tables(): 514 ↛ 518line 514 didn't jump to line 518 because the condition on line 514 was always true
515 table_model = self.tableSchemas[apdb_enum]
516 column_defs += [column.clone() for column in table_model.primary_key]
517 else:
518 assert False, "Above branches have to cover all enum values"
519 return column_defs
521 def _replicaChunkIndices(
522 self,
523 table_enum: ExtraTables,
524 column_map: Mapping[str, schema_model.Column],
525 ) -> list[schema_model.Index]:
526 """Return set of indices for replica chunk table."""
527 index_defs: list[schema_model.Index] = []
528 if table_enum in ExtraTables.replica_chunk_tables(): 528 ↛ 533line 528 didn't jump to line 533 because the condition on line 528 was always true
529 # Non-unique index on replica chunk column.
530 name = self._prefix + table_enum.name + "_apdb_replica_chunk_idx"
531 column = column_map["apdb_replica_chunk"]
532 index_defs.append(schema_model.Index(name=name, id=f"#{name}", columns=[column]))
533 return index_defs
535 def _replicaChunkConstraints(
536 self,
537 table_enum: ExtraTables,
538 apdb_table: schema_model.Table,
539 parent_table: schema_model.Table,
540 column_map: Mapping[str, schema_model.Column],
541 ) -> list[schema_model.Constraint]:
542 """Return set of constraints for replica chunk table."""
543 constraints: list[schema_model.Constraint] = []
544 replica_chunk_tables = ExtraTables.replica_chunk_tables()
545 if table_enum in replica_chunk_tables: 545 ↛ 575line 545 didn't jump to line 575 because the condition on line 545 was always true
546 # Foreign key to original table
547 name = f"{table_enum.table_name()}_fk_{apdb_table.name}"
548 other_columns = apdb_table.primary_key
549 this_columns = [column_map[column.name] for column in apdb_table.primary_key]
550 constraints.append(
551 schema_model.ForeignKeyConstraint(
552 name=name,
553 id=f"#{name}",
554 columns=this_columns,
555 referenced_columns=other_columns,
556 onupdate="CASCADE",
557 ondelete="CASCADE",
558 )
559 )
561 # Foreign key to parent chunk ID table
562 name = f"{table_enum.table_name()}_fk_{parent_table.name}"
563 other_columns = parent_table.primary_key
564 this_columns = [column_map[column.name] for column in parent_table.primary_key]
565 constraints.append(
566 schema_model.ForeignKeyConstraint(
567 name=name,
568 id=f"#{name}",
569 columns=this_columns,
570 referenced_columns=other_columns,
571 onupdate="CASCADE",
572 ondelete="CASCADE",
573 )
574 )
575 return constraints