Coverage for python / lsst / dax / apdb / cassandra / apdbCassandraSchema.py: 18%
332 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-28 01:07 -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/>.
22from __future__ import annotations
24__all__ = ["ApdbCassandraSchema", "CreateTableOptions", "TableOptions"]
26import dataclasses
27import enum
28import logging
29from collections.abc import Mapping
30from typing import TYPE_CHECKING, cast
32import felis.datamodel
33import pydantic
35from .. import schema_model
36from ..apdbSchema import ApdbTables
38if TYPE_CHECKING:
39 import cassandra.cluster
41 from ..schema_model import Table
42 from .config import ApdbCassandraTimePartitionRange
45_LOG = logging.getLogger(__name__)
48class InconsistentSchemaError(RuntimeError):
49 """Exception raised when schema state is inconsistent."""
52class TableOptions(pydantic.BaseModel):
53 """Set of per-table options for creating Cassandra tables."""
55 model_config = pydantic.ConfigDict(extra="forbid")
57 tables: list[str]
58 """List of table names for which the options should be applied."""
60 options: str
63class CreateTableOptions(pydantic.BaseModel):
64 """Set of options for creating Cassandra tables."""
66 model_config = pydantic.ConfigDict(extra="forbid")
68 table_options: list[TableOptions] = pydantic.Field(default_factory=list)
69 """Collection of per-table options."""
71 default_table_options: str = ""
72 """Default options used for tables that are not in the above list."""
74 def get_options(self, table_name: str) -> str:
75 """Find table options for a given table name."""
76 for table_options in self.table_options:
77 if table_name in table_options.tables:
78 return table_options.options
79 return self.default_table_options
82@enum.unique
83class ExtraTables(enum.Enum):
84 """Names of the extra tables used by Cassandra implementation.
86 Chunk tables exist in two versions now to support both old and new schema.
87 Eventually we will drop support for old tables.
88 """
90 ApdbReplicaChunks = "ApdbReplicaChunks"
91 """Name of the table for replica chunk records."""
93 DiaObjectChunks = "DiaObjectChunks"
94 """Name of the table for DIAObject chunk data."""
96 DiaSourceChunks = "DiaSourceChunks"
97 """Name of the table for DIASource chunk data."""
99 DiaForcedSourceChunks = "DiaForcedSourceChunks"
100 """Name of the table for DIAForcedSource chunk data."""
102 DiaObjectChunks2 = "DiaObjectChunks2"
103 """Name of the table for DIAObject chunk data."""
105 DiaSourceChunks2 = "DiaSourceChunks2"
106 """Name of the table for DIASource chunk data."""
108 DiaForcedSourceChunks2 = "DiaForcedSourceChunks2"
109 """Name of the table for DIAForcedSource chunk data."""
111 ApdbUpdateRecordChunks = "ApdbUpdateRecordChunks"
112 """Name of the table for ApdbUpdateRecord chunk data."""
114 DiaSourceToPartition = "DiaSourceToPartition"
115 """Maps diaSourceId to its partition values (pixel and time)."""
117 DiaObjectLastToPartition = "DiaObjectLastToPartition"
118 """Maps last diaObjectId version to its partition (pixel)."""
120 ApdbVisitDetector = "ApdbVisitDetector"
121 """Records attempted processing of visit/detector."""
123 DiaObjectDedup = "DiaObjectDedup"
124 """DiaObject records used for deduplication."""
126 def table_name(self, prefix: str = "", time_partition: int | None = None) -> str:
127 """Return full table name.
129 Parameters
130 ----------
131 prefix : `str`, optional
132 Optional prefix for table name.
133 time_partition : `int`, optional
134 Optional time partition, only used for tables that support time
135 patitioning.
136 """
137 return f"{prefix}{self.value}"
139 @classmethod
140 def replica_chunk_tables(cls, has_subchunks: bool) -> Mapping[ApdbTables, ExtraTables]:
141 """Return mapping of APDB tables to corresponding replica chunks
142 tables.
143 """
144 if has_subchunks:
145 return {
146 ApdbTables.DiaObject: cls.DiaObjectChunks2,
147 ApdbTables.DiaSource: cls.DiaSourceChunks2,
148 ApdbTables.DiaForcedSource: cls.DiaForcedSourceChunks2,
149 }
150 else:
151 return {
152 ApdbTables.DiaObject: cls.DiaObjectChunks,
153 ApdbTables.DiaSource: cls.DiaSourceChunks,
154 ApdbTables.DiaForcedSource: cls.DiaForcedSourceChunks,
155 }
158class ApdbCassandraSchema:
159 """Class for management of APDB schema.
161 Parameters
162 ----------
163 session : `cassandra.cluster.Session`
164 Cassandra session object
165 keyspace : `str`
166 Keyspace name for all tables.
167 schema_file : `str`
168 Name of the YAML schema file.
169 schema_name : `str`, optional
170 Name of the schema in YAML files.
171 prefix : `str`, optional
172 Prefix to add to all schema elements.
173 time_partition_tables : `bool`
174 If `True` then schema will have a separate table for each time
175 partition.
176 enable_replica : `bool`, optional
177 If `True` then use additional tables for replica chunks.
178 has_chunk_sub_partitions : `bool`, optional
179 If `True` then replica chunk tables have sub-partition columns. Only
180 used if ``enable_replica`` is `True`.
181 """
183 _type_map = {
184 felis.datamodel.DataType.double: "DOUBLE",
185 felis.datamodel.DataType.float: "FLOAT",
186 felis.datamodel.DataType.timestamp: "TIMESTAMP",
187 felis.datamodel.DataType.long: "BIGINT",
188 felis.datamodel.DataType.int: "INT",
189 felis.datamodel.DataType.short: "SMALLINT",
190 felis.datamodel.DataType.byte: "TINYINT",
191 felis.datamodel.DataType.binary: "BLOB",
192 felis.datamodel.DataType.char: "TEXT",
193 felis.datamodel.DataType.string: "TEXT",
194 felis.datamodel.DataType.unicode: "TEXT",
195 felis.datamodel.DataType.text: "TEXT",
196 felis.datamodel.DataType.boolean: "BOOLEAN",
197 schema_model.ExtraDataTypes.UUID: "UUID",
198 }
199 """Map YAML column types to Cassandra"""
201 _time_partitioned_tables = [
202 ApdbTables.DiaObject,
203 ApdbTables.DiaSource,
204 ApdbTables.DiaForcedSource,
205 ]
206 _spatially_partitioned_tables = [ApdbTables.DiaObjectLast]
208 def __init__(
209 self,
210 session: cassandra.cluster.Session,
211 keyspace: str,
212 table_schemas: Mapping[ApdbTables, Table],
213 prefix: str = "",
214 time_partition_tables: bool = False,
215 enable_replica: bool = False,
216 replica_skips_diaobjects: bool = False,
217 has_chunk_sub_partitions: bool = True,
218 has_visit_detector_table: bool = True,
219 ):
220 self._session = session
221 self._keyspace = keyspace
222 self._table_schemas = table_schemas
223 self._prefix = prefix
224 self._time_partition_tables = time_partition_tables
225 self._enable_replica = enable_replica
226 self._replica_skips_diaobjects = replica_skips_diaobjects
227 self._has_chunk_sub_partitions = has_chunk_sub_partitions
228 self._has_visit_detector_table = has_visit_detector_table
230 self._apdb_tables = self._apdb_tables_schema(time_partition_tables)
231 self._extra_tables = self._extra_tables_schema(self._apdb_tables)
233 def _apdb_tables_schema(self, time_partition_tables: bool) -> Mapping[ApdbTables, schema_model.Table]:
234 """Generate schema for regular APDB tables."""
235 apdb_tables: dict[ApdbTables, schema_model.Table] = {}
237 # add columns and index for partitioning.
238 for table, apdb_table_def in self._table_schemas.items():
239 part_columns = []
240 add_columns = []
241 primary_key = apdb_table_def.primary_key[:]
242 if table in self._spatially_partitioned_tables:
243 # DiaObjectLast does not need temporal partitioning
244 part_columns = ["apdb_part"]
245 add_columns = part_columns
246 elif table in self._time_partitioned_tables:
247 if time_partition_tables:
248 part_columns = ["apdb_part"]
249 else:
250 part_columns = ["apdb_part", "apdb_time_part"]
251 add_columns = part_columns
252 elif table is ApdbTables.SSObject:
253 # For SSObject there is no natural partition key but we have
254 # to partition it because there are too many of them. I'm
255 # going to partition on its primary key (and drop separate
256 # primary key index).
257 part_columns = ["ssObjectId"]
258 primary_key = []
259 elif table is ApdbTables.metadata:
260 # Metadata is in one partition because we want to read all of
261 # it in one query, add an extra column for partition.
262 part_columns = ["meta_part"]
263 add_columns = part_columns
264 else:
265 # TODO: Do not know what to do with the other tables
266 continue
268 column_defs = []
269 if add_columns:
270 column_defs = [
271 schema_model.Column(
272 id=f"#{name}", name=name, datatype=felis.datamodel.DataType.long, nullable=False
273 )
274 for name in add_columns
275 ]
277 if table is ApdbTables.DiaObjectLast:
278 # In the past DiaObjectLast table did not have validityStart.
279 validity_start_column = "validityStartMjdTai"
280 try:
281 if not self.check_column(ApdbTables.DiaObjectLast, validity_start_column):
282 for column in apdb_table_def.columns:
283 if column.name == validity_start_column:
284 apdb_table_def.columns.remove(column)
285 break
286 except LookupError:
287 # Table has not been created yet.
288 pass
290 annotations = dict(apdb_table_def.annotations)
291 annotations["cassandra:apdb_column_names"] = [column.name for column in apdb_table_def.columns]
292 if part_columns:
293 annotations["cassandra:partitioning_columns"] = part_columns
295 apdb_tables[table] = schema_model.Table(
296 id=apdb_table_def.id,
297 name=apdb_table_def.name,
298 columns=column_defs + apdb_table_def.columns,
299 primary_key=primary_key,
300 indexes=[],
301 constraints=[],
302 annotations=annotations,
303 )
305 return apdb_tables
307 def _extra_tables_schema(
308 self, apdb_tables: Mapping[ApdbTables, schema_model.Table]
309 ) -> Mapping[ExtraTables, schema_model.Table]:
310 """Generate schema for extra tables."""
311 extra_tables: dict[ExtraTables, schema_model.Table] = {}
313 if self._has_visit_detector_table:
314 columns = [
315 schema_model.Column(
316 id="#visit",
317 name="visit",
318 datatype=felis.datamodel.DataType.long,
319 nullable=False,
320 ),
321 schema_model.Column(
322 id="#detector",
323 name="detector",
324 datatype=felis.datamodel.DataType.short,
325 nullable=False,
326 ),
327 ]
328 extra_tables[ExtraTables.ApdbVisitDetector] = schema_model.Table(
329 id="#" + ExtraTables.ApdbVisitDetector.value,
330 name=ExtraTables.ApdbVisitDetector.table_name(self._prefix),
331 columns=columns,
332 primary_key=[],
333 indexes=[],
334 constraints=[],
335 annotations={"cassandra:partitioning_columns": ["visit", "detector"]},
336 )
338 # DiaObjectDedup table contains a subset of columns of DiaObject, the
339 # table is used for deduplication, it is partitioned on some random
340 # key. This column list defines non-PK columns in this table, PK
341 # columns from DiaObject are the same as in DiaObject.
342 dedup_column_names = {"ra", "dec", "nDiaSources", "firstDiaSourceMjdTai"}
343 columns = [
344 schema_model.Column(
345 id="#dedup_part",
346 name="dedup_part",
347 datatype=felis.datamodel.DataType.short,
348 nullable=False,
349 )
350 ]
351 primary_keys = []
352 found_column_names = []
353 # Coly column definitions from DiaObject.
354 dia_object_table = apdb_tables[ApdbTables.DiaObject]
355 for column in dia_object_table.columns:
356 if column in dia_object_table.primary_key or column.name in dedup_column_names:
357 cloned_column = dataclasses.replace(column, table=None)
358 columns.append(cloned_column)
359 found_column_names.append(column.name)
360 if column in dia_object_table.primary_key:
361 primary_keys.append(cloned_column)
363 # Check that we found all expected columns.
364 missing_columns = dedup_column_names - set(found_column_names)
365 if missing_columns:
366 raise LookupError(f"Expected columns not found in DiaObject table: {missing_columns}")
368 extra_tables[ExtraTables.DiaObjectDedup] = schema_model.Table(
369 id="#" + ExtraTables.DiaObjectDedup.value,
370 name=ExtraTables.DiaObjectDedup.table_name(self._prefix),
371 columns=columns,
372 primary_key=primary_keys,
373 indexes=[],
374 constraints=[],
375 annotations={
376 "cassandra:partitioning_columns": ["dedup_part"],
377 "cassandra:apdb_column_names": found_column_names,
378 },
379 )
381 # This table maps DiaSource ID to its partitions in DiaSource table and
382 # DiaSourceChunks tables.
383 extra_tables[ExtraTables.DiaSourceToPartition] = schema_model.Table(
384 id="#" + ExtraTables.DiaSourceToPartition.value,
385 name=ExtraTables.DiaSourceToPartition.table_name(self._prefix),
386 columns=[
387 schema_model.Column(
388 id="#diaSourceId",
389 name="diaSourceId",
390 datatype=felis.datamodel.DataType.long,
391 nullable=False,
392 ),
393 schema_model.Column(
394 id="#apdb_part", name="apdb_part", datatype=felis.datamodel.DataType.long, nullable=False
395 ),
396 schema_model.Column(
397 id="#apdb_time_part",
398 name="apdb_time_part",
399 datatype=felis.datamodel.DataType.int,
400 nullable=False,
401 ),
402 schema_model.Column(
403 id="#apdb_replica_chunk",
404 name="apdb_replica_chunk",
405 datatype=felis.datamodel.DataType.long,
406 nullable=True,
407 ),
408 schema_model.Column(
409 id="#apdb_replica_subchunk",
410 name="apdb_replica_subchunk",
411 datatype=felis.datamodel.DataType.int,
412 nullable=True,
413 ),
414 ],
415 primary_key=[],
416 indexes=[],
417 constraints=[],
418 annotations={"cassandra:partitioning_columns": ["diaSourceId"]},
419 )
421 # This table maps diaObjectId to its partition in DiaObjectLast table.
422 extra_tables[ExtraTables.DiaObjectLastToPartition] = schema_model.Table(
423 id="#" + ExtraTables.DiaObjectLastToPartition.value,
424 name=ExtraTables.DiaObjectLastToPartition.table_name(self._prefix),
425 columns=[
426 schema_model.Column(
427 id="#diaObjectId",
428 name="diaObjectId",
429 datatype=felis.datamodel.DataType.long,
430 nullable=False,
431 ),
432 schema_model.Column(
433 id="#apdb_part", name="apdb_part", datatype=felis.datamodel.DataType.long, nullable=False
434 ),
435 ],
436 primary_key=[],
437 indexes=[],
438 constraints=[],
439 annotations={"cassandra:partitioning_columns": ["diaObjectId"]},
440 )
442 if not self._enable_replica:
443 return extra_tables
445 replica_chunk_column = schema_model.Column(
446 id="#apdb_replica_chunk",
447 name="apdb_replica_chunk",
448 datatype=felis.datamodel.DataType.long,
449 nullable=False,
450 )
452 replica_chunk_columns = [replica_chunk_column]
453 if self._has_chunk_sub_partitions:
454 replica_chunk_columns.append(
455 schema_model.Column(
456 id="#apdb_replica_subchunk",
457 name="apdb_replica_subchunk",
458 datatype=felis.datamodel.DataType.int,
459 nullable=False,
460 )
461 )
463 # Table containing replica chunks, this one is not partitioned, but
464 # partition key must be defined.
465 extra_tables[ExtraTables.ApdbReplicaChunks] = schema_model.Table(
466 id="#" + ExtraTables.ApdbReplicaChunks.value,
467 name=ExtraTables.ApdbReplicaChunks.table_name(self._prefix),
468 columns=[
469 schema_model.Column(
470 id="#partition", name="partition", datatype=felis.datamodel.DataType.int, nullable=False
471 ),
472 replica_chunk_column,
473 schema_model.Column(
474 id="#last_update_time",
475 name="last_update_time",
476 datatype=felis.datamodel.DataType.timestamp,
477 nullable=False,
478 ),
479 schema_model.Column(
480 id="#unique_id",
481 name="unique_id",
482 datatype=schema_model.ExtraDataTypes.UUID,
483 nullable=False,
484 ),
485 schema_model.Column(
486 id="#has_subchunks",
487 name="has_subchunks",
488 datatype=felis.datamodel.DataType.boolean,
489 nullable=True,
490 ),
491 ],
492 primary_key=[replica_chunk_column],
493 indexes=[],
494 constraints=[],
495 annotations={"cassandra:partitioning_columns": ["partition"]},
496 )
498 replica_chunk_tables = ExtraTables.replica_chunk_tables(self._has_chunk_sub_partitions)
499 for apdb_table_enum, chunk_table_enum in replica_chunk_tables.items():
500 apdb_table_def = self._table_schemas[apdb_table_enum]
502 extra_tables[chunk_table_enum] = schema_model.Table(
503 id="#" + chunk_table_enum.value,
504 name=chunk_table_enum.table_name(self._prefix),
505 columns=replica_chunk_columns + apdb_table_def.columns,
506 primary_key=apdb_table_def.primary_key[:],
507 indexes=[],
508 constraints=[],
509 annotations={
510 "cassandra:partitioning_columns": [column.name for column in replica_chunk_columns],
511 "cassandra:apdb_column_names": [column.name for column in apdb_table_def.columns],
512 },
513 )
515 # Table with replica chunk data for ApdbUpdateRecord.
516 columns = [
517 schema_model.Column(
518 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_time_ns",
519 name="update_time_ns",
520 datatype=felis.datamodel.DataType.long,
521 nullable=False,
522 ),
523 schema_model.Column(
524 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_order",
525 name="update_order",
526 datatype=felis.datamodel.DataType.int,
527 nullable=False,
528 ),
529 schema_model.Column(
530 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_unique_id",
531 name="update_unique_id",
532 datatype=schema_model.ExtraDataTypes.UUID,
533 nullable=False,
534 ),
535 schema_model.Column(
536 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}.update_payload",
537 name="update_payload",
538 datatype=felis.datamodel.DataType.string,
539 nullable=False,
540 ),
541 ]
542 extra_tables[ExtraTables.ApdbUpdateRecordChunks] = schema_model.Table(
543 id=f"#{ExtraTables.ApdbUpdateRecordChunks.value}",
544 name=ExtraTables.ApdbUpdateRecordChunks.table_name(self._prefix),
545 columns=replica_chunk_columns + columns,
546 primary_key=columns[:3],
547 indexes=[],
548 constraints=[],
549 annotations={
550 "cassandra:partitioning_columns": [column.name for column in replica_chunk_columns],
551 },
552 )
554 return extra_tables
556 @property
557 def replication_enabled(self) -> bool:
558 """True when replication is enabled (`bool`)."""
559 return self._enable_replica
561 def empty(self) -> bool:
562 """Return True if database schema is empty.
564 Returns
565 -------
566 empty : `bool`
567 `True` if none of the required APDB tables exist in the database,
568 `False` if all required tables exist.
570 Raises
571 ------
572 InconsistentSchemaError
573 Raised when some of the required tables exist but not all.
574 """
575 query = "SELECT table_name FROM system_schema.tables WHERE keyspace_name = %s"
576 result = self._session.execute(query, (self._keyspace,))
577 table_names = {row[0] for row in result.all()}
579 existing_tables = []
580 missing_tables = []
581 for table_enum in self._apdb_tables:
582 table_name = table_enum.table_name(self._prefix)
583 if self._time_partition_tables and table_enum in self._time_partitioned_tables:
584 # Check prefix for time-partitioned tables.
585 exists = any(table.startswith(f"{table_name}_") for table in table_names)
586 else:
587 exists = table_name in table_names
588 if exists:
589 existing_tables.append(table_name)
590 else:
591 missing_tables.append(table_name)
593 if not missing_tables:
594 return False
595 elif not existing_tables:
596 return True
597 else:
598 raise InconsistentSchemaError(
599 f"Only some required APDB tables exist: {existing_tables}, missing tables: {missing_tables}"
600 )
602 def existing_tables(self, *args: ApdbTables) -> dict[ApdbTables, list[str]]:
603 """Return the list of existing table names for given table.
605 Parameters
606 ----------
607 *args : `ApdbTables`
608 Tables for which to return their existing table names.
610 Returns
611 -------
612 tables : `dict` [`ApdbTables`, `list`[`str`]]
613 Mapping of the APDB table to the list of the existing table names.
614 More than one name can be present in the list if configuration
615 specifies per-partition tables.
616 """
617 if self._time_partition_tables and not set(args).isdisjoint(self._time_partitioned_tables):
618 # Some of the tables should have per-partition tables.
619 query = "SELECT table_name FROM system_schema.tables WHERE keyspace_name = %s"
620 result = self._session.execute(query, (self._keyspace,))
621 table_names = {row[0] for row in result.all()}
623 tables = {}
624 for table_enum in args:
625 base_name = table_enum.table_name(self._prefix)
626 if table_enum in self._time_partitioned_tables:
627 tables[table_enum] = [table for table in table_names if table.startswith(f"{base_name}_")]
628 else:
629 tables[table_enum] = [base_name]
630 return tables
631 else:
632 # Do not check that they exist, we know that they should.
633 return {table_enum: [table_enum.table_name(self._prefix)] for table_enum in args}
635 def check_column(self, table_enum: ApdbTables | ExtraTables, column: str) -> bool:
636 """Check for the existence of the column in a given table.
638 Parameters
639 ----------
640 table_enum : `ApdbTables` or `ExtraTables`
641 Table to check for a column.
642 column : `str`
643 Name of the column to check.
645 Returns
646 -------
647 exists : `bool`
648 True if column exists, False otherwise.
650 Raises
651 ------
652 LookupError
653 Raised if table does not exist.
654 """
655 if self._time_partition_tables and table_enum in self._time_partitioned_tables:
656 query = (
657 "SELECT table_name FROM system_schema.columns WHERE keyspace_name = %s AND column_name = %s "
658 "ALLOW FILTERING"
659 )
660 result = self._session.execute(query, (self._keyspace, column))
661 base_name = table_enum.table_name(self._prefix)
662 for row in result.all():
663 table_name = row[0]
664 if table_name.startswith(f"{base_name}_"):
665 return True
666 # Check that there is any table with matching name.
667 assert isinstance(table_enum, ApdbTables), "Can only be ApdbTables"
668 tables = self.existing_tables(table_enum)
669 if not tables[table_enum]:
670 raise LookupError(f"Table {base_name} does not exist.")
671 return False
672 else:
673 table_name = table_enum.table_name(self._prefix)
674 query = (
675 "SELECT column_name FROM system_schema.columns WHERE keyspace_name = %s AND table_name = %s"
676 )
677 result = self._session.execute(query, (self._keyspace, table_name))
678 rows = list(result)
679 if not rows:
680 raise LookupError(f"Table {table_name} does not exist.")
681 for row in rows:
682 if row.column_name == column:
683 return True
684 return False
686 def tableName(self, table_name: ApdbTables | ExtraTables, time_partition: int | None = None) -> str:
687 """Return Cassandra table name for APDB table.
689 Parameters
690 ----------
691 table_name : `ApdbTables` or `ExtraTables`
692 Table enum for which to generate table name.
693 time_partition : `int`, optional
694 Optional time partition, only used for tables that support time
695 patitioning.
696 """
697 return table_name.table_name(self._prefix, time_partition)
699 def keyspace(self) -> str:
700 """Return Cassandra keyspace for APDB tables."""
701 return self._keyspace
703 def getColumnMap(self, table_name: ApdbTables | ExtraTables) -> Mapping[str, schema_model.Column]:
704 """Return mapping of column names to Column definitions.
706 Parameters
707 ----------
708 table_name : `ApdbTables`
709 One of known APDB table names.
711 Returns
712 -------
713 column_map : `dict`
714 Mapping of column names to `ColumnDef` instances.
715 """
716 table_schema = self._table_schema(table_name)
717 cmap = {column.name: column for column in table_schema.columns}
718 return cmap
720 def apdbColumnNames(self, table_name: ApdbTables | ExtraTables) -> list[str]:
721 """Return a list of columns names for a table as defined in APDB
722 schema.
724 Parameters
725 ----------
726 table_name : `ApdbTables` or `ExtraTables`
727 Enum for a table in APDB schema.
729 Returns
730 -------
731 columns : `list` of `str`
732 Names of regular columns in the table.
733 """
734 table_schema = self._table_schema(table_name)
735 return table_schema.annotations["cassandra:apdb_column_names"]
737 def partitionColumns(self, table_name: ApdbTables | ExtraTables) -> list[str]:
738 """Return a list of columns used for table partitioning.
740 Parameters
741 ----------
742 table_name : `ApdbTables`
743 Table name in APDB schema
745 Returns
746 -------
747 columns : `list` of `str`
748 Names of columns used for partitioning.
749 """
750 table_schema = self._table_schema(table_name)
751 return table_schema.annotations.get("cassandra:partitioning_columns", [])
753 def clusteringColumns(self, table_name: ApdbTables | ExtraTables) -> list[str]:
754 """Return a list of columns used for clustering.
756 Parameters
757 ----------
758 table_name : `ApdbTables`
759 Table name in APDB schema
761 Returns
762 -------
763 columns : `list` of `str`
764 Names of columns for used for clustering.
765 """
766 table_schema = self._table_schema(table_name)
767 return [column.name for column in table_schema.primary_key]
769 def makeSchema(
770 self,
771 *,
772 drop: bool = False,
773 part_range: ApdbCassandraTimePartitionRange | None = None,
774 replication_factor: int | None = None,
775 table_options: CreateTableOptions | None = None,
776 ) -> None:
777 """Create or re-create all tables.
779 Parameters
780 ----------
781 drop : `bool`
782 If True then drop tables before creating new ones. Note that
783 only tables are dropped and not the whole keyspace.
784 part_range : `ApdbCassandraTimePartitionRange` or `None`
785 Start and end partition number for time partitions. Used to create
786 per-partition DiaObject, DiaSource, and DiaForcedSource tables. If
787 `None` then per-partition tables are not created.
788 replication_factor : `int`, optional
789 Replication factor used when creating new keyspace, if keyspace
790 already exists its replication factor is not changed.
791 """
792 # Try to create keyspace if it does not exist
793 if replication_factor is None:
794 replication_factor = 1
796 # If keyspace exists check its replication factor.
797 query = "SELECT replication FROM system_schema.keyspaces WHERE keyspace_name = %s"
798 result = self._session.execute(query, (self._keyspace,))
799 if row := result.one():
800 # Check replication factor, this depends strategy class.
801 repl_config = cast(Mapping[str, str], row[0])
802 _, _, repl_class = repl_config["class"].rpartition(".")
803 if repl_class == "SimpleStrategy":
804 current_repl = {int(repl_config["replication_factor"])}
805 elif repl_class == "NetworkTopologyStrategy":
806 # There may by multiple datacenters with different replication
807 # factors.
808 current_repl = {
809 int(val) for key, val in repl_config.items() if key != "class" and val.isdecimal()
810 }
811 else:
812 raise ValueError(f"Unexpected replication strategy: {repl_class}")
813 if replication_factor not in current_repl:
814 raise ValueError(
815 f"New replication factor {replication_factor} differs from the replication factor "
816 f"for already existing keyspace: {current_repl}"
817 )
818 else:
819 # Need a new keyspace.
820 query = (
821 f'CREATE KEYSPACE "{self._keyspace}"'
822 " WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': "
823 f"{replication_factor}"
824 "}"
825 )
826 self._session.execute(query)
828 table_options = self._update_table_options(table_options)
829 for table in self._apdb_tables:
830 if table is ApdbTables.DiaObject and self._enable_replica and self._replica_skips_diaobjects:
831 continue
832 if table in (ApdbTables.SSObject, ApdbTables.SSSource):
833 # SSObject/SSSource do not exist in APDB, but are defined in
834 # ApdbTables. The reason is that AP wants to have schema of
835 # these tables for alert-related business.
836 continue
837 self._makeTableSchema(table, drop, part_range, table_options)
838 for extra_table in self._extra_tables:
839 self._makeTableSchema(extra_table, drop, part_range, table_options)
841 def _update_table_options(self, options: CreateTableOptions | None) -> CreateTableOptions | None:
842 """Extend table options with options for internal tables."""
843 # We want to add TTL option to ApdbVisitDetector table.
844 if not self._has_visit_detector_table:
845 return options
847 if not options:
848 options = CreateTableOptions()
850 # set both TTL and gc_grace_seconds to 24h.
851 options.table_options.append(
852 TableOptions(
853 tables=[ExtraTables.ApdbVisitDetector.table_name(self._prefix)],
854 options="default_time_to_live=86400 AND gc_grace_seconds=86400",
855 )
856 )
858 return options
860 def _makeTableSchema(
861 self,
862 table: ApdbTables | ExtraTables,
863 drop: bool = False,
864 part_range: ApdbCassandraTimePartitionRange | None = None,
865 table_options: CreateTableOptions | None = None,
866 ) -> None:
867 _LOG.debug("Making table %s", table)
869 fullTable = table.table_name(self._prefix)
871 table_list = [fullTable]
872 if part_range is not None:
873 if table in self._time_partitioned_tables:
874 table_list = [table.table_name(self._prefix, part) for part in part_range.range()]
876 if drop:
877 queries = [f'DROP TABLE IF EXISTS "{self._keyspace}"."{table_name}"' for table_name in table_list]
878 futures = [self._session.execute_async(query, timeout=None) for query in queries]
879 for future in futures:
880 _LOG.debug("wait for query: %s", future.query)
881 future.result()
882 _LOG.debug("query finished: %s", future.query)
884 queries = []
885 options = table_options.get_options(fullTable).strip() if table_options else None
886 for table_name in table_list:
887 if_not_exists = "" if drop else "IF NOT EXISTS"
888 columns = ", ".join(self._tableColumns(table))
889 query = f'CREATE TABLE {if_not_exists} "{self._keyspace}"."{table_name}" ({columns})'
890 if options:
891 query = f"{query} WITH {options}"
892 _LOG.debug("query: %s", query)
893 queries.append(query)
894 futures = [self._session.execute_async(query, timeout=None) for query in queries]
895 for future in futures:
896 _LOG.debug("wait for query: %s", future.query)
897 future.result()
898 _LOG.debug("query finished: %s", future.query)
900 def _tableColumns(self, table_name: ApdbTables | ExtraTables) -> list[str]:
901 """Return set of columns in a table
903 Parameters
904 ----------
905 table_name : `ApdbTables`
906 Name of the table.
908 Returns
909 -------
910 column_defs : `list`
911 List of strings in the format "column_name type".
912 """
913 table_schema = self._table_schema(table_name)
915 # must have partition columns and clustering columns
916 part_columns = table_schema.annotations.get("cassandra:partitioning_columns", [])
917 clust_columns = [column.name for column in table_schema.primary_key]
918 _LOG.debug("part_columns: %s", part_columns)
919 _LOG.debug("clust_columns: %s", clust_columns)
920 if not part_columns:
921 raise ValueError(f"Table {table_name} configuration is missing partition index")
923 # all columns
924 column_defs = []
925 for column in table_schema.columns:
926 ctype = self._type_map[column.datatype]
927 column_defs.append(f'"{column.name}" {ctype}')
929 # primary key definition
930 part_columns = [f'"{col}"' for col in part_columns]
931 clust_columns = [f'"{col}"' for col in clust_columns]
932 if len(part_columns) > 1:
933 columns = ", ".join(part_columns)
934 part_columns = [f"({columns})"]
935 pkey = ", ".join(part_columns + clust_columns)
936 _LOG.debug("pkey: %s", pkey)
937 column_defs.append(f"PRIMARY KEY ({pkey})")
939 return column_defs
941 def _table_schema(self, table: ApdbTables | ExtraTables) -> schema_model.Table:
942 """Return schema definition for a table."""
943 if isinstance(table, ApdbTables):
944 table_schema = self._apdb_tables[table]
945 else:
946 table_schema = self._extra_tables[table]
947 return table_schema
949 def table_row_size(self, table: ApdbTables | ExtraTables) -> int:
950 """Return an estimate of the row size of a given table.
952 Parameters
953 ----------
954 table : `ApdbTables` or `ExtraTables`
956 Returns
957 -------
958 size : `int`
959 An estimate of a table row size.
961 Notes
962 -----
963 Returned size is not exact. When table has variable-size columns (e.g.
964 strings) may be incorrect. Stored data size or wire-level protocol size
965 can be smaller if some columns are not set or set to NULL.
966 """
967 table_schema = self._table_schema(table)
968 size = sum(column.size() for column in table_schema.columns)
969 return size
971 def time_partitioned_tables(self) -> list[ApdbTables]:
972 """Make the list of time-partitioned tables.
974 Returns
975 -------
976 tables : `list` [`ApdbTables`]
977 Tables the are time-partitioned.
978 """
979 if not self._time_partition_tables:
980 return []
981 has_dia_object_table = not (self._enable_replica and self._replica_skips_diaobjects)
982 tables = [ApdbTables.DiaSource, ApdbTables.DiaForcedSource]
983 if has_dia_object_table:
984 tables.append(ApdbTables.DiaObject)
985 return tables