Coverage for python/lsst/daf/butler/registry/collections/synthIntKey.py: 98%
58 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-03-30 02:32 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2023-03-30 02:32 -0700
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__ = ["SynthIntKeyCollectionManager"]
25from collections.abc import Iterable
26from typing import TYPE_CHECKING, Any
28import sqlalchemy
30from ...core import TimespanDatabaseRepresentation, ddl
31from ..interfaces import CollectionRecord, VersionTuple
32from ._base import (
33 CollectionTablesTuple,
34 DefaultCollectionManager,
35 makeCollectionChainTableSpec,
36 makeRunTableSpec,
37)
39if TYPE_CHECKING:
40 from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext
43_KEY_FIELD_SPEC = ddl.FieldSpec(
44 "collection_id", dtype=sqlalchemy.BigInteger, primaryKey=True, autoincrement=True
45)
48# This has to be updated on every schema change
49_VERSION = VersionTuple(2, 0, 0)
52def _makeTableSpecs(TimespanReprClass: type[TimespanDatabaseRepresentation]) -> CollectionTablesTuple:
53 return CollectionTablesTuple(
54 collection=ddl.TableSpec(
55 fields=[
56 _KEY_FIELD_SPEC,
57 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=64, nullable=False),
58 ddl.FieldSpec("type", dtype=sqlalchemy.SmallInteger, nullable=False),
59 ddl.FieldSpec("doc", dtype=sqlalchemy.Text, nullable=True),
60 ],
61 unique=[("name",)],
62 ),
63 run=makeRunTableSpec("collection_id", sqlalchemy.BigInteger, TimespanReprClass),
64 collection_chain=makeCollectionChainTableSpec("collection_id", sqlalchemy.BigInteger),
65 )
68class SynthIntKeyCollectionManager(DefaultCollectionManager):
69 """A `CollectionManager` implementation that uses synthetic primary key
70 (auto-incremented integer) for collections table.
72 Most of the logic, including caching policy, is implemented in the base
73 class, this class only adds customizations specific to this particular
74 table schema.
76 Parameters
77 ----------
78 db : `Database`
79 Interface to the underlying database engine and namespace.
80 tables : `NamedTuple`
81 Named tuple of SQLAlchemy table objects.
82 collectionIdName : `str`
83 Name of the column in collections table that identifies it (PK).
84 dimensions : `DimensionRecordStorageManager`
85 Manager object for the dimensions in this `Registry`.
86 """
88 def __init__(
89 self,
90 db: Database,
91 tables: CollectionTablesTuple,
92 collectionIdName: str,
93 dimensions: DimensionRecordStorageManager,
94 registry_schema_version: VersionTuple | None = None,
95 ):
96 super().__init__(
97 db=db,
98 tables=tables,
99 collectionIdName=collectionIdName,
100 dimensions=dimensions,
101 registry_schema_version=registry_schema_version,
102 )
103 self._nameCache: dict[str, CollectionRecord] = {} # indexed by collection name
105 @classmethod
106 def initialize(
107 cls,
108 db: Database,
109 context: StaticTablesContext,
110 *,
111 dimensions: DimensionRecordStorageManager,
112 registry_schema_version: VersionTuple | None = None,
113 ) -> SynthIntKeyCollectionManager:
114 # Docstring inherited from CollectionManager.
115 return cls(
116 db,
117 tables=context.addTableTuple(_makeTableSpecs(db.getTimespanRepresentation())), # type: ignore
118 collectionIdName="collection_id",
119 dimensions=dimensions,
120 registry_schema_version=registry_schema_version,
121 )
123 @classmethod
124 def getCollectionForeignKeyName(cls, prefix: str = "collection") -> str:
125 # Docstring inherited from CollectionManager.
126 return f"{prefix}_id"
128 @classmethod
129 def getRunForeignKeyName(cls, prefix: str = "run") -> str:
130 # Docstring inherited from CollectionManager.
131 return f"{prefix}_id"
133 @classmethod
134 def addCollectionForeignKey(
135 cls,
136 tableSpec: ddl.TableSpec,
137 *,
138 prefix: str = "collection",
139 onDelete: str | None = None,
140 constraint: bool = True,
141 **kwargs: Any,
142 ) -> ddl.FieldSpec:
143 # Docstring inherited from CollectionManager.
144 original = _KEY_FIELD_SPEC
145 copy = ddl.FieldSpec(
146 cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
147 )
148 tableSpec.fields.add(copy)
149 if constraint:
150 tableSpec.foreignKeys.append(
151 ddl.ForeignKeySpec(
152 "collection", source=(copy.name,), target=(original.name,), onDelete=onDelete
153 )
154 )
155 return copy
157 @classmethod
158 def addRunForeignKey(
159 cls,
160 tableSpec: ddl.TableSpec,
161 *,
162 prefix: str = "run",
163 onDelete: str | None = None,
164 constraint: bool = True,
165 **kwargs: Any,
166 ) -> ddl.FieldSpec:
167 # Docstring inherited from CollectionManager.
168 original = _KEY_FIELD_SPEC
169 copy = ddl.FieldSpec(
170 cls.getRunForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
171 )
172 tableSpec.fields.add(copy)
173 if constraint: 173 ↛ 177line 173 didn't jump to line 177, because the condition on line 173 was never false
174 tableSpec.foreignKeys.append(
175 ddl.ForeignKeySpec("run", source=(copy.name,), target=(original.name,), onDelete=onDelete)
176 )
177 return copy
179 def _setRecordCache(self, records: Iterable[CollectionRecord]) -> None:
180 """Set internal record cache to contain given records,
181 old cached records will be removed.
182 """
183 self._records = {}
184 self._nameCache = {}
185 for record in records:
186 self._records[record.key] = record
187 self._nameCache[record.name] = record
189 def _addCachedRecord(self, record: CollectionRecord) -> None:
190 """Add single record to cache."""
191 self._records[record.key] = record
192 self._nameCache[record.name] = record
194 def _removeCachedRecord(self, record: CollectionRecord) -> None:
195 """Remove single record from cache."""
196 del self._records[record.key]
197 del self._nameCache[record.name]
199 def _getByName(self, name: str) -> CollectionRecord | None:
200 # Docstring inherited from DefaultCollectionManager.
201 return self._nameCache.get(name)
203 @classmethod
204 def currentVersions(cls) -> list[VersionTuple]:
205 # Docstring inherited from VersionedExtension.
206 return [_VERSION]