Coverage for python/lsst/daf/butler/registry/collections/synthIntKey.py: 99%
113 statements
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-12 10:05 +0000
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-12 10:05 +0000
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
29from ... import ddl
31__all__ = ["SynthIntKeyCollectionManager"]
33import logging
34from collections.abc import Iterable, Mapping
35from typing import TYPE_CHECKING, Any
37import sqlalchemy
39from ...column_spec import COLLECTION_NAME_MAX_LENGTH
40from ...timespan_database_representation import TimespanDatabaseRepresentation
41from .._collection_type import CollectionType
42from ..interfaces import ChainedCollectionRecord, CollectionRecord, RunRecord, VersionTuple
43from ._base import (
44 CollectionTablesTuple,
45 DefaultCollectionManager,
46 makeCollectionChainTableSpec,
47 makeRunTableSpec,
48)
50if TYPE_CHECKING:
51 from .._caching_context import CachingContext
52 from ..interfaces import Database, StaticTablesContext
55_KEY_FIELD_SPEC = ddl.FieldSpec(
56 "collection_id", dtype=sqlalchemy.BigInteger, primaryKey=True, autoincrement=True
57)
60# This has to be updated on every schema change
61_VERSION = VersionTuple(2, 0, 0)
63_LOG = logging.getLogger(__name__)
66def _makeTableSpecs(TimespanReprClass: type[TimespanDatabaseRepresentation]) -> CollectionTablesTuple:
67 return CollectionTablesTuple(
68 collection=ddl.TableSpec(
69 fields=[
70 _KEY_FIELD_SPEC,
71 ddl.FieldSpec(
72 "name", dtype=sqlalchemy.String, length=COLLECTION_NAME_MAX_LENGTH, nullable=False
73 ),
74 ddl.FieldSpec("type", dtype=sqlalchemy.SmallInteger, nullable=False),
75 ddl.FieldSpec("doc", dtype=sqlalchemy.Text, nullable=True),
76 ],
77 unique=[("name",)],
78 ),
79 run=makeRunTableSpec("collection_id", sqlalchemy.BigInteger, TimespanReprClass),
80 collection_chain=makeCollectionChainTableSpec("collection_id", sqlalchemy.BigInteger),
81 )
84class SynthIntKeyCollectionManager(DefaultCollectionManager[int]):
85 """A `CollectionManager` implementation that uses synthetic primary key
86 (auto-incremented integer) for collections table.
87 """
89 @classmethod
90 def initialize(
91 cls,
92 db: Database,
93 context: StaticTablesContext,
94 *,
95 caching_context: CachingContext,
96 registry_schema_version: VersionTuple | None = None,
97 ) -> SynthIntKeyCollectionManager:
98 # Docstring inherited from CollectionManager.
99 return cls(
100 db,
101 tables=context.addTableTuple(_makeTableSpecs(db.getTimespanRepresentation())), # type: ignore
102 collectionIdName="collection_id",
103 caching_context=caching_context,
104 registry_schema_version=registry_schema_version,
105 )
107 def clone(self, db: Database, caching_context: CachingContext) -> SynthIntKeyCollectionManager:
108 return SynthIntKeyCollectionManager(
109 db,
110 tables=self._tables,
111 collectionIdName=self._collectionIdName,
112 caching_context=caching_context,
113 registry_schema_version=self._registry_schema_version,
114 )
116 @classmethod
117 def getCollectionForeignKeyName(cls, prefix: str = "collection") -> str:
118 # Docstring inherited from CollectionManager.
119 return f"{prefix}_id"
121 @classmethod
122 def getRunForeignKeyName(cls, prefix: str = "run") -> str:
123 # Docstring inherited from CollectionManager.
124 return f"{prefix}_id"
126 @classmethod
127 def addCollectionForeignKey(
128 cls,
129 tableSpec: ddl.TableSpec,
130 *,
131 prefix: str = "collection",
132 onDelete: str | None = None,
133 constraint: bool = True,
134 **kwargs: Any,
135 ) -> ddl.FieldSpec:
136 # Docstring inherited from CollectionManager.
137 original = _KEY_FIELD_SPEC
138 copy = ddl.FieldSpec(
139 cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
140 )
141 tableSpec.fields.add(copy)
142 if constraint:
143 tableSpec.foreignKeys.append(
144 ddl.ForeignKeySpec(
145 "collection", source=(copy.name,), target=(original.name,), onDelete=onDelete
146 )
147 )
148 return copy
150 @classmethod
151 def addRunForeignKey(
152 cls,
153 tableSpec: ddl.TableSpec,
154 *,
155 prefix: str = "run",
156 onDelete: str | None = None,
157 constraint: bool = True,
158 **kwargs: Any,
159 ) -> ddl.FieldSpec:
160 # Docstring inherited from CollectionManager.
161 original = _KEY_FIELD_SPEC
162 copy = ddl.FieldSpec(
163 cls.getRunForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
164 )
165 tableSpec.fields.add(copy)
166 if constraint: 166 ↛ 170line 166 didn't jump to line 170, because the condition on line 166 was never false
167 tableSpec.foreignKeys.append(
168 ddl.ForeignKeySpec("run", source=(copy.name,), target=(original.name,), onDelete=onDelete)
169 )
170 return copy
172 def getParentChains(self, key: int) -> set[str]:
173 # Docstring inherited from CollectionManager.
174 chain = self._tables.collection_chain
175 collection = self._tables.collection
176 sql = (
177 sqlalchemy.sql.select(collection.columns["name"])
178 .select_from(collection)
179 .join(chain, onclause=collection.columns[self._collectionIdName] == chain.columns["parent"])
180 .where(chain.columns["child"] == key)
181 )
182 with self._db.query(sql) as sql_result:
183 parent_names = set(sql_result.scalars().all())
184 return parent_names
186 def _fetch_by_name(self, names: Iterable[str]) -> list[CollectionRecord[int]]:
187 # Docstring inherited from base class.
188 _LOG.debug("Fetching collection records using names %s.", names)
189 return self._fetch("name", names)
191 def _fetch_by_key(self, collection_ids: Iterable[int] | None) -> list[CollectionRecord[int]]:
192 # Docstring inherited from base class.
193 _LOG.debug("Fetching collection records using IDs %s.", collection_ids)
194 return self._fetch(self._collectionIdName, collection_ids)
196 def _fetch(
197 self, column_name: str, collections: Iterable[int | str] | None
198 ) -> list[CollectionRecord[int]]:
199 collection_chain = self._tables.collection_chain
200 collection = self._tables.collection
201 sql = sqlalchemy.sql.select(*collection.columns, *self._tables.run.columns).select_from(
202 collection.join(self._tables.run, isouter=True)
203 )
205 chain_sql = (
206 sqlalchemy.sql.select(
207 collection_chain.columns["parent"],
208 collection_chain.columns["position"],
209 collection.columns["name"].label("child_name"),
210 )
211 .select_from(collection_chain)
212 .join(
213 collection,
214 onclause=collection_chain.columns["child"] == collection.columns[self._collectionIdName],
215 )
216 )
218 records: list[CollectionRecord[int]] = []
219 # We want to keep transactions as short as possible. When we fetch
220 # everything we want to quickly fetch things into memory and finish
221 # transaction. When we fetch just few records we need to process first
222 # query before wi can run second one,
223 if collections is not None:
224 sql = sql.where(collection.columns[column_name].in_(collections))
225 with self._db.transaction():
226 with self._db.query(sql) as sql_result:
227 sql_rows = sql_result.mappings().fetchall()
229 records, chained_ids = self._rows_to_records(sql_rows)
231 if chained_ids:
232 chain_sql = chain_sql.where(collection_chain.columns["parent"].in_(list(chained_ids)))
234 with self._db.query(chain_sql) as sql_result:
235 chain_rows = sql_result.mappings().fetchall()
237 records += self._rows_to_chains(chain_rows, chained_ids)
239 else:
240 with self._db.transaction():
241 with self._db.query(sql) as sql_result:
242 sql_rows = sql_result.mappings().fetchall()
243 with self._db.query(chain_sql) as sql_result:
244 chain_rows = sql_result.mappings().fetchall()
246 records, chained_ids = self._rows_to_records(sql_rows)
247 records += self._rows_to_chains(chain_rows, chained_ids)
249 return records
251 def _rows_to_records(self, rows: Iterable[Mapping]) -> tuple[list[CollectionRecord[int]], dict[int, str]]:
252 """Convert rows returned from collection query to a list of records
253 and a dict chained collection names.
254 """
255 records: list[CollectionRecord[int]] = []
256 chained_ids: dict[int, str] = {}
257 TimespanReprClass = self._db.getTimespanRepresentation()
258 for row in rows:
259 key: int = row[self._collectionIdName]
260 name: str = row[self._tables.collection.columns.name]
261 type = CollectionType(row["type"])
262 record: CollectionRecord[int]
263 if type is CollectionType.RUN:
264 record = RunRecord[int](
265 key=key,
266 name=name,
267 host=row[self._tables.run.columns.host],
268 timespan=TimespanReprClass.extract(row),
269 )
270 records.append(record)
271 elif type is CollectionType.CHAINED:
272 # Need to delay chained collection construction until to
273 # fetch their children names.
274 chained_ids[key] = name
275 else:
276 record = CollectionRecord[int](key=key, name=name, type=type)
277 records.append(record)
278 return records, chained_ids
280 def _rows_to_chains(
281 self, rows: Iterable[Mapping], chained_ids: dict[int, str]
282 ) -> list[CollectionRecord[int]]:
283 """Convert rows returned from collection chain query to a list of
284 records.
285 """
286 chains_defs: dict[int, list[tuple[int, str]]] = {chain_id: [] for chain_id in chained_ids}
287 for row in rows:
288 chains_defs[row["parent"]].append((row["position"], row["child_name"]))
290 records: list[CollectionRecord[int]] = []
291 for key, children in chains_defs.items():
292 name = chained_ids[key]
293 children_names = [child for _, child in sorted(children)]
294 record = ChainedCollectionRecord[int](
295 key=key,
296 name=name,
297 children=children_names,
298 )
299 records.append(record)
301 return records
303 @classmethod
304 def currentVersions(cls) -> list[VersionTuple]:
305 # Docstring inherited from VersionedExtension.
306 return [_VERSION]