Coverage for python/lsst/dax/apdb/cassandra/apdbCassandraReplica.py: 14%

1# This file is part of dax_apdb.

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.

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__ = ["ApdbCassandraReplica"]

26import logging

27from collections.abc import Iterable, Mapping, Sequence

28from typing import TYPE_CHECKING, cast

30import astropy.time

31import felis.datamodel

33from ..apdbReplica import ApdbReplica, ApdbTableData, ReplicaChunk

34from ..apdbSchema import ApdbTables

35from ..apdbUpdateRecord import ApdbUpdateRecord

36from ..monitor import MonAgent

37from ..schema_model import ExtraDataTypes

38from ..timer import Timer

39from ..versionTuple import VersionTuple

40from .apdbCassandraSchema import ExtraTables

41from .cassandra_utils import (

42 ApdbCassandraTableData,

43 execute_concurrent,

44 select_concurrent,

45)

46from .queries import Column as C # noqa: N817

47from .queries import ColumnExpr, Delete, Select

49if TYPE_CHECKING:

50 from .apdbCassandra import ApdbCassandra

52_LOG = logging.getLogger(__name__)

54_MON = MonAgent(__name__)

56VERSION = VersionTuple(1, 1, 1)

57"""Version for the code controlling replication tables. This needs to be

58updated following compatibility rules when schema produced by this code

59changes.

60"""

63class ApdbCassandraReplica(ApdbReplica):

64 """Implementation of `ApdbReplica` for Cassandra backend.

66 Parameters

67 ----------

68 apdb : `ApdbCassandra`

69 Instance of ApbdCassandra for database.

70 """

72 def __init__(self, apdb: ApdbCassandra):

73 # Note that ApdbCassandra instance must stay alive while this object

74 # exists, so we keep reference to it.

75 self._apdb = apdb

77 def close(self) -> None:

78 # docstring is inherited from a base class

79 self._apdb.close()

81 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:

82 """Create `Timer` instance given its name."""

83 return Timer(name, _MON, tags=tags)

85 def schemaVersion(self) -> VersionTuple:

86 # Docstring inherited from base class.

87 context = self._apdb._context

88 return context.db_versions.schema_version

90 @classmethod

91 def apdbReplicaImplementationVersion(cls) -> VersionTuple:

92 # Docstring inherited from base class.

93 return VERSION

95 @classmethod

96 def hasChunkSubPartitions(cls, version: VersionTuple) -> bool:

97 """Return True if replica chunk tables have sub-partitions."""

98 return version >= VersionTuple(1, 1, 0)

100 @classmethod

101 def hasUpdateRecordChunks(cls, version: VersionTuple) -> bool:

102 """Return True if ApdbUpdateRecordChunks should exists."""

103 return version >= VersionTuple(1, 1, 1)

104

105 def getReplicaChunks(self) -> list[ReplicaChunk] | None:

106 # docstring is inherited from a base class

107 context = self._apdb._context

108 config = context.config

109

110 if not context.schema.replication_enabled:

111 return None

112

113 # everything goes into a single partition

114 partition = 0

115

116 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)

117 # We want to avoid timezone mess so return timestamps as milliseconds.

118 columns = (ColumnExpr("toUnixTimestamp(last_update_time)"), "apdb_replica_chunk", "unique_id")

119 query = Select(config.keyspace, table_name, columns).where(C("partition") == partition)

120 statement, params = context.stmt_factory.with_params(query, prepare=False)

121

122 with self._timer("chunks_select_time") as timer:

123 result = context.session.execute(

124 statement,

125 params,

126 timeout=config.connection_config.read_timeout,

127 execution_profile="read_tuples",

128 )

129 # order by last_update_time

130 rows = sorted(result)

131 timer.add_values(row_count=len(rows))

132 return [

133 ReplicaChunk(

134 id=row[1],

135 last_update_time=astropy.time.Time(row[0] / 1000, format="unix_tai"),

136 unique_id=row[2],

137 )

138 for row in rows

139 ]

140

141 def deleteReplicaChunks(self, chunks: Iterable[int]) -> None:

142 # docstring is inherited from a base class

143 context = self._apdb._context

144 config = context.config

145

146 if not context.schema.replication_enabled:

147 raise ValueError("APDB is not configured for replication")

148

149 # everything goes into a single partition

150 partition = 0

151

152 # Iterable can be single pass, make everything that we need from it

153 # in a single loop.

154 repl_table_params = []

155 chunk_table_params: list[tuple] = []

156 for chunk in chunks:

157 repl_table_params.append((partition, chunk))

158 if context.has_chunk_sub_partitions:

159 for subchunk in range(config.replica_sub_chunk_count):

160 chunk_table_params.append((chunk, subchunk))

161 else:

162 chunk_table_params.append((chunk,))

163 # Anything to do att all?

164 if not repl_table_params:

165 return

166

167 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)

168 query = Delete(config.keyspace, table_name)

169 query = query.where(C("partition") == -1)

170 query = query.where(C("apdb_replica_chunk") == -1)

171 statement = context.stmt_factory(query)

172

173 queries = [(statement, param) for param in repl_table_params]

174 with self._timer("chunks_delete_time") as timer:

175 execute_concurrent(context.session, queries)

176 timer.add_values(row_count=len(queries))

177

178 # Also remove those chunk_ids from Dia*Chunks tables.

179 tables = list(ExtraTables.replica_chunk_tables(context.has_chunk_sub_partitions).values())

180 if context.has_update_record_chunks_table:

181 tables.append(ExtraTables.ApdbUpdateRecordChunks)

182 for table in tables:

183 table_name = context.schema.tableName(table)

184 query = Delete(config.keyspace, table_name)

185 query = query.where(C("apdb_replica_chunk") == -1)

186 if context.has_chunk_sub_partitions:

187 query = query.where(C("apdb_replica_subchunk") == -1)

188 statement = context.stmt_factory(query)

189

190 queries = [(statement, param) for param in chunk_table_params]

191 with self._timer("table_chunk_detele_time", tags={"table": table_name}) as timer:

192 execute_concurrent(context.session, queries)

193 timer.add_values(row_count=len(queries))

194

195 def getTableDataChunks(self, table: ApdbTables, chunks: Iterable[int]) -> ApdbTableData:

196 # docstring is inherited from a base class

197 context = self._apdb._context

198 config = context.config

199

200 if not context.schema.replication_enabled:

201 raise ValueError("APDB is not configured for replication")

202 if table not in ExtraTables.replica_chunk_tables(False):

203 raise ValueError(f"Table {table} does not support replica chunks.")

204

205 # We need to iterate few times.

206 chunks = list(chunks)

207

208 # If schema was migrated then a chunk can appear in either old or new

209 # chunk table (e.g. DiaObjectChunks or DiaObjectChunks2). Chunk table

210 # has a column which will be set to true for new table.

211 has_chunk_sub_partitions: dict[int, bool] = {}

212 if context.has_chunk_sub_partitions:

213 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)

214 partition = 0

215 query = Select(config.keyspace, table_name, ("apdb_replica_chunk", "has_subchunks"))

216 query = query.where(C("partition") == partition)

217 query = query.where(C("apdb_replica_chunk").in_(chunks))

218 stmt, params = context.stmt_factory.with_params(query, prepare=False)

219 result = context.session.execute(

220 stmt,

221 params,

222 timeout=config.connection_config.read_timeout,

223 execution_profile="read_tuples",

224 )

225 has_chunk_sub_partitions = dict(result)

226 else:

227 has_chunk_sub_partitions = dict.fromkeys(chunks, False)

228

229 # Check what kind of tables we want to query, if chunk list is empty

230 # then use tables which should exist in the schema.

231 if has_chunk_sub_partitions:

232 have_subchunks = any(has_chunk_sub_partitions.values())

233 have_non_subchunks = not all(has_chunk_sub_partitions.values())

234 else:

235 have_subchunks = context.has_chunk_sub_partitions

236 have_non_subchunks = not have_subchunks

237

238 # NOTE: if an existing database is migrated and has both types of chunk

239 # tables (e.g. DiaObjectChunks and DiaObjectChunks2) it is possible

240 # that the same chunk can appear in both tables. In reality schema

241 # migration should only happen during the downtime, so there will be

242 # suffient gap and a different chunk ID will be used for new chunks.

243

244 table_data: ApdbCassandraTableData | None = None

245 table_data_subchunk: ApdbCassandraTableData | None = None

246

247 table_name = context.schema.tableName(ExtraTables.replica_chunk_tables(False)[table])

248 with self._timer("table_chunk_select_time", tags={"table": table_name}) as timer:

249 if have_subchunks:

250 replica_table = ExtraTables.replica_chunk_tables(True)[table]

251 table_name = context.schema.tableName(replica_table)

252 query = Select(config.keyspace, table_name, ["*"])

253 query = query.where(C("apdb_replica_chunk") == 0)

254 query = query.where(C("apdb_replica_subchunk") == 0)

255 statement = context.stmt_factory(query, prepare=True)

256

257 queries: list[tuple] = []

258 for chunk in chunks:

259 if has_chunk_sub_partitions.get(chunk, False):

260 for subchunk in range(config.replica_sub_chunk_count):

261 queries.append((statement, (chunk, subchunk)))

262 if not queries and not have_non_subchunks:

263 # Add a dummy query to return correct set of columns.

264 queries.append((statement, (-1, -1)))

265

266 if queries:

267 table_data_subchunk = cast(

268 ApdbCassandraTableData,

269 select_concurrent(

270 context.session,

271 queries,

272 "read_raw_multi",

273 config.connection_config.read_concurrency,

274 ),

275 )

276

277 if have_non_subchunks:

278 replica_table = ExtraTables.replica_chunk_tables(False)[table]

279 table_name = context.schema.tableName(replica_table)

280 query = Select(config.keyspace, table_name, ["*"]).where(C("apdb_replica_chunk") == 0)

281 statement = context.stmt_factory(query, prepare=True)

282

283 queries = []

284 for chunk in chunks:

285 if not has_chunk_sub_partitions.get(chunk, True):

286 queries.append((statement, (chunk,)))

287 if not queries and not table_data_subchunk:

288 # Add a dummy query to return correct set of columns.

289 queries.append((statement, (-1,)))

290

291 if queries:

292 table_data = cast(

293 ApdbCassandraTableData,

294 select_concurrent(

295 context.session,

296 queries,

297 "read_raw_multi",

298 config.connection_config.read_concurrency,

299 ),

300 )

301

302 # Merge if both are non-empty.

303 if table_data and table_data_subchunk:

304 table_data_subchunk.project(drop=["apdb_replica_subchunk"])

305 table_data.append(table_data_subchunk)

306 elif table_data_subchunk:

307 table_data = table_data_subchunk

308 elif not table_data:

309 raise AssertionError("above logic is incorrect")

310

311 timer.add_values(row_count=len(table_data.rows()))

312

313 table_schema = self._apdb._schema.tableSchemas[table]

314 # Regular tables should never have columns of ExtraDataTypes, this

315 # is just to make mypy happy.

316 column_types = {

317 column.name: column.datatype

318 for column in table_schema.columns

319 if not isinstance(column.datatype, ExtraDataTypes)

320 }

321 column_types["apdb_replica_chunk"] = felis.datamodel.DataType.long

322 # It may also have subchunk column, we do not always drop it, and

323 # clients should not need it, but we need to provide type for it.

324 column_types["apdb_replica_subchunk"] = felis.datamodel.DataType.int

325 table_data.set_column_types(column_types)

326

327 self._fix_null_columns(table, table_data)

328

329 return table_data

330

331 def getUpdateRecordChunks(self, chunks: Iterable[int]) -> Sequence[ApdbUpdateRecord]:

332 # docstring is inherited from a base class

333 context = self._apdb._context

334 config = context.config

335

336 if not context.schema.replication_enabled:

337 raise ValueError("APDB is not configured for replication")

338

339 if not context.has_update_record_chunks_table:

340 # Table does not exist yet.

341 return []

342

343 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)

344

345 records = []

346 if context.has_chunk_sub_partitions:

347 subchunks = list(range(config.replica_sub_chunk_count))

348 query = Select(config.keyspace, table_name, ["*"])

349 query = query.where(C("apdb_replica_chunk") == 0)

350 query = query.where(C("apdb_replica_subchunk").in_(subchunks))

351 statement = context.stmt_factory(query, prepare=False)

352

353 with self._timer("select_update_record_time", tags={"table": table_name}) as timer:

354 for chunk in chunks:

355 result = context.session.execute(statement, [chunk] + subchunks)

356 for row in result:

357 records.append(

358 ApdbUpdateRecord.from_json(

359 row.update_time_ns, row.update_order, row.update_payload

360 )

361 )

362 timer.add_values(row_count=len(records))

363

364 else:

365 query = Select(config.keyspace, table_name, ["*"])

366 query = query.where(C("apdb_replica_chunk").in_(chunks))

367 statement, params = context.stmt_factory(query, prepare=False)

368

369 with self._timer("select_update_record_time", tags={"table": table_name}) as timer:

370 result = context.session.execute(statement, params)

371 for row in result:

372 records.append(

373 ApdbUpdateRecord.from_json(row.update_time_ns, row.update_order, row.update_payload)

374 )

375 timer.add_values(row_count=len(records))

376

377 records.sort()

378 return records

379

380 def _fix_null_columns(self, table: ApdbTables, table_data: ApdbCassandraTableData) -> None:

381 """For non-null columns replace NULL values with column defaults.

382

383 When new non-null columns are added to the schema we do not populate

384 values for existing records. Instead when reading old data back we

385 fill NULLs with the defaults. Table data is updated in-place.

386 """

387 # Find all non-null columns which also have default values defined.

388 table_schema = self._apdb._schema.tableSchemas[table]

389 # Defaults are indexed by column index.

390 defaults: dict[int, object] = {

391 table_data.column_index(column.name): column.value

392 for column in table_schema.columns

393 if not column.nullable and column.value is not None

394 }

395 if defaults:

396 # Rows are tuples, need to convert to list to modify.

397 for row_idx, row in enumerate(table_data._rows):

398 if any(row[def_idx] is None for def_idx in defaults):

399 fixed = tuple(

400 defaults[idx] if idx in defaults and item is None else item

401 for idx, item in enumerate(list(row))

402 )

403 table_data._rows[row_idx] = fixed