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