Coverage for python / lsst / dax / apdb / cassandra / connectionContext.py: 38%
106 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-12 01:19 -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 metadataDedupKey = "status:deduplication.json"
97 """Name of the metadata key to store code version number."""
99 frozen_parameters = (
100 "enable_replica",
101 "ra_dec_columns",
102 "replica_skips_diaobjects",
103 "replica_sub_chunk_count",
104 "partitioning.part_pixelization",
105 "partitioning.part_pix_level",
106 "partitioning.time_partition_tables",
107 "partitioning.time_partition_days",
108 "partitioning.num_part_dedup",
109 )
110 """Names of the config parameters to be frozen in metadata table."""
112 def __init__(
113 self,
114 session: Session,
115 config: ApdbCassandraConfig,
116 table_schemas: Mapping[ApdbTables, Table],
117 current_versions: DbVersions,
118 ):
119 self.session = session
121 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
122 self.metadata = ApdbMetadataCassandra(
123 self.session, meta_table_name, config.keyspace, "read_tuples", "write"
124 )
126 # Read versions stored in database.
127 self.db_versions = self._readVersions(self.metadata)
128 _LOG.debug("Database versions: %s", self.db_versions)
130 self._versionCheck(current_versions, self.db_versions)
132 # Read frozen config from metadata.
133 config_json = self.metadata.get(self.metadataConfigKey)
134 if config_json is not None:
135 # Update config from metadata.
136 freezer = ApdbConfigFreezer[ApdbCassandraConfig](self.frozen_parameters)
137 self.config = freezer.update(config, config_json)
138 else:
139 self.config = config
140 del config
142 # Since replica version 1.1.0 we use finer partitioning for replica
143 # chunk tables.
144 self.has_chunk_sub_partitions = False
145 self.has_update_record_chunks_table = False
146 if self.config.enable_replica:
147 assert self.db_versions.replica_version is not None, "Replica version must be defined"
148 self.has_chunk_sub_partitions = ApdbCassandraReplica.hasChunkSubPartitions(
149 self.db_versions.replica_version
150 )
151 self.has_update_record_chunks_table = ApdbCassandraReplica.hasUpdateRecordChunks(
152 self.db_versions.replica_version
153 )
155 # Since version 1.3.0 we have metadata for time partitions.
156 self.has_dedup_table = self.db_versions.code_version >= VersionTuple(1, 3, 0)
158 # Since version 0.1.3 we have metadata for time partitions.
159 self.has_time_partition_meta = self.db_versions.code_version >= VersionTuple(0, 1, 3)
161 # Since version 0.1.2 we have an extra table for visit/detector.
162 self.has_visit_detector_table = self.db_versions.code_version >= VersionTuple(0, 1, 2)
164 # Support for DiaObjectLastToPartition was added at code version 0.1.1
165 # in a backward-compatible way (we only use the table if it is there).
166 self.has_dia_object_last_to_partition = self.db_versions.code_version >= VersionTuple(0, 1, 1)
168 # Cache for prepared statements
169 self.preparer = PreparedStatementCache(self.session)
171 self.partitioner = Partitioner(self.config)
173 self.schema = ApdbCassandraSchema(
174 session=self.session,
175 keyspace=self.config.keyspace,
176 table_schemas=table_schemas,
177 prefix=self.config.prefix,
178 time_partition_tables=self.config.partitioning.time_partition_tables,
179 enable_replica=self.config.enable_replica,
180 replica_skips_diaobjects=self.config.replica_skips_diaobjects,
181 has_chunk_sub_partitions=self.has_chunk_sub_partitions,
182 has_visit_detector_table=self.has_visit_detector_table,
183 )
185 @property
186 def time_partitions_range(self) -> ApdbCassandraTimePartitionRange | None:
187 """Time partition range or None if instance does not use
188 time-partitioned tables (`ApdbCassandraTimePartitionRange` or `None`).
189 """
190 if not self.config.partitioning.time_partition_tables:
191 return None
193 if self.has_time_partition_meta:
194 return ApdbCassandraTimePartitionRange.from_meta(self.metadata)
195 else:
196 # Scan DiaSource tables and see which partitions are present.
197 partitions = set()
198 tables = self.schema.existing_tables(ApdbTables.DiaSource)
199 for table_name in tables[ApdbTables.DiaSource]:
200 _, _, part_str = table_name.rpartition("_")
201 partitions.add(int(part_str))
202 if not partitions:
203 raise LookupError("Failed to find any partitioned DiaSource table.")
204 return ApdbCassandraTimePartitionRange(start=min(partitions), end=max(partitions))
206 @classmethod
207 def _readVersions(cls, metadata: ApdbMetadataCassandra) -> DbVersions:
208 """Read versions of all objects from metadata."""
210 def _get_version(key: str) -> VersionTuple:
211 """Retrieve version number from given metadata key."""
212 version_str = metadata.get(key)
213 if version_str is None:
214 # Should not happen with existing metadata table.
215 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
216 return VersionTuple.fromString(version_str)
218 db_schema_version = _get_version(cls.metadataSchemaVersionKey)
219 db_code_version = _get_version(cls.metadataCodeVersionKey)
221 db_replica_version: VersionTuple | None = None
222 if version_str := metadata.get(cls.metadataReplicaVersionKey):
223 db_replica_version = VersionTuple.fromString(version_str)
225 return DbVersions(
226 schema_version=db_schema_version, code_version=db_code_version, replica_version=db_replica_version
227 )
229 @classmethod
230 def _versionCheck(cls, current_versions: DbVersions, db_versions: DbVersions) -> None:
231 """Check schema version compatibility."""
232 if not current_versions.schema_version.checkCompatibility(db_versions.schema_version):
233 raise IncompatibleVersionError(
234 f"Configured schema version {current_versions.schema_version} "
235 f"is not compatible with database version {db_versions.schema_version}"
236 )
237 if not current_versions.code_version.checkCompatibility(db_versions.code_version):
238 raise IncompatibleVersionError(
239 f"Current code version {current_versions.code_version} "
240 f"is not compatible with database version {db_versions.code_version}"
241 )
243 # Check replica code version only if replica is enabled. Sort of
244 # chicken and egg problem - `enable_replica` is a part of frozen
245 # configuration, but we cannot read frozen configuration until we
246 # validate versions. Assume that if the replica version is present
247 # then replication is enabled.
248 if db_versions.replica_version is not None:
249 assert current_versions.replica_version is not None, "Do not expect None"
250 if not current_versions.replica_version.checkCompatibility(db_versions.replica_version):
251 raise IncompatibleVersionError(
252 f"Current replication code version {current_versions.replica_version} "
253 f"is not compatible with database version {db_versions.replica_version}"
254 )