Coverage for python/lsst/daf/butler/registry/dimensions/spatial.py : 88%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 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/>.
21from __future__ import annotations
23__all__ = ["SpatialDimensionRecordStorage"]
25from typing import List, Optional
27import sqlalchemy
29from ...core import (
30 addDimensionForeignKey,
31 DatabaseTimespanRepresentation,
32 ddl,
33 DimensionElement,
34 DimensionRecord,
35 NamedKeyDict,
36 NamedValueSet,
37 REGION_FIELD_SPEC,
38)
39from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext
40from ..queries import QueryBuilder
41from .table import TableDimensionRecordStorage
44_OVERLAP_TABLE_NAME_PATTERN = "{0}_{1}_overlap"
47def _makeOverlapTableSpec(a: DimensionElement, b: DimensionElement) -> ddl.TableSpec:
48 """Create a specification for a table that represents a many-to-many
49 relationship between two `DimensionElement` tables.
51 Parameters
52 ----------
53 a : `DimensionElement`
54 First element in the relationship.
55 b : `DimensionElement`
56 Second element in the relationship.
58 Returns
59 -------
60 spec : `TableSpec`
61 Database-agnostic specification for a table.
62 """
63 tableSpec = ddl.TableSpec(
64 fields=NamedValueSet(),
65 unique=set(),
66 foreignKeys=[],
67 )
68 for dimension in a.required:
69 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
70 for dimension in b.required:
71 addDimensionForeignKey(tableSpec, dimension, primaryKey=True)
72 return tableSpec
75class SpatialDimensionRecordStorage(TableDimensionRecordStorage):
76 """A record storage implementation for spatial dimension elements that uses
77 a regular database table.
79 Parameters
80 ----------
81 db : `Database`
82 Interface to the database engine and namespace that will hold these
83 dimension records.
84 element : `DimensionElement`
85 The element whose records this storage will manage.
86 table : `sqlalchemy.schema.Table`
87 The logical table for the element.
88 commonSkyPixOverlapTable : `sqlalchemy.schema.Table`, optional
89 The logical table for the overlap table with the dimension universe's
90 common skypix dimension.
91 """
92 def __init__(self, db: Database, element: DimensionElement, *, table: sqlalchemy.schema.Table,
93 commonSkyPixOverlapTable: sqlalchemy.schema.Table):
94 super().__init__(db, element, table=table)
95 self._commonSkyPixOverlapTable = commonSkyPixOverlapTable
96 assert element.spatial is not None
98 @classmethod
99 def initialize(cls, db: Database, element: DimensionElement, *,
100 context: Optional[StaticTablesContext] = None) -> DimensionRecordStorage:
101 # Docstring inherited from DimensionRecordStorage.
102 if context is not None: 102 ↛ 105line 102 didn't jump to line 105, because the condition on line 102 was never false
103 method = context.addTable
104 else:
105 method = db.ensureTableExists
106 return cls(
107 db,
108 element,
109 table=method(
110 element.name,
111 element.RecordClass.fields.makeTableSpec(tsRepr=db.getTimespanRepresentation())
112 ),
113 commonSkyPixOverlapTable=method(
114 _OVERLAP_TABLE_NAME_PATTERN.format(element.name, element.universe.commonSkyPix.name),
115 _makeOverlapTableSpec(element, element.universe.commonSkyPix)
116 )
117 )
119 def join(
120 self,
121 builder: QueryBuilder, *,
122 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
123 timespans: Optional[NamedKeyDict[DimensionElement, DatabaseTimespanRepresentation]] = None,
124 ) -> None:
125 # Docstring inherited from DimensionRecordStorage.
126 if regions is not None:
127 dimensions = NamedValueSet(self.element.required)
128 dimensions.add(self.element.universe.commonSkyPix)
129 builder.joinTable(self._commonSkyPixOverlapTable, dimensions)
130 regions[self.element] = self._table.columns[REGION_FIELD_SPEC.name]
131 return super().join(builder, regions=None, timespans=timespans)
133 def _computeCommonSkyPixRows(self, *records: DimensionRecord) -> List[dict]:
134 commonSkyPixRows = []
135 commonSkyPix = self.element.universe.commonSkyPix
136 for record in records:
137 # MyPy can't tell that some DimensionRecords have regions, because
138 # they're dynamically-created types.
139 region = record.region # type: ignore
140 if region is None:
141 # TODO: should we warn about this case?
142 continue
143 base = record.dataId.byName()
144 for begin, end in commonSkyPix.pixelization.envelope(region):
145 for skypix in range(begin, end):
146 row = base.copy()
147 row[commonSkyPix.name] = skypix
148 commonSkyPixRows.append(row)
149 return commonSkyPixRows
151 def insert(self, *records: DimensionRecord) -> None:
152 # Docstring inherited from DimensionRecordStorage.insert.
153 commonSkyPixRows = self._computeCommonSkyPixRows(*records)
154 with self._db.transaction():
155 super().insert(*records)
156 if commonSkyPixRows:
157 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
159 def sync(self, record: DimensionRecord) -> bool:
160 # Docstring inherited from DimensionRecordStorage.sync.
161 with self._db.transaction():
162 inserted = super().sync(record)
163 if inserted:
164 commonSkyPixRows = self._computeCommonSkyPixRows(record)
165 self._db.insert(self._commonSkyPixOverlapTable, *commonSkyPixRows)
166 return inserted