Coverage for python/lsst/daf/butler/registry/dimensions/table.py : 95%

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__ = ["TableDimensionRecordStorage"]
25import itertools
26from typing import Dict, Iterable, Optional
28import sqlalchemy
30from ...core import (
31 DataCoordinateIterable,
32 DimensionElement,
33 DimensionRecord,
34 makeDimensionElementTableSpec,
35 NamedKeyDict,
36 SimpleQuery,
37 Timespan,
38 TIMESPAN_FIELD_SPECS,
39)
40from ..interfaces import Database, DimensionRecordStorage, StaticTablesContext
41from ..queries import QueryBuilder
44MAX_FETCH_CHUNK = 1000
45"""Maximum number of data IDs we fetch records at a time.
47Barring something database-engine-specific, this sets the size of the actual
48SQL query, not just the number of result rows, because the only way to query
49for multiple data IDs in a single SELECT query via SQLAlchemy is to have an OR
50term in the WHERE clause for each one.
51"""
54class TableDimensionRecordStorage(DimensionRecordStorage):
55 """A record storage implementation uses a regular database table.
57 For spatial dimension elements, use `SpatialDimensionRecordStorage`
58 instead.
60 Parameters
61 ----------
62 db : `Database`
63 Interface to the database engine and namespace that will hold these
64 dimension records.
65 element : `DimensionElement`
66 The element whose records this storage will manage.
67 table : `sqlalchemy.schema.Table`
68 The logical table for the element.
69 """
70 def __init__(self, db: Database, element: DimensionElement, *, table: sqlalchemy.schema.Table):
71 self._db = db
72 self._table = table
73 self._element = element
74 self._fetchColumns: Dict[str, sqlalchemy.sql.ColumnElement] = {
75 dimension.name: self._table.columns[name]
76 for dimension, name in zip(itertools.chain(self._element.required, self._element.implied),
77 self._element.RecordClass.__slots__)
78 }
80 @classmethod
81 def initialize(cls, db: Database, element: DimensionElement, *,
82 context: Optional[StaticTablesContext] = None) -> DimensionRecordStorage:
83 # Docstring inherited from DimensionRecordStorage.
84 spec = makeDimensionElementTableSpec(element)
85 if context is not None: 85 ↛ 88line 85 didn't jump to line 88, because the condition on line 85 was never false
86 table = context.addTable(element.name, spec)
87 else:
88 table = db.ensureTableExists(element.name, spec)
89 return cls(db, element, table=table)
91 @property
92 def element(self) -> DimensionElement:
93 # Docstring inherited from DimensionRecordStorage.element.
94 return self._element
96 def clearCaches(self) -> None:
97 # Docstring inherited from DimensionRecordStorage.clearCaches.
98 pass
100 def join(
101 self,
102 builder: QueryBuilder, *,
103 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
104 timespans: Optional[NamedKeyDict[DimensionElement, Timespan[sqlalchemy.sql.ColumnElement]]] = None,
105 ) -> None:
106 # Docstring inherited from DimensionRecordStorage.
107 assert regions is None, "This implementation does not handle spatial joins."
108 joinDimensions = list(self.element.required)
109 joinDimensions.extend(self.element.implied)
110 joinOn = builder.startJoin(self._table, joinDimensions, self.element.RecordClass.__slots__)
111 if timespans is not None:
112 timespanInTable: Timespan[sqlalchemy.sql.ColumnElement] = Timespan(
113 begin=self._table.columns[TIMESPAN_FIELD_SPECS.begin.name],
114 end=self._table.columns[TIMESPAN_FIELD_SPECS.end.name],
115 )
116 for timespanInQuery in timespans.values(): 116 ↛ 117line 116 didn't jump to line 117, because the loop on line 116 never started
117 joinOn.append(timespanInQuery.overlaps(timespanInTable, ops=sqlalchemy.sql))
118 timespans[self.element] = timespanInTable
119 builder.finishJoin(self._table, joinOn)
120 return self._table
122 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
123 # Docstring inherited from DimensionRecordStorage.fetch.
124 RecordClass = self.element.RecordClass
125 query = SimpleQuery()
126 query.columns.extend(self._table.columns[name] for name in RecordClass.__slots__)
127 query.join(self._table)
128 dataIds.constrain(query, lambda name: self._fetchColumns[name])
129 for row in self._db.query(query.combine()):
130 yield RecordClass(*row)
132 def insert(self, *records: DimensionRecord) -> None:
133 # Docstring inherited from DimensionRecordStorage.insert.
134 elementRows = [record.toDict() for record in records]
135 with self._db.transaction():
136 self._db.insert(self._table, *elementRows)
138 def sync(self, record: DimensionRecord) -> bool:
139 # Docstring inherited from DimensionRecordStorage.sync.
140 n = len(self.element.required)
141 _, inserted = self._db.sync(
142 self._table,
143 keys={k: getattr(record, k) for k in record.__slots__[:n]},
144 compared={k: getattr(record, k) for k in record.__slots__[n:]},
145 )
146 return inserted
148 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
149 # Docstring inherited from DimensionRecordStorage.digestTables.
150 return [self._table]