Coverage for python / lsst / dax / apdb / cassandra / connectionContext.py: 38%
103 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 01:38 -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__ = ["ConnectionContext", "DbVersions"]
26import dataclasses
27import logging
28from collections.abc import Mapping
30# If cassandra-driver is not there the module can still be imported
31# but ApdbCassandra cannot be instantiated.
32try:
33 from cassandra.cluster import Session
34except ImportError:
35 pass
37from ..apdbConfigFreezer import ApdbConfigFreezer
38from ..apdbSchema import ApdbTables
39from ..monitor import MonAgent
40from ..schema_model import Table
41from ..versionTuple import IncompatibleVersionError, VersionTuple
42from .apdbCassandraReplica import ApdbCassandraReplica
43from .apdbCassandraSchema import ApdbCassandraSchema
44from .apdbMetadataCassandra import ApdbMetadataCassandra
45from .cassandra_utils import PreparedStatementCache
46from .config import ApdbCassandraConfig, ApdbCassandraTimePartitionRange
47from .partitioner import Partitioner
49_LOG = logging.getLogger(__name__)
51_MON = MonAgent(__name__)
54@dataclasses.dataclass
55class DbVersions:
56 """Versions defined in APDB metadata table."""
58 schema_version: VersionTuple
59 """Version of the schema from which database was created."""
61 code_version: VersionTuple
62 """Version of ApdbCassandra with which database was created."""
64 replica_version: VersionTuple | None
65 """Version of ApdbCassandraReplica with which database was created, None
66 if replication was not configured.
67 """
70class ConnectionContext:
71 """Container for all kinds ob objects that are instantiated once the
72 connection to Cassandra is established.
74 Parameters
75 ----------
76 session : `cassandra.cluster.Sesion`
77 Cassandra session.
78 config : `ApdbCassandraConfig`
79 Configuration object.
80 table_schemas : `~collection.abc.Mapping` [`ApdbTables`, `Table`]
81 Schema definitions for regular APDB tables.
82 """
84 metadataSchemaVersionKey = "version:schema"
85 """Name of the metadata key to store schema version number."""
87 metadataCodeVersionKey = "version:ApdbCassandra"
88 """Name of the metadata key to store code version number."""
90 metadataReplicaVersionKey = "version:ApdbCassandraReplica"
91 """Name of the metadata key to store replica code version number."""
93 metadataConfigKey = "config:apdb-cassandra.json"
94 """Name of the metadata key to store frozen part of the configuration."""
96 frozen_parameters = (
97 "enable_replica",
98 "ra_dec_columns",
99 "replica_skips_diaobjects",
100 "replica_sub_chunk_count",
101 "partitioning.part_pixelization",
102 "partitioning.part_pix_level",
103 "partitioning.time_partition_tables",
104 "partitioning.time_partition_days",
105 )
106 """Names of the config parameters to be frozen in metadata table."""
108 def __init__(
109 self,
110 session: Session,
111 config: ApdbCassandraConfig,
112 table_schemas: Mapping[ApdbTables, Table],
113 current_versions: DbVersions,
114 ):
115 self.session = session
117 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
118 self.metadata = ApdbMetadataCassandra(
119 self.session, meta_table_name, config.keyspace, "read_tuples", "write"
120 )
122 # Read versions stored in database.
123 self.db_versions = self._readVersions(self.metadata)
124 _LOG.debug("Database versions: %s", self.db_versions)
126 self._versionCheck(current_versions, self.db_versions)
128 # Read frozen config from metadata.
129 config_json = self.metadata.get(self.metadataConfigKey)
130 if config_json is not None:
131 # Update config from metadata.
132 freezer = ApdbConfigFreezer[ApdbCassandraConfig](self.frozen_parameters)
133 self.config = freezer.update(config, config_json)
134 else:
135 self.config = config
136 del config
138 # Since replica version 1.1.0 we use finer partitioning for replica
139 # chunk tables.
140 self.has_chunk_sub_partitions = False
141 self.has_update_record_chunks_table = False
142 if self.config.enable_replica:
143 assert self.db_versions.replica_version is not None, "Replica version must be defined"
144 self.has_chunk_sub_partitions = ApdbCassandraReplica.hasChunkSubPartitions(
145 self.db_versions.replica_version
146 )
147 self.has_update_record_chunks_table = ApdbCassandraReplica.hasUpdateRecordChunks(
148 self.db_versions.replica_version
149 )
151 # Since version 0.1.3 we have metadata for time partitions.
152 self.has_time_partition_meta = self.db_versions.code_version >= VersionTuple(0, 1, 3)
154 # Since version 0.1.2 we have an extra table for visit/detector.
155 self.has_visit_detector_table = self.db_versions.code_version >= VersionTuple(0, 1, 2)
157 # Support for DiaObjectLastToPartition was added at code version 0.1.1
158 # in a backward-compatible way (we only use the table if it is there).
159 self.has_dia_object_last_to_partition = self.db_versions.code_version >= VersionTuple(0, 1, 1)
161 # Cache for prepared statements
162 self.preparer = PreparedStatementCache(self.session)
164 self.partitioner = Partitioner(self.config)
166 self.schema = ApdbCassandraSchema(
167 session=self.session,
168 keyspace=self.config.keyspace,
169 table_schemas=table_schemas,
170 prefix=self.config.prefix,
171 time_partition_tables=self.config.partitioning.time_partition_tables,
172 enable_replica=self.config.enable_replica,
173 replica_skips_diaobjects=self.config.replica_skips_diaobjects,
174 has_chunk_sub_partitions=self.has_chunk_sub_partitions,
175 has_visit_detector_table=self.has_visit_detector_table,
176 )
178 @property
179 def time_partitions_range(self) -> ApdbCassandraTimePartitionRange | None:
180 """Time partition range or None if instance does not use
181 time-partitioned tables (`ApdbCassandraTimePartitionRange` or `None`).
182 """
183 if not self.config.partitioning.time_partition_tables:
184 return None
186 if self.has_time_partition_meta:
187 return ApdbCassandraTimePartitionRange.from_meta(self.metadata)
188 else:
189 # Scan DiaSource tables and see which partitions are present.
190 partitions = set()
191 tables = self.schema.existing_tables(ApdbTables.DiaSource)
192 for table_name in tables[ApdbTables.DiaSource]:
193 _, _, part_str = table_name.rpartition("_")
194 partitions.add(int(part_str))
195 if not partitions:
196 raise LookupError("Failed to find any partitioned DiaSource table.")
197 return ApdbCassandraTimePartitionRange(start=min(partitions), end=max(partitions))
199 @classmethod
200 def _readVersions(cls, metadata: ApdbMetadataCassandra) -> DbVersions:
201 """Read versions of all objects from metadata."""
203 def _get_version(key: str) -> VersionTuple:
204 """Retrieve version number from given metadata key."""
205 version_str = metadata.get(key)
206 if version_str is None:
207 # Should not happen with existing metadata table.
208 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
209 return VersionTuple.fromString(version_str)
211 db_schema_version = _get_version(cls.metadataSchemaVersionKey)
212 db_code_version = _get_version(cls.metadataCodeVersionKey)
214 db_replica_version: VersionTuple | None = None
215 if version_str := metadata.get(cls.metadataReplicaVersionKey):
216 db_replica_version = VersionTuple.fromString(version_str)
218 return DbVersions(
219 schema_version=db_schema_version, code_version=db_code_version, replica_version=db_replica_version
220 )
222 @classmethod
223 def _versionCheck(cls, current_versions: DbVersions, db_versions: DbVersions) -> None:
224 """Check schema version compatibility."""
225 if not current_versions.schema_version.checkCompatibility(db_versions.schema_version):
226 raise IncompatibleVersionError(
227 f"Configured schema version {current_versions.schema_version} "
228 f"is not compatible with database version {db_versions.schema_version}"
229 )
230 if not current_versions.code_version.checkCompatibility(db_versions.code_version):
231 raise IncompatibleVersionError(
232 f"Current code version {current_versions.code_version} "
233 f"is not compatible with database version {db_versions.code_version}"
234 )
236 # Check replica code version only if replica is enabled. Sort of
237 # chicken and egg problem - `enable_replica` is a part of frozen
238 # configuration, but we cannot read frozen configuration until we
239 # validate versions. Assume that if the replica version is present
240 # then replication is enabled.
241 if db_versions.replica_version is not None:
242 assert current_versions.replica_version is not None, "Do not expect None"
243 if not current_versions.replica_version.checkCompatibility(db_versions.replica_version):
244 raise IncompatibleVersionError(
245 f"Current replication code version {current_versions.replica_version} "
246 f"is not compatible with database version {db_versions.replica_version}"
247 )