Coverage for python / lsst / dax / apdb / cassandra / sessionFactory.py: 26%
119 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__ = ["SessionContext", "SessionFactory"]
26import contextlib
27import cProfile
28import io
29import logging
30import os
31import pstats
32from collections.abc import Mapping
33from contextlib import ExitStack
34from typing import TYPE_CHECKING, Any
36# If cassandra-driver is not there the module can still be imported
37# but ApdbCassandra cannot be instantiated.
38try:
39 import cassandra
40 import cassandra.query
41 from cassandra.auth import AuthProvider, PlainTextAuthProvider
42 from cassandra.cluster import EXEC_PROFILE_DEFAULT, Cluster, ExecutionProfile, Session
43 from cassandra.policies import AddressTranslator, RoundRobinPolicy, WhiteListRoundRobinPolicy
45 CASSANDRA_IMPORTED = True
46except ImportError:
47 CASSANDRA_IMPORTED = False
49from lsst.utils.db_auth import DbAuth, DbAuthNotFoundError
51from ..monitor import MonAgent
52from ..timer import Timer
53from .cassandra_utils import raw_data_factory
55if TYPE_CHECKING:
56 from .config import ApdbCassandraConfig
58_LOG = logging.getLogger(__name__)
60_MON = MonAgent(__name__)
62_DO_PROFILE = os.environ.get("DAX_APDB_DM53808_PROFILE")
65def _profile_context() -> Any:
66 if _DO_PROFILE:
67 return cProfile.Profile()
68 else:
69 return contextlib.nullcontext()
72def _profile_cpu_threshold() -> float:
73 if _DO_PROFILE:
74 try:
75 return float(_DO_PROFILE)
76 except Exception:
77 return 2.5
78 else:
79 return 1_000_000.0
82def _dump_query(rf: Any) -> None:
83 """Dump cassandra query to debug log."""
84 _LOG.debug("Cassandra query: %s", rf.query)
87if CASSANDRA_IMPORTED: 87 ↛ 89line 87 didn't jump to line 89 because the condition on line 87 was never true
89 class _AddressTranslator(AddressTranslator):
90 """Translate internal IP address to external.
92 Only used for docker-based setup, not a viable long-term solution.
93 """
95 def __init__(self, public_ips: tuple[str, ...], private_ips: tuple[str, ...]):
96 self._map = dict(zip(private_ips, public_ips))
98 def translate(self, private_ip: str) -> str:
99 return self._map.get(private_ip, private_ip)
102class SessionFactory:
103 """Implementation of SessionFactory that uses parameters from Apdb
104 configuration.
106 Parameters
107 ----------
108 config : `ApdbCassandraConfig`
109 Configuration object.
110 """
112 def __init__(self, config: ApdbCassandraConfig):
113 self._config = config
114 self._cluster: Cluster | None = None
115 self._session: Session | None = None
117 def __del__(self) -> None:
118 # Need to call Cluster.shutdown() to avoid warnings.
119 if hasattr(self, "_cluster"):
120 if self._cluster:
121 self._cluster.shutdown()
123 def session(self) -> Session:
124 """Return Cassandra Session, making new connection if necessary.
126 Returns
127 -------
128 session : `cassandra.cluster.Sesion`
129 Cassandra session object.
130 """
131 if self._session is None:
132 self._cluster, self._session = self._make_session()
133 return self._session
135 def _make_session(self) -> tuple[Cluster, Session]:
136 """Make Cassandra session.
138 Returns
139 -------
140 cluster : `cassandra.cluster.Cluster`
141 Cassandra Cluster object
142 session : `cassandra.cluster.Session`
143 Cassandra session object
144 """
145 addressTranslator: AddressTranslator | None = None
146 if self._config.connection_config.private_ips:
147 addressTranslator = _AddressTranslator(
148 self._config.contact_points, self._config.connection_config.private_ips
149 )
151 extra_parameters = {
152 "idle_heartbeat_interval": 0,
153 "idle_heartbeat_timeout": 30,
154 "control_connection_timeout": 100,
155 "executor_threads": 10,
156 }
157 extra_parameters.update(self._config.connection_config.extra_parameters)
158 with Timer("cluster_connect", _MON) as timer:
159 with _profile_context() as cluster_connect_prf:
160 cluster = Cluster(
161 execution_profiles=self._make_profiles(),
162 contact_points=self._config.contact_points,
163 port=self._config.connection_config.port,
164 address_translator=addressTranslator,
165 protocol_version=self._config.connection_config.protocol_version,
166 auth_provider=self._make_auth_provider(),
167 **extra_parameters,
168 )
169 session = cluster.connect()
171 if cluster_connect_prf and timer.accumulated()[1] > _profile_cpu_threshold():
172 out = io.StringIO()
173 ps = (
174 pstats.Stats(cluster_connect_prf, stream=out)
175 .sort_stats(pstats.SortKey.CUMULATIVE)
176 .strip_dirs()
177 )
178 ps.print_stats(10)
179 # Make it a single line.
180 text = out.getvalue().strip().replace("\n", "\\n")
181 _LOG.info(text)
183 # Dump queries if debug level is enabled.
184 if _LOG.isEnabledFor(logging.DEBUG):
185 session.add_request_init_listener(_dump_query)
187 # Disable result paging
188 session.default_fetch_size = None
190 return cluster, session
192 def _make_auth_provider(self) -> AuthProvider | None:
193 """Make Cassandra authentication provider instance."""
194 try:
195 dbauth = DbAuth()
196 except DbAuthNotFoundError:
197 # Credentials file doesn't exist, use anonymous login.
198 return None
200 empty_username = True
201 # Try every contact point in turn.
202 for hostname in self._config.contact_points:
203 try:
204 username, password = dbauth.getAuth(
205 "cassandra",
206 self._config.connection_config.username,
207 hostname,
208 self._config.connection_config.port,
209 self._config.keyspace,
210 )
211 if not username:
212 # Password without user name, try next hostname, but give
213 # warning later if no better match is found.
214 empty_username = True
215 else:
216 return PlainTextAuthProvider(username=username, password=password)
217 except DbAuthNotFoundError:
218 pass
220 if empty_username:
221 _LOG.warning(
222 f"Credentials file ({dbauth.db_auth_path}) provided password but not "
223 "user name, anonymous Cassandra logon will be attempted."
224 )
226 return None
228 def _make_profiles(self) -> Mapping[Any, ExecutionProfile]:
229 """Make all execution profiles used in the code."""
230 config = self._config
231 if config.connection_config.private_ips:
232 loadBalancePolicy = WhiteListRoundRobinPolicy(hosts=config.contact_points)
233 else:
234 loadBalancePolicy = RoundRobinPolicy()
236 read_tuples_profile = ExecutionProfile(
237 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
238 request_timeout=config.connection_config.read_timeout,
239 row_factory=cassandra.query.tuple_factory,
240 load_balancing_policy=loadBalancePolicy,
241 )
242 read_named_tuples_profile = ExecutionProfile(
243 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
244 request_timeout=config.connection_config.read_timeout,
245 row_factory=cassandra.query.named_tuple_factory,
246 load_balancing_policy=loadBalancePolicy,
247 )
248 read_raw_profile = ExecutionProfile(
249 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
250 request_timeout=config.connection_config.read_timeout,
251 row_factory=raw_data_factory,
252 load_balancing_policy=loadBalancePolicy,
253 )
254 # Profile to use with select_concurrent to return raw data (columns and
255 # rows)
256 read_raw_multi_profile = ExecutionProfile(
257 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
258 request_timeout=config.connection_config.read_timeout,
259 row_factory=raw_data_factory,
260 load_balancing_policy=loadBalancePolicy,
261 )
262 # Profile to use with select_concurrent to return raw data,
263 # this also has very long timeout, to be be use for querying
264 # DiaObjectDedup table that can return a lot of data.
265 read_raw_multi_dedup_profile = ExecutionProfile(
266 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.read_consistency),
267 request_timeout=3600.0,
268 row_factory=raw_data_factory,
269 load_balancing_policy=loadBalancePolicy,
270 )
271 write_profile = ExecutionProfile(
272 consistency_level=getattr(cassandra.ConsistencyLevel, config.connection_config.write_consistency),
273 request_timeout=config.connection_config.write_timeout,
274 load_balancing_policy=loadBalancePolicy,
275 )
276 # To replace default DCAwareRoundRobinPolicy
277 default_profile = ExecutionProfile(
278 load_balancing_policy=loadBalancePolicy,
279 )
280 return {
281 "read_tuples": read_tuples_profile,
282 "read_named_tuples": read_named_tuples_profile,
283 "read_raw": read_raw_profile,
284 "read_raw_multi": read_raw_multi_profile,
285 "read_raw_multi_dedup": read_raw_multi_dedup_profile,
286 "write": write_profile,
287 EXEC_PROFILE_DEFAULT: default_profile,
288 }
291class SessionContext(ExitStack):
292 """Context manager for creating short-lived Cassandra sessions.
294 Parameters
295 ----------
296 config : `ApdbCassandraConfig`
297 Configuration object.
298 """
300 def __init__(self, config: ApdbCassandraConfig):
301 super().__init__()
302 self._session_factory = SessionFactory(config)
304 def __enter__(self) -> Session:
305 super().__enter__()
306 cluster, session = self._session_factory._make_session()
307 self.enter_context(cluster)
308 self.enter_context(session)
309 return session