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