Coverage for python/lsst/daf/butler/registry/collections/synthIntKey.py: 97%
62 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-15 02:03 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-15 02:03 -0800
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: 39 ↛ 40line 39 didn't jump to line 40, because the condition on line 39 was never true
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 ):
95 super().__init__(db=db, tables=tables, collectionIdName=collectionIdName, dimensions=dimensions)
96 self._nameCache: dict[str, CollectionRecord] = {} # indexed by collection name
98 @classmethod
99 def initialize(
100 cls,
101 db: Database,
102 context: StaticTablesContext,
103 *,
104 dimensions: DimensionRecordStorageManager,
105 ) -> SynthIntKeyCollectionManager:
106 # Docstring inherited from CollectionManager.
107 return cls(
108 db,
109 tables=context.addTableTuple(_makeTableSpecs(db.getTimespanRepresentation())), # type: ignore
110 collectionIdName="collection_id",
111 dimensions=dimensions,
112 )
114 @classmethod
115 def getCollectionForeignKeyName(cls, prefix: str = "collection") -> str:
116 # Docstring inherited from CollectionManager.
117 return f"{prefix}_id"
119 @classmethod
120 def getRunForeignKeyName(cls, prefix: str = "run") -> str:
121 # Docstring inherited from CollectionManager.
122 return f"{prefix}_id"
124 @classmethod
125 def addCollectionForeignKey(
126 cls,
127 tableSpec: ddl.TableSpec,
128 *,
129 prefix: str = "collection",
130 onDelete: str | None = None,
131 constraint: bool = True,
132 **kwargs: Any,
133 ) -> ddl.FieldSpec:
134 # Docstring inherited from CollectionManager.
135 original = _KEY_FIELD_SPEC
136 copy = ddl.FieldSpec(
137 cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
138 )
139 tableSpec.fields.add(copy)
140 if constraint:
141 tableSpec.foreignKeys.append(
142 ddl.ForeignKeySpec(
143 "collection", source=(copy.name,), target=(original.name,), onDelete=onDelete
144 )
145 )
146 return copy
148 @classmethod
149 def addRunForeignKey(
150 cls,
151 tableSpec: ddl.TableSpec,
152 *,
153 prefix: str = "run",
154 onDelete: str | None = None,
155 constraint: bool = True,
156 **kwargs: Any,
157 ) -> ddl.FieldSpec:
158 # Docstring inherited from CollectionManager.
159 original = _KEY_FIELD_SPEC
160 copy = ddl.FieldSpec(
161 cls.getRunForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs
162 )
163 tableSpec.fields.add(copy)
164 if constraint:
165 tableSpec.foreignKeys.append(
166 ddl.ForeignKeySpec("run", source=(copy.name,), target=(original.name,), onDelete=onDelete)
167 )
168 return copy
170 def _setRecordCache(self, records: Iterable[CollectionRecord]) -> None:
171 """Set internal record cache to contain given records,
172 old cached records will be removed.
173 """
174 self._records = {}
175 self._nameCache = {}
176 for record in records:
177 self._records[record.key] = record
178 self._nameCache[record.name] = record
180 def _addCachedRecord(self, record: CollectionRecord) -> None:
181 """Add single record to cache."""
182 self._records[record.key] = record
183 self._nameCache[record.name] = record
185 def _removeCachedRecord(self, record: CollectionRecord) -> None:
186 """Remove single record from cache."""
187 del self._records[record.key]
188 del self._nameCache[record.name]
190 def _getByName(self, name: str) -> CollectionRecord | None:
191 # Docstring inherited from DefaultCollectionManager.
192 return self._nameCache.get(name)
194 @classmethod
195 def currentVersion(cls) -> VersionTuple | None:
196 # Docstring inherited from VersionedExtension.
197 return _VERSION
199 def schemaDigest(self) -> str | None:
200 # Docstring inherited from VersionedExtension.
201 return self._defaultSchemaDigest(self._tables, self._db.dialect)