Coverage for python/lsst/daf/butler/registry/dimensions/static.py: 91%

1# This file is part of daf_butler. 

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27from __future__ import annotations 

28 

29import dataclasses 

30import itertools 

31import logging 

32from collections import defaultdict 

33from collections.abc import Iterable, Mapping, Sequence, Set 

34from typing import TYPE_CHECKING, Any 

35 

36import sqlalchemy 

37 

38from lsst.daf.relation import Calculation, ColumnExpression, Join, Relation, sql 

39from lsst.sphgeom import Region 

40 

41from ... import ddl 

42from ..._column_tags import DimensionKeyColumnTag, DimensionRecordColumnTag 

43from ..._column_type_info import LogicalColumn 

44from ..._exceptions import UnimplementedQueryError 

45from ..._named import NamedKeyDict 

46from ...dimensions import ( 

47 DatabaseDimensionElement, 

48 DatabaseTopologicalFamily, 

49 DataCoordinate, 

50 DataIdValue, 

51 Dimension, 

52 DimensionElement, 

53 DimensionGroup, 

54 DimensionRecord, 

55 DimensionRecordSet, 

56 DimensionUniverse, 

57 SkyPixDimension, 

58 addDimensionForeignKey, 

59) 

60from ...dimensions.record_cache import DimensionRecordCache 

61from ...direct_query_driver import ( # Future query system (direct,server). 

62 Postprocessing, 

63 SqlJoinsBuilder, 

64 SqlSelectBuilder, 

65) 

66from ...queries import tree as qt # Future query system (direct,client,server) 

67from ...queries.overlaps import OverlapsVisitor 

68from ...queries.visitors import PredicateVisitFlags 

69from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext, VersionTuple 

70 

71if TYPE_CHECKING: 

72 from .. import queries # Current Registry.query* system. 

73 

74 

75# This has to be updated on every schema change 

76_VERSION = VersionTuple(6, 0, 2) 

77 

78_LOG = logging.getLogger(__name__) 

79 

80 

81class StaticDimensionRecordStorageManager(DimensionRecordStorageManager): 

82 """An implementation of `DimensionRecordStorageManager` for single-layer 

83 `Registry` and the base layers of multi-layer `Registry`. 

84 

85 This manager creates `DimensionRecordStorage` instances for all elements 

86 in the `DimensionUniverse` in its own `initialize` method, as part of 

87 static table creation, so it never needs to manage any dynamic registry 

88 tables. 

89 

90 Parameters 

91 ---------- 

92 db : `Database` 

93 Interface to the underlying database engine and namespace. 

94 tables : `dict` [ `str`, `sqlalchemy.Table` ] 

95 Mapping from dimension element name to SQL table, for all elements that 

96 have `DimensionElement.has_own_table` `True`. 

97 overlap_tables : `dict` [ `str`, `tuple` [ `sqlalchemy.Table`, \ 

98 `sqlalchemy.Table` ] ] 

99 Mapping from dimension element name to SQL table holding overlaps 

100 between the common skypix dimension and that element, for all elements 

101 that have `DimensionElement.has_own_table` `True` and 

102 `DimensionElement.spatial` not `None`. 

103 dimension_group_storage : `_DimensionGroupStorage` 

104 Object that manages saved `DimensionGroup` definitions. 

105 universe : `DimensionUniverse` 

106 All known dimensions. 

107 registry_schema_version : `VersionTuple` or `None`, optional 

108 Version of registry schema. 

109 """ 

110 

111 def __init__( 

112 self, 

113 db: Database, 

114 *, 

115 tables: dict[str, sqlalchemy.Table], 

116 overlap_tables: dict[str, tuple[sqlalchemy.Table, sqlalchemy.Table]], 

117 dimension_group_storage: _DimensionGroupStorage, 

118 universe: DimensionUniverse, 

119 registry_schema_version: VersionTuple | None = None, 

120 ): 

121 super().__init__(universe=universe, registry_schema_version=registry_schema_version) 

122 self._db = db 

123 self._tables = tables 

124 self._overlap_tables = overlap_tables 

125 self._dimension_group_storage = dimension_group_storage 

126 

127 def clone(self, db: Database) -> StaticDimensionRecordStorageManager: 

128 return StaticDimensionRecordStorageManager( 

129 db, 

130 tables=self._tables, 

131 overlap_tables=self._overlap_tables, 

132 dimension_group_storage=self._dimension_group_storage.clone(db), 

133 universe=self.universe, 

134 registry_schema_version=self._registry_schema_version, 

135 ) 

136 

137 @classmethod 

138 def initialize( 

139 cls, 

140 db: Database, 

141 context: StaticTablesContext, 

142 *, 

143 universe: DimensionUniverse, 

144 registry_schema_version: VersionTuple | None = None, 

145 ) -> DimensionRecordStorageManager: 

146 # Docstring inherited from DimensionRecordStorageManager. 

147 tables: dict[str, sqlalchemy.Table] = {} 

148 # Define tables for governor dimensions, which are never spatial or 

149 # temporal and always have tables. 

150 for dimension in universe.governor_dimensions: 

151 spec = dimension.RecordClass.fields.makeTableSpec( 

152 TimespanReprClass=db.getTimespanRepresentation() 

153 ) 

154 tables[dimension.name] = context.addTable(dimension.name, spec) 

155 # Define tables for database dimension elements, which may or may not 

156 # have their own tables and may be spatial or temporal. 

157 spatial = NamedKeyDict[DatabaseTopologicalFamily, list[DimensionElement]]() 

158 overlap_tables: dict[str, tuple[sqlalchemy.Table, sqlalchemy.Table]] = {} 

159 for element in universe.database_elements: 

160 if not element.has_own_table: 

161 continue 

162 spec = element.RecordClass.fields.makeTableSpec(TimespanReprClass=db.getTimespanRepresentation()) 

163 tables[element.name] = context.addTable(element.name, spec) 

164 if element.spatial is not None: 

165 spatial.setdefault(element.spatial, []).append(element) 

166 overlap_tables[element.name] = cls._make_skypix_overlap_tables(context, element) 

167 for field_name in spec.fields.names: 

168 if ( 168 ↛ 176line 168 didn't jump to line 176 because the condition on line 168 was never true

169 len(qt.ColumnSet.get_qualified_name(element.name, field_name)) 

170 >= db.dialect.max_identifier_length 

171 ): 

172 # Being able to assume that all dimension fields fit inside 

173 # the DB's identifier limit is really convenient and very 

174 # unlikely to cause trouble in practice. We'll just make 

175 # sure we catch any such trouble as early as possible. 

176 raise RuntimeError( 

177 f"Dimension filed '{element.name}.{field_name}' is too long for this database. " 

178 "Please file a ticket for long-field support if this was not a mistake." 

179 ) 

180 # Add some tables for materialized overlaps between database 

181 # dimensions. We've never used these and no longer plan to, but we 

182 # have to keep creating them to keep schema versioning consistent. 

183 cls._make_legacy_overlap_tables(context, spatial) 

184 # Create tables that store DimensionGroup definitions. 

185 dimension_group_storage = _DimensionGroupStorage.initialize(db, context, universe=universe) 

186 return cls( 

187 db=db, 

188 tables=tables, 

189 overlap_tables=overlap_tables, 

190 universe=universe, 

191 dimension_group_storage=dimension_group_storage, 

192 registry_schema_version=registry_schema_version, 

193 ) 

194 

195 def fetch_cache_dict(self) -> dict[str, DimensionRecordSet]: 

196 # Docstring inherited. 

197 result: dict[str, DimensionRecordSet] = {} 

198 with self._db.transaction(): 

199 for element in self.universe.elements: 

200 if not element.is_cached: 

201 continue 

202 assert not element.temporal, ( 

203 "Cached dimension elements should not be spatial or temporal, as that " 

204 "suggests a large number of records." 

205 ) 

206 if element.implied_union_target is not None: 

207 assert isinstance(element, Dimension), "Only dimensions can be implied dependencies." 

208 table = self._tables[element.implied_union_target.name] 

209 sql = sqlalchemy.select( 

210 table.columns[element.name].label(element.primary_key.name) 

211 ).distinct() 

212 else: 

213 table = self._tables[element.name] 

214 sql = table.select() 

215 with self._db.query(sql) as results: 

216 result[element.name] = DimensionRecordSet( 

217 element=element, 

218 records=[element.RecordClass(**row) for row in results.mappings()], 

219 ) 

220 return result 

221 

222 def insert( 

223 self, 

224 element: DimensionElement, 

225 *records: DimensionRecord, 

226 replace: bool = False, 

227 skip_existing: bool = False, 

228 ) -> None: 

229 # Docstring inherited. 

230 if not element.has_own_table: 

231 raise TypeError(f"Cannot insert {element.name} records.") 

232 db_rows = self._make_record_db_rows(element, records, replace=replace) 

233 table = self._tables[element.name] 

234 with self._db.transaction(): 

235 if replace: 

236 self._db.replace(table, *db_rows.main_rows) 

237 elif skip_existing: 

238 self._db.ensure(table, *db_rows.main_rows, primary_key_only=True) 

239 else: 

240 self._db.insert(table, *db_rows.main_rows) 

241 self._insert_overlaps( 

242 element, db_rows.overlap_insert_rows, db_rows.overlap_delete_rows, skip_existing=skip_existing 

243 ) 

244 for related_element_name, summary_rows in db_rows.overlap_summary_rows.items(): 

245 self._db.ensure(self._overlap_tables[related_element_name][0], *summary_rows) 

246 

247 def sync(self, record: DimensionRecord, update: bool = False) -> bool | dict[str, Any]: 

248 # Docstring inherited. 

249 if not record.definition.has_own_table: 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true

250 raise TypeError(f"Cannot sync {record.definition.name} records.") 

251 # We might not need the overlap rows at all; we won't know until we try 

252 # to insert the main row. But we figure it's better to spend the time 

253 # to compute them in advance always *outside* the database transaction 

254 # than to compute them only as-needed inside the database transaction, 

255 # since in-transaction time is especially precious. 

256 db_rows = self._make_record_db_rows(record.definition, [record], replace=True) 

257 (compared,) = db_rows.main_rows 

258 keys = {} 

259 for name in record.fields.required.names: 

260 keys[name] = compared.pop(name) 

261 with self._db.transaction(): 

262 _, inserted_or_updated = self._db.sync( 

263 self._tables[record.definition.name], 

264 keys=keys, 

265 compared=compared, 

266 update=update, 

267 ) 

268 if inserted_or_updated: 

269 if inserted_or_updated is True: 

270 # Inserted a new row, so we just need to insert new 

271 # overlap rows (if there are any). 

272 self._insert_overlaps( 

273 record.definition, db_rows.overlap_insert_rows, overlap_delete_rows=[] 

274 ) 

275 elif "region" in inserted_or_updated: 275 ↛ 281line 275 didn't jump to line 281 because the condition on line 275 was always true

276 # Updated the region, so we need to delete old overlap 

277 # rows and insert new ones. 

278 self._insert_overlaps( 

279 record.definition, db_rows.overlap_insert_rows, db_rows.overlap_delete_rows 

280 ) 

281 for related_element_name, summary_rows in db_rows.overlap_summary_rows.items(): 

282 self._db.ensure(self._overlap_tables[related_element_name][0], *summary_rows) 

283 return inserted_or_updated 

284 

285 def fetch_one( 

286 self, 

287 element_name: str, 

288 data_id: DataCoordinate, 

289 cache: DimensionRecordCache, 

290 ) -> DimensionRecord | None: 

291 # Docstring inherited. 

292 element = self.universe[element_name] 

293 if element_name in cache: 

294 try: 

295 return cache[element_name].find(data_id) 

296 except LookupError: 

297 return None 

298 if element.implied_union_target is not None: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true

299 assert isinstance(element, Dimension), "Only dimensions can be implied dependencies." 

300 table = self._tables[element.implied_union_target.name] 

301 sql = sqlalchemy.select(table.columns[element.name].label(element.primary_key.name)).where( 

302 table.columns[element_name] == data_id[element_name] 

303 ) 

304 elif isinstance(element, SkyPixDimension): 

305 id = data_id[element_name] 

306 return element.RecordClass(id=id, region=element.pixelization.pixel(id)) 

307 else: 

308 table = self._tables[element.name] 

309 sql = table.select().where( 

310 *[ 

311 table.columns[column_name] == data_id[dimension_name] 

312 for column_name, dimension_name in zip( 

313 element.schema.required.names, element.required.names 

314 ) 

315 ] 

316 ) 

317 with self._db.query(sql) as results: 

318 row = results.fetchone() 

319 if row is None: 

320 return None 

321 mapping: Mapping 

322 if element.temporal is not None: 

323 mapping = dict(**row._mapping) 

324 timespan = self._db.getTimespanRepresentation().extract(mapping) 

325 for name in self._db.getTimespanRepresentation().getFieldNames(): 

326 del mapping[name] 

327 mapping["timespan"] = timespan 

328 else: 

329 mapping = row._mapping 

330 return element.RecordClass(**mapping) 

331 

332 def save_dimension_group(self, group: DimensionGroup) -> int: 

333 # Docstring inherited from DimensionRecordStorageManager. 

334 return self._dimension_group_storage.save(group) 

335 

336 def load_dimension_group(self, key: int) -> DimensionGroup: 

337 # Docstring inherited from DimensionRecordStorageManager. 

338 return self._dimension_group_storage.load(key) 

339 

340 def join( 

341 self, 

342 element_name: str, 

343 target: Relation, 

344 join: Join, 

345 context: queries.SqlQueryContext, 

346 ) -> Relation: 

347 # Docstring inherited. 

348 element = self.universe[element_name] 

349 # We use Join.partial(...).apply(...) instead of Join.apply(..., ...) 

350 # for the "backtracking" insertion capabilities of the former; more 

351 # specifically, if `target` is a tree that starts with SQL relations 

352 # and ends with iteration-engine operations (e.g. region-overlap 

353 # postprocessing), this will try to perform the join upstream in the 

354 # SQL engine before the transfer to iteration. 

355 if element.has_own_table: 355 ↛ 357line 355 didn't jump to line 357 because the condition on line 355 was always true

356 return join.partial(self._make_relation(element, context)).apply(target) 

357 elif element.implied_union_target is not None: 

358 columns = DimensionKeyColumnTag(element.name) 

359 return join.partial( 

360 self._make_relation(element.implied_union_target, context) 

361 .with_only_columns( 

362 {columns}, 

363 preferred_engine=context.preferred_engine, 

364 require_preferred_engine=True, 

365 ) 

366 .without_duplicates() 

367 ).apply(target) 

368 elif isinstance(element, SkyPixDimension): 

369 assert join.predicate.as_trivial(), "Expected trivial join predicate for skypix relation." 

370 id_column = DimensionKeyColumnTag(element.name) 

371 assert id_column in target.columns, "Guaranteed by QueryBuilder.make_dimension_target." 

372 function_name = f"{element.name}_region" 

373 context.iteration_engine.functions[function_name] = element.pixelization.pixel 

374 calculation = Calculation( 

375 tag=DimensionRecordColumnTag(element.name, "region"), 

376 expression=ColumnExpression.function(function_name, ColumnExpression.reference(id_column)), 

377 ) 

378 return calculation.apply( 

379 target, preferred_engine=context.iteration_engine, transfer=True, backtrack=True 

380 ) 

381 else: 

382 raise AssertionError(f"Unexpected definition of {element_name!r}.") 

383 

384 def make_joins_builder(self, element: DimensionElement, fields: Set[str]) -> SqlJoinsBuilder: 

385 if element.implied_union_target is not None: 

386 assert not fields, "Dimensions with implied-union storage never have fields." 

387 return SqlSelectBuilder( 

388 self.make_joins_builder(element.implied_union_target, fields), 

389 columns=qt.ColumnSet(element.minimal_group).drop_implied_dimension_keys(), 

390 distinct=True, 

391 ).into_joins_builder(postprocessing=None) 

392 if not element.has_own_table: 392 ↛ 393line 392 didn't jump to line 393 because the condition on line 392 was never true

393 raise UnimplementedQueryError(f"Cannot join dimension element {element} with no table.") 

394 table = self._tables[element.name] 

395 result = SqlJoinsBuilder(db=self._db, from_clause=table) 

396 for dimension_name, column_name in zip(element.required.names, element.schema.required.names): 

397 result.dimension_keys[dimension_name].append(table.columns[column_name]) 

398 result.extract_dimensions(element.implied.names) 

399 for field in fields: 

400 if field == "timespan": 

401 result.timespans[element.name] = self._db.getTimespanRepresentation().from_columns( 

402 table.columns 

403 ) 

404 else: 

405 result.fields[element.name][field] = table.columns[field] 

406 return result 

407 

408 def process_query_overlaps( 

409 self, 

410 dimensions: DimensionGroup, 

411 predicate: qt.Predicate, 

412 join_operands: Iterable[DimensionGroup], 

413 calibration_dataset_types: Set[str | qt.AnyDatasetType], 

414 allow_duplicates: bool, 

415 constraint_data_id: Mapping[str, DataIdValue], 

416 ) -> tuple[qt.Predicate, SqlSelectBuilder, Postprocessing]: 

417 overlaps_visitor = _CommonSkyPixMediatedOverlapsVisitor( 

418 self._db, 

419 dimensions, 

420 calibration_dataset_types, 

421 self._overlap_tables, 

422 allow_duplicates, 

423 constraint_data_id=constraint_data_id, 

424 ) 

425 new_predicate = overlaps_visitor.run(predicate, join_operands) 

426 return new_predicate, overlaps_visitor.builder, overlaps_visitor.postprocessing 

427 

428 def _make_relation( 

429 self, 

430 element: DimensionElement, 

431 context: queries.SqlQueryContext, 

432 ) -> Relation: 

433 table = self._tables[element.name] 

434 payload = sql.Payload[LogicalColumn](table) 

435 for tag, field_name in element.RecordClass.fields.columns.items(): 

436 if field_name == "timespan": 436 ↛ 437line 436 didn't jump to line 437 because the condition on line 436 was never true

437 payload.columns_available[tag] = self._db.getTimespanRepresentation().from_columns( 

438 table.columns, name=field_name 

439 ) 

440 else: 

441 payload.columns_available[tag] = table.columns[field_name] 

442 return context.sql_engine.make_leaf( 

443 payload.columns_available.keys(), 

444 name=element.name, 

445 payload=payload, 

446 ) 

447 

448 @classmethod 

449 def currentVersions(cls) -> list[VersionTuple]: 

450 # Docstring inherited from VersionedExtension. 

451 return [_VERSION] 

452 

453 @classmethod 

454 def _make_skypix_overlap_tables( 

455 cls, context: StaticTablesContext, element: DimensionElement 

456 ) -> tuple[sqlalchemy.Table, sqlalchemy.Table]: 

457 assert element.governor is not None 

458 summary_spec = ddl.TableSpec( 

459 fields=[ 

460 ddl.FieldSpec( 

461 name="skypix_system", 

462 dtype=sqlalchemy.String, 

463 length=16, 

464 nullable=False, 

465 primaryKey=True, 

466 ), 

467 ddl.FieldSpec( 

468 name="skypix_level", 

469 dtype=sqlalchemy.SmallInteger, 

470 nullable=False, 

471 primaryKey=True, 

472 ), 

473 ] 

474 ) 

475 addDimensionForeignKey(summary_spec, element.governor, primaryKey=True) 

476 overlap_spec = ddl.TableSpec( 

477 fields=[ 

478 ddl.FieldSpec( 

479 name="skypix_system", 

480 dtype=sqlalchemy.String, 

481 length=16, 

482 nullable=False, 

483 primaryKey=True, 

484 ), 

485 ddl.FieldSpec( 

486 name="skypix_level", 

487 dtype=sqlalchemy.SmallInteger, 

488 nullable=False, 

489 primaryKey=True, 

490 ), 

491 # (more columns added below) 

492 ], 

493 unique=set(), 

494 indexes={ 

495 # This index has the same fields as the PK, in a different 

496 # order, to facilitate queries that know skypix_index and want 

497 # to find the other element. 

498 ddl.IndexSpec( 

499 "skypix_system", 

500 "skypix_level", 

501 "skypix_index", 

502 *element.minimal_group.required, 

503 ), 

504 }, 

505 foreignKeys=[ 

506 # Foreign key to summary table. This makes sure we don't 

507 # materialize any overlaps without remembering that we've done 

508 # so in the summary table, though it can't prevent the converse 

509 # of adding a summary row without adding overlap row (either of 

510 # those is a logic bug, of course, but we want to be defensive 

511 # about those). Using ON DELETE CASCADE, it'd be very easy to 

512 # implement "disabling" an overlap materialization, because we 

513 # can just delete the summary row. 

514 # Note that the governor dimension column is added below, in 

515 # the call to addDimensionForeignKey. 

516 ddl.ForeignKeySpec( 

517 f"{element.name}_skypix_overlap_summary", 

518 source=("skypix_system", "skypix_level", element.governor.name), 

519 target=("skypix_system", "skypix_level", element.governor.name), 

520 onDelete="CASCADE", 

521 ), 

522 ], 

523 ) 

524 # Add fields for the standard element this class manages overlaps for. 

525 # This is guaranteed to add a column for the governor dimension, 

526 # because that's a required dependency of element. 

527 for dimension in element.required: 

528 addDimensionForeignKey(overlap_spec, dimension, primaryKey=True) 

529 # Add field for the actual skypix index. We do this later because I 

530 # think we care (at least a bit) about the order in which the primary 

531 # key is defined, in that we want a non-summary column like this one 

532 # to appear after the governor dimension column. 

533 overlap_spec.fields.add( 

534 ddl.FieldSpec( 

535 name="skypix_index", 

536 dtype=sqlalchemy.BigInteger, 

537 nullable=False, 

538 primaryKey=True, 

539 ) 

540 ) 

541 return ( 

542 context.addTable(f"{element.name}_skypix_overlap_summary", summary_spec), 

543 context.addTable(f"{element.name}_skypix_overlap", overlap_spec), 

544 ) 

545 

546 @classmethod 

547 def _make_legacy_overlap_tables( 

548 cls, 

549 context: StaticTablesContext, 

550 spatial: NamedKeyDict[DatabaseTopologicalFamily, list[DimensionElement]], 

551 ) -> None: 

552 for (_, elements1), (_, elements2) in itertools.combinations(spatial.items(), 2): 

553 for element1, element2 in itertools.product(elements1, elements2): 

554 if element1 > element2: 554 ↛ 555line 554 didn't jump to line 555 because the condition on line 554 was never true

555 (element2, element1) = (element1, element2) 

556 assert element1.spatial is not None and element2.spatial is not None 

557 assert element1.governor != element2.governor 

558 assert element1.governor is not None and element2.governor is not None 

559 summary_spec = ddl.TableSpec(fields=[]) 

560 addDimensionForeignKey(summary_spec, element1.governor, primaryKey=True) 

561 addDimensionForeignKey(summary_spec, element2.governor, primaryKey=True) 

562 context.addTable(f"{element1.name}_{element2.name}_overlap_summary", summary_spec) 

563 overlap_spec = ddl.TableSpec(fields=[]) 

564 addDimensionForeignKey(overlap_spec, element1.governor, primaryKey=True) 

565 addDimensionForeignKey(overlap_spec, element2.governor, primaryKey=True) 

566 for dimension in element1.required: 

567 if dimension != element1.governor: 

568 addDimensionForeignKey(overlap_spec, dimension, primaryKey=True) 

569 for dimension in element2.required: 

570 if dimension != element2.governor: 

571 addDimensionForeignKey(overlap_spec, dimension, primaryKey=True) 

572 context.addTable(f"{element1.name}_{element2.name}_overlap", overlap_spec) 

573 

574 def _make_record_db_rows( 

575 self, element: DimensionElement, records: Sequence[DimensionRecord], replace: bool 

576 ) -> _DimensionRecordDatabaseRows: 

577 result = _DimensionRecordDatabaseRows() 

578 result.main_rows = [record.toDict() for record in records] 

579 if element.temporal is not None: 

580 TimespanReprClass = self._db.getTimespanRepresentation() 

581 for row in result.main_rows: 

582 timespan = row.pop("timespan") 

583 TimespanReprClass.update(timespan, result=row) 

584 if element.spatial is not None: 

585 result.overlap_insert_rows = self._compute_common_skypix_overlap_inserts(element, records) 

586 if replace: 

587 result.overlap_delete_rows = self._compute_common_skypix_overlap_deletes(records) 

588 if element in self.universe.governor_dimensions: 

589 for related_element_name in self._overlap_tables.keys(): 

590 if self.universe[related_element_name].governor == element: 

591 result.overlap_summary_rows[related_element_name] = [ 

592 { 

593 "skypix_system": self.universe.commonSkyPix.system.name, 

594 "skypix_level": self.universe.commonSkyPix.level, 

595 element.name: record.dataId[element.name], 

596 } 

597 for record in records 

598 ] 

599 return result 

600 

601 def _compute_common_skypix_overlap_deletes( 

602 self, records: Sequence[DimensionRecord] 

603 ) -> list[dict[str, Any]]: 

604 return [ 

605 { 

606 "skypix_system": self.universe.commonSkyPix.system.name, 

607 "skypix_level": self.universe.commonSkyPix.level, 

608 **record.dataId.required, 

609 } 

610 for record in records 

611 ] 

612 

613 def _compute_common_skypix_overlap_inserts( 

614 self, 

615 element: DimensionElement, 

616 records: Sequence[DimensionRecord], 

617 ) -> list[dict[str, Any]]: 

618 _LOG.debug("Precomputing common skypix overlaps for %s.", element.name) 

619 overlap_records: list[dict[str, Any]] = [] 

620 for record in records: 

621 if record.region is None: 

622 continue 

623 base_overlap_record = dict(record.dataId.required) 

624 base_overlap_record["skypix_system"] = self.universe.commonSkyPix.system.name 

625 base_overlap_record["skypix_level"] = self.universe.commonSkyPix.level 

626 for begin, end in self.universe.commonSkyPix.pixelization.envelope(record.region): 

627 for index in range(begin, end): 

628 overlap_records.append({"skypix_index": index, **base_overlap_record}) 

629 return overlap_records 

630 

631 def _insert_overlaps( 

632 self, 

633 element: DimensionElement, 

634 overlap_insert_rows: list[dict[str, Any]], 

635 overlap_delete_rows: list[dict[str, Any]], 

636 skip_existing: bool = False, 

637 ) -> None: 

638 if overlap_delete_rows: 

639 # Since any of the new records might have replaced existing ones 

640 # that already have overlap records, and we don't know which, we 

641 # have no choice but to delete all overlaps for these records and 

642 # recompute them. We include the skypix_system and skypix_level 

643 # column values explicitly instead of just letting the query search 

644 # for all of those related to the given records, because they are 

645 # the first columns in the primary key, and hence searching with 

646 # them will be way faster (and we don't want to add a new index 

647 # just for this operation). 

648 _LOG.debug("Deleting old common skypix overlaps for %s.", element.name) 

649 self._db.delete( 

650 self._overlap_tables[element.name][1], 

651 ["skypix_system", "skypix_level"] + list(element.minimal_group.required), 

652 *overlap_delete_rows, 

653 ) 

654 if overlap_insert_rows: 

655 _LOG.debug("Inserting %d new skypix overlap rows for %s.", len(overlap_insert_rows), element.name) 

656 if skip_existing: 

657 self._db.ensure( 

658 self._overlap_tables[element.name][1], *overlap_insert_rows, primary_key_only=True 

659 ) 

660 else: 

661 self._db.insert(self._overlap_tables[element.name][1], *overlap_insert_rows) 

662 # We have only ever put overlaps with the commonSkyPix system into 

663 # this table, and *probably* only ever will. But the schema leaves 

664 # open the possibility that we should be inserting overlaps for 

665 # some other skypix system, as we once thought we'd support. In 

666 # case that door opens again in the future, we need to check the 

667 # "overlap summary" table to see if are any skypix systems other 

668 # than the common skypix system and raise (rolling back the entire 

669 # transaction) if there are. 

670 summary_table = self._overlap_tables[element.name][0] 

671 check_sql = ( 

672 sqlalchemy.sql.select(summary_table.columns.skypix_system, summary_table.columns.skypix_level) 

673 .select_from(summary_table) 

674 .where( 

675 sqlalchemy.sql.not_( 

676 sqlalchemy.sql.and_( 

677 summary_table.columns.skypix_system == self.universe.commonSkyPix.system.name, 

678 summary_table.columns.skypix_level == self.universe.commonSkyPix.level, 

679 ) 

680 ) 

681 ) 

682 ) 

683 with self._db.query(check_sql) as sql_result: 

684 bad_summary_rows = sql_result.fetchall() 

685 if bad_summary_rows: 685 ↛ 686line 685 didn't jump to line 686 because the condition on line 685 was never true

686 bad_skypix_names = [f"{row.skypix_system}{row.skypix.level}" for row in bad_summary_rows] 

687 raise RuntimeError( 

688 f"Data repository has overlaps between {element} and {bad_skypix_names} that " 

689 "are not supported by this version of daf_butler. Please use a newer version." 

690 ) 

691 

692 

693@dataclasses.dataclass 

694class _DimensionRecordDatabaseRows: 

695 """Rows to be inserted into the database whenever a DimensionRecord is 

696 added. 

697 """ 

698 

699 main_rows: list[dict[str, Any]] = dataclasses.field(default_factory=list) 

700 """Rows for the dimension element table itself.""" 

701 

702 overlap_insert_rows: list[dict[str, Any]] = dataclasses.field(default_factory=list) 

703 """Rows for overlaps with the common skypix dimension.""" 

704 

705 overlap_delete_rows: list[dict[str, Any]] = dataclasses.field(default_factory=list) 

706 """Rows for overlaps with the common skypix dimension that should be 

707 deleted before inserting new ones. 

708 """ 

709 

710 overlap_summary_rows: dict[str, list[dict[str, Any]]] = dataclasses.field(default_factory=dict) 

711 """Rows that record which overlaps between skypix dimensiosn and other 

712 dimension elements are stored. 

713 

714 This is populated when inserting governor dimension rows, with keys being 

715 the names of spatial dimension elements associated with that governor. 

716 """ 

717 

718 

719class _DimensionGroupStorage: 

720 """Helper object that manages saved DimensionGroup definitions. 

721 

722 Should generally be constructed by calling `initialize` instead of invoking 

723 the constructor directly. 

724 

725 Parameters 

726 ---------- 

727 db : `Database` 

728 Interface to the underlying database engine and namespace. 

729 idTable : `sqlalchemy.schema.Table` 

730 Table that just holds unique IDs for dimension graphs. 

731 definitionTable : `sqlalchemy.schema.Table` 

732 Table that maps dimension names to the IDs of the dimension graphs to 

733 which they belong. 

734 universe : `DimensionUniverse` 

735 All known dimensions. 

736 """ 

737 

738 def __init__( 

739 self, 

740 db: Database, 

741 idTable: sqlalchemy.schema.Table, 

742 definitionTable: sqlalchemy.schema.Table, 

743 universe: DimensionUniverse, 

744 ): 

745 self._db = db 

746 self._idTable = idTable 

747 self._definitionTable = definitionTable 

748 self._universe = universe 

749 self._keysByGroup: dict[DimensionGroup, int] = {universe.empty: 0} 

750 self._groupsByKey: dict[int, DimensionGroup] = {0: universe.empty} 

751 

752 def clone(self, db: Database) -> _DimensionGroupStorage: 

753 """Make an independent copy of this manager instance bound to a new 

754 `Database` instance. 

755 

756 Parameters 

757 ---------- 

758 db : `Database` 

759 New `Database` object to use when instantiating the manager. 

760 

761 Returns 

762 ------- 

763 instance : `_DimensionGroupStorage` 

764 New manager instance with the same configuration as this instance, 

765 but bound to a new Database object. 

766 """ 

767 return _DimensionGroupStorage( 

768 db=db, idTable=self._idTable, definitionTable=self._definitionTable, universe=self._universe 

769 ) 

770 

771 @classmethod 

772 def initialize( 

773 cls, 

774 db: Database, 

775 context: StaticTablesContext, 

776 *, 

777 universe: DimensionUniverse, 

778 ) -> _DimensionGroupStorage: 

779 """Construct a new instance, including creating tables if necessary. 

780 

781 Parameters 

782 ---------- 

783 db : `Database` 

784 Interface to the underlying database engine and namespace. 

785 context : `StaticTablesContext`