Coverage for python/lsst/dax/apdb/sql/apdbSqlSchema.py: 17%
192 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-24 01:24 -0800
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(ApdbSchema):
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 engine : `sqlalchemy.engine.Engine`
101 SQLAlchemy engine instance
102 dia_object_index : `str`
103 Indexing mode for DiaObject table, see `ApdbSqlConfig.dia_object_index`
104 for details.
105 htm_index_column : `str`
106 Name of a HTM index column for DiaObject and DiaSource tables.
107 schema_file : `str`
108 Name of the YAML schema file.
109 schema_name : `str`, optional
110 Name of the schema in YAML files.
111 prefix : `str`, optional
112 Prefix to add to all schema elements.
113 namespace : `str`, optional
114 Namespace (or schema name) to use for all APDB tables.
115 enable_replica : `bool`, optional
116 If `True` then use additional tables for replica chunks.
117 """
119 pixel_id_tables = (ApdbTables.DiaObject, ApdbTables.DiaObjectLast, ApdbTables.DiaSource)
120 """Tables that need pixelId column for spatial indexing."""
122 def __init__(
123 self,
124 engine: sqlalchemy.engine.Engine,
125 dia_object_index: str,
126 htm_index_column: str,
127 schema_file: str,
128 schema_name: str = "ApdbSchema",
129 prefix: str = "",
130 namespace: str | None = None,
131 enable_replica: bool = False,
132 ):
133 super().__init__(schema_file, schema_name)
135 self._engine = engine
136 self._dia_object_index = dia_object_index
137 self._htm_index_column = htm_index_column
138 self._prefix = prefix
139 self._enable_replica = enable_replica
141 self._metadata = sqlalchemy.schema.MetaData(schema=namespace)
143 # Add pixelId column and index to tables that need it
144 for table in self.pixel_id_tables:
145 tableDef = self.tableSchemas.get(table)
146 if not tableDef:
147 continue
148 column = schema_model.Column(
149 id=f"#{htm_index_column}",
150 name=htm_index_column,
151 datatype=felis.datamodel.DataType.long,
152 nullable=False,
153 value=None,
154 description="Pixelization index column.",
155 table=tableDef,
156 )
157 tableDef.columns.append(column)
159 # Adjust index if needed
160 if table == ApdbTables.DiaObject and self._dia_object_index == "pix_id_iov":
161 tableDef.primary_key.insert(0, column)
163 if table is ApdbTables.DiaObjectLast:
164 # use it as a leading PK column
165 tableDef.primary_key.insert(0, column)
166 else:
167 # make a regular index
168 name = f"IDX_{tableDef.name}_{htm_index_column}"
169 index = schema_model.Index(id=f"#{name}", name=name, columns=[column])
170 tableDef.indexes.append(index)
172 # generate schema for all tables, must be called last
173 apdb_tables = self._make_apdb_tables()
174 extra_tables = self._make_extra_tables(apdb_tables)
176 converter = ModelToSql(metadata=self._metadata, prefix=self._prefix)
177 id_to_table = converter.make_tables(itertools.chain(apdb_tables.values(), extra_tables.values()))
179 self._apdb_tables = {
180 apdb_enum: id_to_table[table_model.id] for apdb_enum, table_model in apdb_tables.items()
181 }
182 self._extra_tables = {
183 extra_enum: id_to_table[table_model.id] for extra_enum, table_model in extra_tables.items()
184 }
186 self._has_replica_chunks: bool | None = None
187 self._metadata_check: bool | None = None
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:
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)
246 _LOG.info("creating all tables")
247 self._metadata.create_all(self._engine)
249 # Reset possibly cached value.
250 self._has_replica_chunks = None
251 self._metadata_check = None
253 def get_table(self, table_enum: ApdbTables | ExtraTables) -> sqlalchemy.schema.Table:
254 """Return SQLAlchemy table instance for a specified table type/enum.
256 Parameters
257 ----------
258 table_enum : `ApdbTables` or `ExtraTables`
259 Type of table to return.
261 Returns
262 -------
263 table : `sqlalchemy.schema.Table`
264 Table instance.
266 Raises
267 ------
268 ValueError
269 Raised if ``table_enum`` is not valid for this database.
270 """
271 try:
272 if isinstance(table_enum, ApdbTables):
273 if table_enum is ApdbTables.metadata:
274 # There may be cases when schema is configured with the
275 # metadata table but database is still missing it. Check
276 # that table actually exists in the database. Note that
277 # this may interact with `makeSchema`.
278 if self._metadata_check is None:
279 inspector = sqlalchemy.inspect(self._engine)
280 table_name = table_enum.table_name(self._prefix)
281 self._metadata_check = inspector.has_table(table_name, schema=self._metadata.schema)
282 if not self._metadata_check:
283 # this will be caught below
284 raise LookupError("metadata table is missing")
285 return self._apdb_tables[table_enum]
286 else:
287 return self._extra_tables[table_enum]
288 except LookupError:
289 raise ValueError(f"Table type {table_enum} does not exist in the schema") from None
291 def check_column(self, table_enum: ApdbTables | ExtraTables, column: str) -> bool:
292 """Check for the existence of the column in a given table, checking is
293 done against database, not APDB schema.
295 Parameters
296 ----------
297 table_enum : `ApdbTables` or `ExtraTables`
298 Table to check for a column.
299 column : `str`
300 Name of the column to check.
302 Returns
303 -------
304 exists : `bool`
305 True if column exists, False otherwise.
306 """
307 inspector = sqlalchemy.inspect(self._engine)
308 table_name = table_enum.table_name(self._prefix)
309 columns = inspector.get_columns(table_name, schema=self._metadata.schema)
310 for col in columns:
311 if col["name"] == column:
312 return True
313 return False
315 def get_apdb_columns(self, table_enum: ApdbTables | ExtraTables) -> list[sqlalchemy.schema.Column]:
316 """Return list of columns defined for a table in APDB schema.
318 Returned list excludes columns that are implementation-specific, e.g.
319 ``pixelId`` column is not include in the returned list.
321 Parameters
322 ----------
323 table_enum : `ApdbTables` or `ExtraTables`
324 Type of table.
326 Returns
327 -------
328 table : `list` [`sqlalchemy.schema.Column`]
329 Table instance.
331 Raises
332 ------
333 ValueError
334 Raised if ``table_enum`` is not valid for this database.
335 """
336 table = self.get_table(table_enum)
337 exclude_columns = set()
338 if table_enum in self.pixel_id_tables:
339 exclude_columns.add(self._htm_index_column)
340 return [column for column in table.columns if column.name not in exclude_columns]
342 @property
343 def replication_enabled(self) -> bool:
344 """True if replication is enabled (`bool`)."""
345 return self._enable_replica
347 def _make_apdb_tables(self, mysql_engine: str = "InnoDB") -> Mapping[ApdbTables, schema_model.Table]:
348 """Generate schema for regular tables.
350 Parameters
351 ----------
352 mysql_engine : `str`, optional
353 MySQL engine type to use for new tables.
354 """
355 tables: dict[ApdbTables, schema_model.Table] = {}
356 for table_enum in ApdbTables:
357 if table_enum is ApdbTables.DiaObjectLast and self._dia_object_index != "last_object_table":
358 continue
359 if table_enum is ApdbTables.metadata and table_enum not in self.tableSchemas:
360 # Schema does not define metadata.
361 continue
362 if table_enum is ApdbTables.SSSource:
363 # We do not support SSSource table yet.
364 continue
365 table = self.tableSchemas[table_enum]
367 if table_enum is ApdbTables.DiaObjectLast:
368 # In the past DiaObjectLast table did not have validityStart.
369 validity_start_column = "validityStartMjdTai"
370 try:
371 if not self.check_column(ApdbTables.DiaObjectLast, validity_start_column):
372 for column in table.columns:
373 if column.name == validity_start_column:
374 table.columns.remove(column)
375 break
376 except sqlalchemy.exc.NoSuchTableError:
377 # Table does not exist yet, will be created later.
378 pass
380 tables[table_enum] = table
382 return tables
384 def _make_extra_tables(
385 self, apdb_tables: Mapping[ApdbTables, schema_model.Table]
386 ) -> Mapping[ExtraTables, schema_model.Table]:
387 """Generate schema for insert ID tables."""
388 if not self._enable_replica:
389 return {}
391 tables = {}
392 column_defs: list[schema_model.Column] = [
393 schema_model.Column(
394 name="apdb_replica_chunk",
395 id="#ApdbReplicaChunks.apdb_replica_chunk",
396 datatype=felis.datamodel.DataType.long,
397 ),
398 schema_model.Column(
399 name="last_update_time",
400 id="#ApdbReplicaChunks.last_update_time",
401 datatype=felis.datamodel.DataType.timestamp,
402 nullable=False,
403 ),
404 schema_model.Column(
405 name="unique_id",
406 id="#ApdbReplicaChunks.unique_id",
407 datatype=schema_model.ExtraDataTypes.UUID,
408 nullable=False,
409 ),
410 ]
411 parent_table = schema_model.Table(
412 name=ExtraTables.ApdbReplicaChunks.table_name(self._prefix),
413 id="#ApdbReplicaChunks",
414 columns=column_defs,
415 primary_key=[column_defs[0]],
416 constraints=[],
417 indexes=[],
418 )
419 tables[ExtraTables.ApdbReplicaChunks] = parent_table
421 for table_enum, apdb_enum in ExtraTables.replica_chunk_tables().items():
422 apdb_table = apdb_tables[apdb_enum]
423 table_name = table_enum.table_name(self._prefix)
425 columns = self._replicaChunkColumns(table_enum, apdb_enum)
426 column_map = {column.name: column for column in columns}
427 # PK is the same as for original table
428 pk_columns = [column_map[column.name] for column in apdb_table.primary_key]
430 indices = self._replicaChunkIndices(table_enum, column_map)
431 constraints = self._replicaChunkConstraints(table_enum, apdb_table, parent_table, column_map)
432 table = schema_model.Table(
433 name=table_name,
434 id=f"#{table_name}",
435 columns=columns,
436 primary_key=pk_columns,
437 indexes=indices,
438 constraints=constraints,
439 )
440 tables[table_enum] = table
442 # ApdbUpdateRecordChunks table.
443 table_name = ExtraTables.ApdbUpdateRecordChunks.table_name(self._prefix)
444 columns = [
445 schema_model.Column(
446 name="apdb_replica_chunk",
447 id=f"#{table_name}.apdb_replica_chunk",
448 datatype=felis.datamodel.DataType.long,
449 nullable=False,
450 ),
451 schema_model.Column(
452 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_time_ns",
453 name="update_time_ns",
454 datatype=felis.datamodel.DataType.long,
455 nullable=False,
456 ),
457 schema_model.Column(
458 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_order",
459 name="update_order",
460 datatype=felis.datamodel.DataType.int,
461 nullable=False,
462 ),
463 schema_model.Column(
464 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_unique_id",
465 name="update_unique_id",
466 datatype=schema_model.ExtraDataTypes.UUID,
467 nullable=False,
468 ),
469 schema_model.Column(
470 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_payload",
471 name="update_payload",
472 datatype=felis.datamodel.DataType.string,
473 nullable=False,
474 ),
475 ]
476 tables[ExtraTables.ApdbUpdateRecordChunks] = schema_model.Table(
477 name=table_name,
478 id=f"#{table_name}",
479 columns=columns,
480 primary_key=columns[:4],
481 indexes=[],
482 constraints=[],
483 )
485 return tables
487 def _replicaChunkColumns(
488 self, table_enum: ExtraTables, apdb_enum: ApdbTables
489 ) -> list[schema_model.Column]:
490 """Return list of columns for replica chunks tables."""
491 table_name = table_enum.table_name()
492 column_defs: list[schema_model.Column] = [
493 schema_model.Column(
494 name="apdb_replica_chunk",
495 id=f"#{table_name}.apdb_replica_chunk",
496 datatype=felis.datamodel.DataType.long,
497 nullable=False,
498 )
499 ]
500 if table_enum in ExtraTables.replica_chunk_tables():
501 table_model = self.tableSchemas[apdb_enum]
502 column_defs += [column.clone() for column in table_model.primary_key]
503 else:
504 assert False, "Above branches have to cover all enum values"
505 return column_defs
507 def _replicaChunkIndices(
508 self,
509 table_enum: ExtraTables,
510 column_map: Mapping[str, schema_model.Column],
511 ) -> list[schema_model.Index]:
512 """Return set of indices for replica chunk table."""
513 index_defs: list[schema_model.Index] = []
514 if table_enum in ExtraTables.replica_chunk_tables():
515 # Non-unique index on replica chunk column.
516 name = self._prefix + table_enum.name + "_apdb_replica_chunk_idx"
517 column = column_map["apdb_replica_chunk"]
518 index_defs.append(schema_model.Index(name=name, id=f"#{name}", columns=[column]))
519 return index_defs
521 def _replicaChunkConstraints(
522 self,
523 table_enum: ExtraTables,
524 apdb_table: schema_model.Table,
525 parent_table: schema_model.Table,
526 column_map: Mapping[str, schema_model.Column],
527 ) -> list[schema_model.Constraint]:
528 """Return set of constraints for replica chunk table."""
529 constraints: list[schema_model.Constraint] = []
530 replica_chunk_tables = ExtraTables.replica_chunk_tables()
531 if table_enum in replica_chunk_tables:
532 # Foreign key to original table
533 name = f"{table_enum.table_name()}_fk_{apdb_table.name}"
534 other_columns = apdb_table.primary_key
535 this_columns = [column_map[column.name] for column in apdb_table.primary_key]
536 constraints.append(
537 schema_model.ForeignKeyConstraint(
538 name=name,
539 id=f"#{name}",
540 columns=this_columns,
541 referenced_columns=other_columns,
542 onupdate="CASCADE",
543 ondelete="CASCADE",
544 )
545 )
547 # Foreign key to parent chunk ID table
548 name = f"{table_enum.table_name()}_fk_{parent_table.name}"
549 other_columns = parent_table.primary_key
550 this_columns = [column_map[column.name] for column in parent_table.primary_key]
551 constraints.append(
552 schema_model.ForeignKeyConstraint(
553 name=name,
554 id=f"#{name}",
555 columns=this_columns,
556 referenced_columns=other_columns,
557 onupdate="CASCADE",
558 ondelete="CASCADE",
559 )
560 )
561 return constraints