Coverage for tests/test_postgresql.py: 42%
100 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 18:11 +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/>.
28from __future__ import annotations
30import itertools
31import os
32import unittest
33import warnings
34from contextlib import contextmanager
35from typing import cast
37import astropy.time
38import sqlalchemy
40from lsst.daf.butler import Butler, ButlerConfig, StorageClassFactory, Timespan, ddl
41from lsst.daf.butler.datastore import NullDatastore
42from lsst.daf.butler.direct_butler import DirectButler
43from lsst.daf.butler.registry import RegistryConfig, _RegistryFactory
44from lsst.daf.butler.tests.postgresql import setup_postgres_test_db
46try:
47 from lsst.daf.butler.registry.databases.postgresql import PostgresqlDatabase, _RangeTimespanType
48except ImportError:
49 PostgresqlDatabase = None
50from lsst.daf.butler.registry.tests import DatabaseTests, RegistryTests
52TESTDIR = os.path.abspath(os.path.dirname(__file__))
55@unittest.skipUnless(PostgresqlDatabase is not None, "Couldn't load PostgresqlDatabase")
56class PostgresqlDatabaseTestCase(unittest.TestCase, DatabaseTests):
57 """Test a postgres Registry."""
59 @classmethod
60 def setUpClass(cls):
61 super().setUpClass()
62 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
64 def makeEmptyDatabase(self, origin: int = 0) -> PostgresqlDatabase:
65 db = PostgresqlDatabase.fromUri(
66 origin=origin, uri=self.postgres.url, namespace=self.postgres.generate_namespace_name()
67 )
68 self.addCleanup(db.dispose)
69 return db
71 def getNewConnection(self, database: PostgresqlDatabase, *, writeable: bool) -> PostgresqlDatabase:
72 db = PostgresqlDatabase.fromUri(
73 origin=database.origin, uri=self.postgres.url, namespace=database.namespace, writeable=writeable
74 )
75 self.addCleanup(db.dispose)
76 return db
78 @contextmanager
79 def asReadOnly(self, database: PostgresqlDatabase) -> PostgresqlDatabase:
80 yield self.getNewConnection(database, writeable=False)
82 def testNameShrinking(self):
83 """Test that too-long names for database entities other than tables
84 and columns (which we preserve, and just expect to fit) are shrunk.
85 """
86 db = self.makeEmptyDatabase(origin=1)
87 with db.declareStaticTables(create=True) as context:
88 # Table and field names are each below the 63-char limit even when
89 # accounting for the prefix, but their combination (which will
90 # appear in sequences and constraints) is not.
91 tableName = "a_table_with_a_very_very_long_42_char_name"
92 fieldName1 = "a_column_with_a_very_very_long_43_char_name"
93 fieldName2 = "another_column_with_a_very_very_long_49_char_name"
94 context.addTable(
95 tableName,
96 ddl.TableSpec(
97 fields=[
98 ddl.FieldSpec(
99 fieldName1, dtype=sqlalchemy.BigInteger, autoincrement=True, primaryKey=True
100 ),
101 ddl.FieldSpec(
102 fieldName2,
103 dtype=sqlalchemy.String,
104 length=16,
105 nullable=False,
106 ),
107 ],
108 unique={(fieldName2,)},
109 ),
110 )
111 # Add another table, this time dynamically, with a foreign key to the
112 # first table.
113 db.ensureTableExists(
114 tableName + "_b",
115 ddl.TableSpec(
116 fields=[
117 ddl.FieldSpec(
118 fieldName1 + "_b", dtype=sqlalchemy.BigInteger, autoincrement=True, primaryKey=True
119 ),
120 ddl.FieldSpec(
121 fieldName2 + "_b",
122 dtype=sqlalchemy.String,
123 length=16,
124 nullable=False,
125 ),
126 ],
127 foreignKeys=[
128 ddl.ForeignKeySpec(tableName, source=(fieldName2 + "_b",), target=(fieldName2,)),
129 ],
130 ),
131 )
133 def test_RangeTimespanType(self):
134 start = astropy.time.Time("2020-01-01T00:00:00", format="isot", scale="tai")
135 offset = astropy.time.TimeDelta(60, format="sec")
136 timestamps = [start + offset * n for n in range(3)]
137 timespans = [Timespan(begin=None, end=None)]
138 timespans.extend(Timespan(begin=None, end=t) for t in timestamps)
139 timespans.extend(Timespan(begin=t, end=None) for t in timestamps)
140 timespans.extend(Timespan(begin=a, end=b) for a, b in itertools.combinations(timestamps, 2))
141 db = self.makeEmptyDatabase(origin=1)
142 with db.declareStaticTables(create=True) as context:
143 tbl = context.addTable(
144 "tbl",
145 ddl.TableSpec(
146 fields=[
147 ddl.FieldSpec(name="id", dtype=sqlalchemy.Integer, primaryKey=True),
148 ddl.FieldSpec(name="timespan", dtype=_RangeTimespanType),
149 ],
150 ),
151 )
152 rows = [{"id": n, "timespan": t} for n, t in enumerate(timespans)]
153 db.insert(tbl, *rows)
155 # Test basic round-trip through database.
156 with db.query(tbl.select().order_by(tbl.columns.id)) as sql_result:
157 self.assertEqual(rows, [row._asdict() for row in sql_result])
159 # Test that Timespan's Python methods are consistent with our usage of
160 # half-open ranges and PostgreSQL operators on ranges.
161 def subquery(alias: str) -> sqlalchemy.sql.FromClause:
162 return (
163 sqlalchemy.sql.select(tbl.columns.id.label("id"), tbl.columns.timespan.label("timespan"))
164 .select_from(tbl)
165 .alias(alias)
166 )
168 sq1 = subquery("sq1")
169 sq2 = subquery("sq2")
170 query = sqlalchemy.sql.select(
171 sq1.columns.id.label("n1"),
172 sq2.columns.id.label("n2"),
173 sq1.columns.timespan.overlaps(sq2.columns.timespan).label("overlaps"),
174 )
176 # `columns` is deprecated since 1.4, but
177 # `selected_columns` method did not exist in 1.3.
178 if hasattr(query, "selected_columns"):
179 columns = query.selected_columns
180 else:
181 columns = query.columns
183 # SQLAlchemy issues a warning about cartesian product of two tables,
184 # which we do intentionally. Disable that warning temporarily.
185 with warnings.catch_warnings():
186 warnings.filterwarnings(
187 "ignore", message=".*cartesian product", category=sqlalchemy.exc.SAWarning
188 )
189 with db.query(query) as sql_result:
190 dbResults = {
191 (row[columns.n1], row[columns.n2]): row[columns.overlaps] for row in sql_result.mappings()
192 }
194 pyResults = {
195 (n1, n2): t1.overlaps(t2)
196 for (n1, t1), (n2, t2) in itertools.product(enumerate(timespans), enumerate(timespans))
197 }
198 self.assertEqual(pyResults, dbResults)
201class PostgresqlRegistryTests(RegistryTests):
202 """Tests for `Registry` backed by a PostgreSQL database.
204 Notes
205 -----
206 This is not a subclass of `unittest.TestCase` but to avoid repetition it
207 defines methods that override `unittest.TestCase` methods. To make this
208 work subclasses have to have this class first in the bases list.
209 """
211 @classmethod
212 def setUpClass(cls):
213 super().setUpClass()
214 cls.postgres = cls.enterClassContext(setup_postgres_test_db())
216 @classmethod
217 def getDataDir(cls) -> str:
218 return os.path.normpath(os.path.join(os.path.dirname(__file__), "data", "registry"))
220 def make_butler(self, config: RegistryConfig | None = None) -> Butler:
221 if config is None:
222 config = self.makeRegistryConfig()
223 self.postgres.patch_registry_config(config)
224 registry = _RegistryFactory(config).create_from_config()
226 butler = DirectButler(
227 config=ButlerConfig(),
228 registry=registry,
229 datastore=NullDatastore(None, None),
230 storageClasses=StorageClassFactory(),
231 )
232 cast(unittest.TestCase, self).enterContext(butler)
234 return butler
236 def testSkipCalibs(self):
237 if self.postgres.server_major_version() < 16:
238 # TODO DM-44875: This test currently fails for older Postgres.
239 self.skipTest("TODO DM-44875")
240 return super().testSkipCalibs()
243class PostgresqlRegistryNameKeyCollMgrUUIDTestCase(PostgresqlRegistryTests, unittest.TestCase):
244 """Tests for `Registry` backed by a PostgreSQL database.
246 This test case uses NameKeyCollectionManager and
247 ByDimensionsDatasetRecordStorageManagerUUID.
248 """
250 collectionsManager = "lsst.daf.butler.registry.collections.nameKey.NameKeyCollectionManager"
251 datasetsManager = (
252 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
253 )
256class PostgresqlRegistrySynthIntKeyCollMgrUUIDTestCase(PostgresqlRegistryTests, unittest.TestCase):
257 """Tests for `Registry` backed by a PostgreSQL database.
259 This test case uses SynthIntKeyCollectionManager and
260 ByDimensionsDatasetRecordStorageManagerUUID.
261 """
263 collectionsManager = "lsst.daf.butler.registry.collections.synthIntKey.SynthIntKeyCollectionManager"
264 datasetsManager = (
265 "lsst.daf.butler.registry.datasets.byDimensions.ByDimensionsDatasetRecordStorageManagerUUID"
266 )
269if __name__ == "__main__":
270 unittest.main()