Coverage for python/lsst/dax/apdb/cassandra/connectionContext.py: 43%
89 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:35 -0700
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 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, session: Session, config: ApdbCassandraConfig, table_schemas: Mapping[ApdbTables, Table]
110 ):
111 self.session = session
113 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
114 self.metadata = ApdbMetadataCassandra(
115 self.session, meta_table_name, config.keyspace, "read_tuples", "write"
116 )
118 # Read frozen config from metadata.
119 config_json = self.metadata.get(self.metadataConfigKey)
120 if config_json is not None:
121 # Update config from metadata.
122 freezer = ApdbConfigFreezer[ApdbCassandraConfig](self.frozen_parameters)
123 self.config = freezer.update(config, config_json)
124 else:
125 self.config = config
126 del config
128 # Read versions stored in database.
129 self.db_versions = self._readVersions(self.metadata)
130 _LOG.debug("Database versions: %s", self.db_versions)
132 # Since replica version 1.1.0 we use finer partitioning for replica
133 # chunk tables.
134 self.has_chunk_sub_partitions = False
135 if self.config.enable_replica:
136 assert self.db_versions.replica_version is not None, "Replica version must be defined"
137 self.has_chunk_sub_partitions = ApdbCassandraReplica.hasChunkSubPartitions(
138 self.db_versions.replica_version
139 )
141 # Since version 0.1.3 we have metadata for time partitions.
142 self.has_time_partition_meta = self.db_versions.code_version >= VersionTuple(0, 1, 3)
144 # Since version 0.1.2 we have an extra table for visit/detector.
145 self.has_visit_detector_table = self.db_versions.code_version >= VersionTuple(0, 1, 2)
147 # Support for DiaObjectLastToPartition was added at code version 0.1.1
148 # in a backward-compatible way (we only use the table if it is there).
149 self.has_dia_object_last_to_partition = self.db_versions.code_version >= VersionTuple(0, 1, 1)
151 # Cache for prepared statements
152 self.preparer = PreparedStatementCache(self.session)
154 self.partitioner = Partitioner(self.config)
156 self.schema = ApdbCassandraSchema(
157 session=self.session,
158 keyspace=self.config.keyspace,
159 table_schemas=table_schemas,
160 prefix=self.config.prefix,
161 time_partition_tables=self.config.partitioning.time_partition_tables,
162 enable_replica=self.config.enable_replica,
163 replica_skips_diaobjects=self.config.replica_skips_diaobjects,
164 has_chunk_sub_partitions=self.has_chunk_sub_partitions,
165 has_visit_detector_table=self.has_visit_detector_table,
166 )
168 @property
169 def time_partitions_range(self) -> ApdbCassandraTimePartitionRange | None:
170 """Time partition range or None if instance does not use
171 time-partitioned tables (`ApdbCassandraTimePartitionRange` or `None`).
172 """
173 if not self.config.partitioning.time_partition_tables:
174 return None
176 if self.has_time_partition_meta:
177 return ApdbCassandraTimePartitionRange.from_meta(self.metadata)
178 else:
179 # Scan DiaSource tables and see which partitions are present.
180 partitions = set()
181 tables = self.schema.existing_tables(ApdbTables.DiaSource)
182 for table_name in tables[ApdbTables.DiaSource]:
183 _, _, part_str = table_name.rpartition("_")
184 partitions.add(int(part_str))
185 if not partitions:
186 raise LookupError("Failed to find any partitioned DiaSource table.")
187 return ApdbCassandraTimePartitionRange(start=min(partitions), end=max(partitions))
189 def _readVersions(self, metadata: ApdbMetadataCassandra) -> DbVersions:
190 """Read versions of all objects from metadata."""
192 def _get_version(key: str) -> VersionTuple:
193 """Retrieve version number from given metadata key."""
194 version_str = metadata.get(key)
195 if version_str is None:
196 # Should not happen with existing metadata table.
197 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
198 return VersionTuple.fromString(version_str)
200 db_schema_version = _get_version(self.metadataSchemaVersionKey)
201 db_code_version = _get_version(self.metadataCodeVersionKey)
203 # Check replica code version only if replica is enabled.
204 db_replica_version: VersionTuple | None = None
205 if self.config.enable_replica:
206 db_replica_version = _get_version(self.metadataReplicaVersionKey)
208 return DbVersions(
209 schema_version=db_schema_version, code_version=db_code_version, replica_version=db_replica_version
210 )