Coverage for python/lsst/dax/apdb/cassandra/sessionFactory.py: 24%
102 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-07 01:39 -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__ = ["SessionContext", "SessionFactory"]
26import logging
27from collections.abc import Mapping
28from contextlib import ExitStack
29from typing import TYPE_CHECKING, Any
31# If cassandra-driver is not there the module can still be imported
32# but ApdbCassandra cannot be instantiated.
33try:
34 import cassandra
35 import cassandra.query
36 from cassandra.auth import AuthProvider, PlainTextAuthProvider
37 from cassandra.cluster import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile, Session
38 from cassandra.policies import AddressTranslator, RoundRobinPolicy, WhiteListRoundRobinPolicy
40 CASSANDRA_IMPORTED = True
41except ImportError:
42 CASSANDRA_IMPORTED = False
44from lsst.utils.db_auth import DbAuth, DbAuthNotFoundError
46from ..monitor import MonAgent
47from ..timer import Timer
48from .cassandra_utils import raw_data_factory
50if TYPE_CHECKING:
51 from .config import ApdbCassandraConfig
53_LOG = logging.getLogger(__name__)
55_MON = MonAgent(__name__)
58def _dump_query(rf: Any) -> None:
59 """Dump cassandra query to debug log."""
60 _LOG.debug("Cassandra query: %s", rf.query)
63if CASSANDRA_IMPORTED: 63 ↛ 65line 63 didn't jump to line 65 because the condition on line 63 was never true
65 class _AddressTranslator(AddressTranslator):
66 """Translate internal IP address to external.
68 Only used for docker-based setup, not a viable long-term solution.
69 """
71 def __init__(self, public_ips: tuple[str, ...], private_ips: tuple[str, ...]):
72 self._map = dict(zip(private_ips, public_ips))
74 def translate(self, private_ip: str) -> str:
75 return self._map.get(private_ip, private_ip)
78class SessionFactory:
79 """Implementation of SessionFactory that uses parameters from Apdb
80 configuration.
82 Parameters
83 ----------
84 config : `ApdbCassandraConfig`
85 Configuration object.
86 """
88 def __init__(self, config: ApdbCassandraConfig):
89 self._config = config
90 self._cluster: Cluster | None = None
91 self._session: Session | None = None
93 def __del__(self) -> None:
94 # Need to call Cluster.shutdown() to avoid warnings.
95 if hasattr(self, "_cluster"):
96 self.close()
98 def close(self) -> None:
99 """Close cluster connection"""
100 if self._cluster:
101 self._cluster.shutdown()
102 self._cluster = None
103 self._session = None
105 def session(self) -> Session:
106 """Return Cassandra Session, making new connection if necessary.
108 Returns
109 -------
110 session : `cassandra.cluster.Sesion`
111 Cassandra session object.
112 """
113 if self._session is None:
114 self._cluster, self._session = self._make_session()
115 return self._session
117 def _make_session(self) -> tuple[Cluster, Session]:
118 """Make Cassandra session.
120 Returns
121 -------
122 cluster : `cassandra.cluster.Cluster`
123 Cassandra Cluster object
124 session : `cassandra.cluster.Session`
125 Cassandra session object
126 """
127 addressTranslator: AddressTranslator | None = None
128 if self._config.connection_config.private_ips:
129 addressTranslator = _AddressTranslator(
130 self._config.contact_points, self._config.connection_config.private_ips
131 )
133 extra_parameters = {
134 "idle_heartbeat_interval": 0,
135 "idle_heartbeat_timeout": 30,
136 "control_connection_timeout": 100,
137 "executor_threads": 10,
138 }
139 extra_parameters.update(self._config.connection_config.extra_parameters)
140 with Timer("cluster_connect", _MON):
141 cluster = Cluster(
142 execution_profiles=self._make_profiles(),
143 contact_points=self._config.contact_points,
144 port=self._config.connection_config.port,
145 address_translator=addressTranslator,
146 protocol_version=self._config.connection_config.protocol_version,
147 auth_provider=self._make_auth_provider(),
148 **extra_parameters,
149 )
150 session = cluster.connect()
152 # Dump queries if debug level is enabled.
153 if _LOG.isEnabledFor(logging.DEBUG):
154 session.add_request_init_listener(_dump_query)
156 # Disable result paging
157 session.default_fetch_size = None
159 return cluster, session
161 def _make_auth_provider(self) -> AuthProvider | None:
162 """Make Cassandra authentication provider instance."""
163 try:
164 dbauth = DbAuth()
165 except DbAuthNotFoundError:
166 # Credentials file doesn't exist, use anonymous login.
167 return None
169 # If dbauth_alias is defined then try it first without port number.
170 hosts: list[tuple[str, int | None]] = [
171 (hostname, self._config.connection_config.port) for hostname in self._config.contact_points
172 ]
173 if dbauth_alias := self._config.get_dbauth_alias():
174 hosts = [(dbauth_alias, None)] + hosts
176 empty_username = False
177 # Try every contact point in turn.
178 for hostname, port in hosts:
179 try:
180 username, password = dbauth.getAuth(
181 "cassandra",
182 self._config.connection_config.username,
183 hostname,
184 port,
185 self._config.keyspace,
186 )
187 if not username:
188 # Password without user name, try next hostname, but give
189 # warning later if no better match is found.
190 empty_username = True
191 else:
192 return PlainTextAuthProvider(username=username, password=password)
193 except DbAuthNotFoundError:
194 pass
196 if empty_username:
197 _LOG.warning(
198 f"Credentials file ({dbauth.db_auth_path}) provided password but not "
199 "user name, anonymous Cassandra logon will be attempted."
200 )
202 return None
204 def _make_profiles(self) -> Mapping[Any, ExecutionProfile]:
205 """Make all execution profiles used in the code."""
206 config = self._config
207 if config.connection_config.private_ips:
208 loadBalancePolicy = WhiteListRoundRobinPolicy(hosts=config.contact_points)
209 else:
210 loadBalancePolicy = RoundRobinPolicy()
212 read_tuples_profile = ExecutionProfile(
213 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
214 request_timeout=config.connection_config.read_timeout,
215 row_factory=cassandra.query.tuple_factory,
216 load_balancing_policy=loadBalancePolicy,
217 )
218 read_named_tuples_profile = ExecutionProfile(
219 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
220 request_timeout=config.connection_config.read_timeout,
221 row_factory=cassandra.query.named_tuple_factory,
222 load_balancing_policy=loadBalancePolicy,
223 )
224 read_raw_profile = ExecutionProfile(
225 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
226 request_timeout=config.connection_config.read_timeout,
227 row_factory=raw_data_factory,
228 load_balancing_policy=loadBalancePolicy,
229 )
230 # Profile to use with select_concurrent to return raw data (columns and
231 # rows)
232 read_raw_multi_profile = ExecutionProfile(
233 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
234 request_timeout=config.connection_config.read_timeout,
235 row_factory=raw_data_factory,
236 load_balancing_policy=loadBalancePolicy,
237 )
238 # Profile to use with select_concurrent to return raw data,
239 # this also has very long timeout, to be be use for querying
240 # DiaObjectDedup table that can return a lot of data.
241 read_raw_multi_dedup_profile = ExecutionProfile(
242 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
243 request_timeout=3600.0,
244 row_factory=raw_data_factory,
245 load_balancing_policy=loadBalancePolicy,
246 )
247 write_profile = ExecutionProfile(
248 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.write_consistency),
249 request_timeout=config.connection_config.write_timeout,
250 load_balancing_policy=loadBalancePolicy,
251 )
252 # To replace default DCAwareRoundRobinPolicy
253 default_profile = ExecutionProfile(
254 load_balancing_policy=loadBalancePolicy,
255 )
256 return {
257 "read_tuples": read_tuples_profile,
258 "read_named_tuples": read_named_tuples_profile,
259 "read_raw": read_raw_profile,
260 "read_raw_multi": read_raw_multi_profile,
261 "read_raw_multi_dedup": read_raw_multi_dedup_profile,
262 "write": write_profile,
263 EXEC_PROFILE_DEFAULT: default_profile,
264 }
267class SessionContext(ExitStack):
268 """Context manager for creating short-lived Cassandra sessions.
270 Parameters
271 ----------
272 config : `ApdbCassandraConfig`
273 Configuration object.
274 """
276 def __init__(self, config: ApdbCassandraConfig):
277 super().__init__()
278 self._session_factory = SessionFactory(config)
280 def __enter__(self) -> Session:
281 super().__enter__()
282 cluster, session = self._session_factory._make_session()
283 self.enter_context(cluster)
284 self.enter_context(session)
285 return session