Coverage for python/lsst/daf/butler/registry/tests/_database.py: 87%
527 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/>.
27from __future__ import annotations
29from ... import ddl
31__all__ = ["DatabaseTests"]
33import asyncio
34import itertools
35import warnings
36from abc import ABC, abstractmethod
37from collections import namedtuple
38from collections.abc import Iterable
39from concurrent.futures import ThreadPoolExecutor
40from contextlib import AbstractContextManager, contextmanager
41from typing import Any
43import astropy.time
44import sqlalchemy
45import sqlalchemy.dialects.postgresql
47from lsst.sphgeom import Circle, ConvexPolygon, Mq3cPixelization, UnionRegion, UnitVector3d
49from ..._timespan import Timespan
50from ...tests.utils import TestCaseMixin
51from ..interfaces import Database, DatabaseConflictError, ReadOnlyDatabaseError, SchemaAlreadyDefinedError
53StaticTablesTuple = namedtuple("StaticTablesTuple", ["a", "b", "c"])
55STATIC_TABLE_SPECS = StaticTablesTuple(
56 a=ddl.TableSpec(
57 fields=[
58 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=16, primaryKey=True),
59 ddl.FieldSpec("region", dtype=ddl.Base64Region, nbytes=128),
60 ]
61 ),
62 b=ddl.TableSpec(
63 fields=[
64 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, autoincrement=True, primaryKey=True),
65 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=16, nullable=False),
66 ddl.FieldSpec("value", dtype=sqlalchemy.SmallInteger, nullable=True),
67 ],
68 unique=[("name",)],
69 ),
70 c=ddl.TableSpec(
71 fields=[
72 ddl.FieldSpec("id", dtype=sqlalchemy.BigInteger, autoincrement=True, primaryKey=True),
73 ddl.FieldSpec("b_id", dtype=sqlalchemy.BigInteger, nullable=True),
74 ],
75 foreignKeys=[
76 ddl.ForeignKeySpec("b", source=("b_id",), target=("id",), onDelete="SET NULL"),
77 ],
78 ),
79)
81DYNAMIC_TABLE_SPEC = ddl.TableSpec(
82 fields=[
83 ddl.FieldSpec("c_id", dtype=sqlalchemy.BigInteger, primaryKey=True),
84 ddl.FieldSpec("a_name", dtype=sqlalchemy.String, length=16, nullable=False),
85 ],
86 foreignKeys=[
87 ddl.ForeignKeySpec("c", source=("c_id",), target=("id",), onDelete="CASCADE"),
88 ddl.ForeignKeySpec("a", source=("a_name",), target=("name",), onDelete="CASCADE"),
89 ],
90)
92TEMPORARY_TABLE_SPEC = ddl.TableSpec(
93 fields=[
94 ddl.FieldSpec("a_name", dtype=sqlalchemy.String, length=16, primaryKey=True),
95 ddl.FieldSpec("b_id", dtype=sqlalchemy.BigInteger, primaryKey=True),
96 ],
97)
100@contextmanager
101def _patch_getExistingTable(db: Database) -> Database:
102 """Patch getExistingTable method in a database instance to test concurrent
103 creation of tables. This patch obviously depends on knowning internals of
104 ``ensureTableExists()`` implementation.
105 """
106 original_method = db.getExistingTable
108 def _getExistingTable(name: str, spec: ddl.TableSpec) -> sqlalchemy.schema.Table | None:
109 # Return None on first call, but forward to original method after that
110 db.getExistingTable = original_method
111 return None
113 db.getExistingTable = _getExistingTable
114 yield db
115 db.getExistingTable = original_method
118class DatabaseTests(ABC, TestCaseMixin):
119 """Generic tests for the `Database` interface that can be subclassed to
120 generate tests for concrete implementations.
121 """
123 @abstractmethod
124 def makeEmptyDatabase(self, origin: int = 0) -> Database:
125 """Return an empty `Database` with the given origin, or an
126 automatically-generated one if ``origin`` is `None`.
128 Parameters
129 ----------
130 origin : `int` or `None`
131 Origin to use for the database.
133 Returns
134 -------
135 db : `Database`
136 Empty database with given origin or auto-generated origin.
137 """
138 raise NotImplementedError()
140 @abstractmethod
141 def asReadOnly(self, database: Database) -> AbstractContextManager[Database]:
142 """Return a context manager for a read-only connection into the given
143 database.
145 Parameters
146 ----------
147 database : `Database`
148 The database to use.
150 Yields
151 ------
152 `Database`
153 The new database connection.
155 Notes
156 -----
157 The original database should be considered unusable within the context
158 but safe to use again afterwards (this allows the context manager to
159 block write access by temporarily changing user permissions to really
160 guarantee that write operations are not performed).
161 """
162 raise NotImplementedError()
164 @abstractmethod
165 def getNewConnection(self, database: Database, *, writeable: bool) -> Database:
166 """Return a new `Database` instance that points to the same underlying
167 storage as the given one.
169 Parameters
170 ----------
171 database : `Database`
172 The current database.
173 writeable : `bool`
174 Whether the connection should be writeable or not.
176 Returns
177 -------
178 db : `Database`
179 The new database connection.
180 """
181 raise NotImplementedError()
183 def query_list(
184 self, database: Database, executable: sqlalchemy.sql.expression.SelectBase
185 ) -> list[sqlalchemy.engine.Row]:
186 """Run a SELECT or other read-only query against the database and
187 return the results as a list.
189 Parameters
190 ----------
191 database : `Database`
192 The database to use.
193 executable : `sqlalchemy.sql.expression.SelectBase`
194 Expression to execute.
196 Returns
197 -------
198 results : `list` of `sqlalchemy.engine.Row`
199 The results.
201 Notes
202 -----
203 This is a thin wrapper around database.query() that just avoids
204 context-manager boilerplate that is usefully verbose in production code
205 but just noise in tests.
206 """
207 with database.transaction(), database.query(executable) as result:
208 return result.fetchall()
210 def query_scalar(self, database: Database, executable: sqlalchemy.sql.expression.SelectBase) -> Any:
211 """Run a SELECT query that yields a single column and row against the
212 database and return its value.
214 Parameters
215 ----------
216 database : `Database`
217 The database to use.
218 executable : `sqlalchemy.sql.expression.SelectBase`
219 Expression to execute.
221 Returns
222 -------
223 results : `~typing.Any`
224 The results.
226 Notes
227 -----
228 This is a thin wrapper around database.query() that just avoids
229 context-manager boilerplate that is usefully verbose in production code
230 but just noise in tests.
231 """
232 with database.query(executable) as result:
233 return result.scalar()
235 def checkTable(self, spec: ddl.TableSpec, table: sqlalchemy.schema.Table):
236 self.assertCountEqual(spec.fields.names, table.columns.keys())
237 # Checking more than this currently seems fragile, as it might restrict
238 # what Database implementations do; we don't care if the spec is
239 # actually preserved in terms of types and constraints as long as we
240 # can use the returned table as if it was.
242 def checkStaticSchema(self, tables: StaticTablesTuple):
243 self.checkTable(STATIC_TABLE_SPECS.a, tables.a)
244 self.checkTable(STATIC_TABLE_SPECS.b, tables.b)
245 self.checkTable(STATIC_TABLE_SPECS.c, tables.c)
247 def testDeclareStaticTables(self):
248 """Tests for `Database.declareStaticSchema` and the methods it
249 delegates to.
250 """
251 # Create the static schema in a new, empty database.
252 newDatabase = self.makeEmptyDatabase()
253 with newDatabase.declareStaticTables(create=True) as context:
254 tables = context.addTableTuple(STATIC_TABLE_SPECS)
255 self.checkStaticSchema(tables)
256 # Check that we can load that schema even from a read-only connection.
257 with self.asReadOnly(newDatabase) as existingReadOnlyDatabase:
258 with existingReadOnlyDatabase.declareStaticTables(create=False) as context:
259 tables = context.addTableTuple(STATIC_TABLE_SPECS)
260 self.checkStaticSchema(tables)
262 def testDeclareStaticTablesTwice(self):
263 """Tests for `Database.declareStaticSchema` being called twice."""
264 # Create the static schema in a new, empty database.
265 newDatabase = self.makeEmptyDatabase()
266 with newDatabase.declareStaticTables(create=True) as context:
267 tables = context.addTableTuple(STATIC_TABLE_SPECS)
268 self.checkStaticSchema(tables)
269 # Second time it should raise
270 with self.assertRaises(SchemaAlreadyDefinedError):
271 with newDatabase.declareStaticTables(create=True) as context:
272 tables = context.addTableTuple(STATIC_TABLE_SPECS)
273 # Check schema, it should still contain all tables, and maybe some
274 # extra.
275 with newDatabase.declareStaticTables(create=False) as context:
276 self.assertLessEqual(frozenset(STATIC_TABLE_SPECS._fields), context._tableNames)
278 def testRepr(self):
279 """Test that repr does not return a generic thing."""
280 newDatabase = self.makeEmptyDatabase()
281 rep = repr(newDatabase)
282 # Check that stringification works and gives us something different
283 self.assertNotEqual(rep, str(newDatabase))
284 self.assertNotIn("object at 0x", rep, "Check default repr was not used")
285 self.assertIn("://", rep)
287 def testDynamicTables(self):
288 """Tests for `Database.ensureTableExists` and
289 `Database.getExistingTable`.
290 """
291 # Need to start with the static schema.
292 newDatabase = self.makeEmptyDatabase()
293 with newDatabase.declareStaticTables(create=True) as context:
294 context.addTableTuple(STATIC_TABLE_SPECS)
295 # Try to ensure the dynamic table exists in a read-only version of that
296 # database, which should fail because we can't create it.
297 with self.asReadOnly(newDatabase) as existingReadOnlyDatabase:
298 with existingReadOnlyDatabase.declareStaticTables(create=False) as context:
299 context.addTableTuple(STATIC_TABLE_SPECS)
300 with self.assertRaises(ReadOnlyDatabaseError):
301 existingReadOnlyDatabase.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
302 # Just getting the dynamic table before it exists should return None.
303 self.assertIsNone(newDatabase.getExistingTable("d", DYNAMIC_TABLE_SPEC))
304 # Ensure the new table exists back in the original database, which
305 # should create it.
306 table = newDatabase.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
307 self.checkTable(DYNAMIC_TABLE_SPEC, table)
308 # Ensuring that it exists should just return the exact same table
309 # instance again.
310 self.assertIs(newDatabase.ensureTableExists("d", DYNAMIC_TABLE_SPEC), table)
311 # Try again from the read-only database.
312 with self.asReadOnly(newDatabase) as existingReadOnlyDatabase:
313 with existingReadOnlyDatabase.declareStaticTables(create=False) as context:
314 context.addTableTuple(STATIC_TABLE_SPECS)
315 # Just getting the dynamic table should now work...
316 self.assertIsNotNone(existingReadOnlyDatabase.getExistingTable("d", DYNAMIC_TABLE_SPEC))
317 # ...as should ensuring that it exists, since it now does.
318 existingReadOnlyDatabase.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
319 self.checkTable(DYNAMIC_TABLE_SPEC, table)
320 # Trying to get the table with a different specification (at least
321 # in terms of what columns are present) should raise.
322 with self.assertRaises(DatabaseConflictError):
323 newDatabase.ensureTableExists(
324 "d",
325 ddl.TableSpec(
326 fields=[ddl.FieldSpec("name", dtype=sqlalchemy.String, length=4, primaryKey=True)]
327 ),
328 )
329 # Calling ensureTableExists inside a transaction block is an error,
330 # even if it would do nothing.
331 with newDatabase.transaction():
332 with self.assertRaises(AssertionError):
333 newDatabase.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
335 def testDynamicTablesConcurrency(self):
336 """Tests for `Database.ensureTableExists` concurrent use."""
337 # We cannot really run things concurrently in a deterministic way, here
338 # we just simulate a situation when the table is created by other
339 # process between the call to getExistingTable() and actual table
340 # creation.
341 db1 = self.makeEmptyDatabase()
342 with db1.declareStaticTables(create=True) as context:
343 context.addTableTuple(STATIC_TABLE_SPECS)
344 self.assertIsNone(db1.getExistingTable("d", DYNAMIC_TABLE_SPEC))
346 # Make a dynamic table using separate connection
347 db2 = self.getNewConnection(db1, writeable=True)
348 with db2.declareStaticTables(create=False) as context:
349 context.addTableTuple(STATIC_TABLE_SPECS)
350 table = db2.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
351 self.checkTable(DYNAMIC_TABLE_SPEC, table)
353 # Call it again but trick it into thinking that table is not there.
354 # This test depends on knowing implementation of ensureTableExists()
355 # which initially calls getExistingTable() to check that table may
356 # exist, the patch intercepts that call and returns None.
357 with _patch_getExistingTable(db1):
358 table = db1.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
360 def testTemporaryTables(self):
361 """Tests for `Database.temporary_table`, and `Database.insert` with the
362 ``select`` argument.
363 """
364 # Need to start with the static schema; also insert some test data.
365 newDatabase = self.makeEmptyDatabase()
366 with newDatabase.declareStaticTables(create=True) as context:
367 static = context.addTableTuple(STATIC_TABLE_SPECS)
368 newDatabase.insert(static.a, {"name": "a1", "region": None}, {"name": "a2", "region": None})
369 bIds = newDatabase.insert(
370 static.b,
371 {"name": "b1", "value": 11},
372 {"name": "b2", "value": 12},
373 {"name": "b3", "value": 13},
374 returnIds=True,
375 )
376 # Create the table.
377 with newDatabase.session():
378 with newDatabase.temporary_table(TEMPORARY_TABLE_SPEC, "e1") as table1:
379 self.checkTable(TEMPORARY_TABLE_SPEC, table1)
380 # Insert via a INSERT INTO ... SELECT query.
381 newDatabase.insert(
382 table1,
383 select=sqlalchemy.sql.select(
384 static.a.columns.name.label("a_name"), static.b.columns.id.label("b_id")
385 )
386 .select_from(static.a.join(static.b, onclause=sqlalchemy.sql.literal(True)))
387 .where(
388 sqlalchemy.sql.and_(
389 static.a.columns.name == "a1",
390 static.b.columns.value <= 12,
391 )
392 ),
393 )
394 # Check that the inserted rows are present.
395 self.assertCountEqual(
396 [{"a_name": "a1", "b_id": bId} for bId in bIds[:2]],
397 [row._asdict() for row in self.query_list(newDatabase, table1.select())],
398 )
399 # Create another one via a read-only connection to the
400 # database. We _do_ allow temporary table modifications in
401 # read-only databases.
402 with self.asReadOnly(newDatabase) as existingReadOnlyDatabase:
403 with existingReadOnlyDatabase.declareStaticTables(create=False) as context:
404 context.addTableTuple(STATIC_TABLE_SPECS)
405 with existingReadOnlyDatabase.temporary_table(TEMPORARY_TABLE_SPEC) as table2:
406 self.checkTable(TEMPORARY_TABLE_SPEC, table2)
407 # Those tables should not be the same, despite having
408 # the same ddl.
409 self.assertIsNot(table1, table2)
410 # Do a slightly different insert into this table, to
411 # check that it works in a read-only database. This
412 # time we pass column names as a kwarg to insert
413 # instead of by labeling the columns in the select.
414 existingReadOnlyDatabase.insert(
415 table2,
416 select=sqlalchemy.sql.select(static.a.columns.name, static.b.columns.id)
417 .select_from(static.a.join(static.b, onclause=sqlalchemy.sql.literal(True)))
418 .where(
419 sqlalchemy.sql.and_(
420 static.a.columns.name == "a2",
421 static.b.columns.value >= 12,
422 )
423 ),
424 names=["a_name", "b_id"],
425 )
426 # Check that the inserted rows are present.
427 self.assertCountEqual(
428 [{"a_name": "a2", "b_id": bId} for bId in bIds[1:]],
429 [row._asdict() for row in self.query_list(existingReadOnlyDatabase, table2.select())],
430 )
431 # Exiting the context managers will drop the temporary tables from
432 # the read-only DB. It's unspecified whether attempting to use it
433 # after this point is an error or just never returns any results,
434 # so we can't test what it does, only that it's not an error.
436 def testSchemaSeparation(self):
437 """Test that creating two different `Database` instances allows us
438 to create different tables with the same name in each.
439 """
440 db1 = self.makeEmptyDatabase(origin=1)
441 with db1.declareStaticTables(create=True) as context:
442 tables = context.addTableTuple(STATIC_TABLE_SPECS)
443 self.checkStaticSchema(tables)
445 db2 = self.makeEmptyDatabase(origin=2)
446 # Make the DDL here intentionally different so we'll definitely
447 # notice if db1 and db2 are pointing at the same schema.
448 spec = ddl.TableSpec(fields=[ddl.FieldSpec("id", dtype=sqlalchemy.Integer, primaryKey=True)])
449 with db2.declareStaticTables(create=True) as context:
450 # Make the DDL here intentionally different so we'll definitely
451 # notice if db1 and db2 are pointing at the same schema.
452 table = context.addTable("a", spec)
453 self.checkTable(spec, table)
455 def testInsertQueryDelete(self):
456 """Test the `Database.insert`, `Database.query`, and `Database.delete`
457 methods, as well as the `Base64Region` type and the ``onDelete``
458 argument to `ddl.ForeignKeySpec`.
459 """
460 db = self.makeEmptyDatabase(origin=1)
461 with db.declareStaticTables(create=True) as context:
462 tables = context.addTableTuple(STATIC_TABLE_SPECS)
463 # Insert a single, non-autoincrement row that contains a region and
464 # query to get it back.
465 region = ConvexPolygon((UnitVector3d(1, 0, 0), UnitVector3d(0, 1, 0), UnitVector3d(0, 0, 1)))
466 row = {"name": "a1", "region": region}
467 db.insert(tables.a, row)
468 self.assertEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row])
469 # Insert multiple autoincrement rows but do not try to get the IDs
470 # back immediately.
471 db.insert(tables.b, {"name": "b1", "value": 10}, {"name": "b2", "value": 20})
472 results = [r._asdict() for r in self.query_list(db, tables.b.select().order_by("id"))]
473 self.assertEqual(len(results), 2)
474 for row in results:
475 self.assertIn(row["name"], ("b1", "b2"))
476 self.assertIsInstance(row["id"], int)
477 self.assertGreater(results[1]["id"], results[0]["id"])
478 # Insert multiple autoincrement rows and get the IDs back from insert.
479 rows = [{"name": "b3", "value": 30}, {"name": "b4", "value": 40}]
480 ids = db.insert(tables.b, *rows, returnIds=True)
481 results = [
482 r._asdict()
483 for r in self.query_list(db, tables.b.select().where(tables.b.columns.id > results[1]["id"]))
484 ]
485 expected = [dict(row, id=id) for row, id in zip(rows, ids, strict=True)]
486 self.assertCountEqual(results, expected)
487 self.assertTrue(all(result["id"] is not None for result in results))
488 # Insert multiple rows into a table with an autoincrement primary key,
489 # then use the returned IDs to insert into a dynamic table.
490 rows = [{"b_id": results[0]["id"]}, {"b_id": None}]
491 ids = db.insert(tables.c, *rows, returnIds=True)
492 results = [r._asdict() for r in self.query_list(db, tables.c.select())]
493 expected = [dict(row, id=id) for row, id in zip(rows, ids, strict=True)]
494 self.assertCountEqual(results, expected)
495 self.assertTrue(all(result["id"] is not None for result in results))
496 # Add the dynamic table.
497 d = db.ensureTableExists("d", DYNAMIC_TABLE_SPEC)
498 # Insert into it.
499 rows = [{"c_id": id, "a_name": "a1"} for id in ids]
500 db.insert(d, *rows)
501 results = [r._asdict() for r in self.query_list(db, d.select())]
502 self.assertCountEqual(rows, results)
503 # Insert multiple rows into a table with an autoincrement primary key,
504 # but pass in a value for the autoincrement key.
505 rows2 = [
506 {"id": 700, "b_id": None},
507 {"id": 701, "b_id": None},
508 ]
509 db.insert(tables.c, *rows2)
510 results = [r._asdict() for r in self.query_list(db, tables.c.select())]
511 self.assertCountEqual(results, expected + rows2)
512 self.assertTrue(all(result["id"] is not None for result in results))
514 # Define 'SELECT COUNT(*)' query for later use.
515 count = sqlalchemy.sql.select(sqlalchemy.sql.func.count())
516 # Get the values we inserted into table b.
517 bValues = [r._asdict() for r in self.query_list(db, tables.b.select())]
518 # Remove two row from table b by ID.
519 n = db.delete(tables.b, ["id"], {"id": bValues[0]["id"]}, {"id": bValues[1]["id"]})
520 self.assertEqual(n, 2)
521 # Remove the other two rows from table b by name.
522 n = db.delete(tables.b, ["name"], {"name": bValues[2]["name"]}, {"name": bValues[3]["name"]})
523 self.assertEqual(n, 2)
524 # There should now be no rows in table b.
525 self.assertEqual(self.query_scalar(db, count.select_from(tables.b)), 0)
526 # All b_id values in table c should now be NULL, because there's an
527 # onDelete='SET NULL' foreign key.
528 self.assertEqual(
529 self.query_scalar(
530 db,
531 count.select_from(tables.c).where(tables.c.columns.b_id != None), # noqa:E711
532 ),
533 0,
534 )
535 # Remove all rows in table a (there's only one); this should remove all
536 # rows in d due to onDelete='CASCADE'.
537 n = db.delete(tables.a, [])
538 self.assertEqual(n, 1)
539 self.assertEqual(self.query_scalar(db, count.select_from(tables.a)), 0)
540 self.assertEqual(self.query_scalar(db, count.select_from(d)), 0)
542 def testDeleteWhere(self):
543 """Tests for `Database.deleteWhere`."""
544 db = self.makeEmptyDatabase(origin=1)
545 with db.declareStaticTables(create=True) as context:
546 tables = context.addTableTuple(STATIC_TABLE_SPECS)
547 db.insert(tables.b, *[{"id": i, "name": f"b{i}"} for i in range(10)])
548 count = sqlalchemy.sql.select(sqlalchemy.sql.func.count())
550 n = db.deleteWhere(tables.b, tables.b.columns.id.in_([0, 1, 2]))
551 self.assertEqual(n, 3)
552 self.assertEqual(self.query_scalar(db, count.select_from(tables.b)), 7)
554 n = db.deleteWhere(
555 tables.b,
556 tables.b.columns.id.in_(
557 sqlalchemy.sql.select(tables.b.columns.id).where(tables.b.columns.id > 5)
558 ),
559 )
560 self.assertEqual(n, 4)
561 self.assertEqual(self.query_scalar(db, count.select_from(tables.b)), 3)
563 n = db.deleteWhere(tables.b, tables.b.columns.name == "b5")
564 self.assertEqual(n, 1)
565 self.assertEqual(self.query_scalar(db, count.select_from(tables.b)), 2)
567 n = db.deleteWhere(tables.b, sqlalchemy.sql.literal(True))
568 self.assertEqual(n, 2)
569 self.assertEqual(self.query_scalar(db, count.select_from(tables.b)), 0)
571 def testUpdate(self):
572 """Tests for `Database.update`."""
573 db = self.makeEmptyDatabase(origin=1)
574 with db.declareStaticTables(create=True) as context:
575 tables = context.addTableTuple(STATIC_TABLE_SPECS)
576 # Insert two rows into table a, both without regions.
577 db.insert(tables.a, {"name": "a1"}, {"name": "a2"})
578 # Update one of the rows with a region.
579 region = ConvexPolygon((UnitVector3d(1, 0, 0), UnitVector3d(0, 1, 0), UnitVector3d(0, 0, 1)))
580 n = db.update(tables.a, {"name": "k"}, {"k": "a2", "region": region})
581 self.assertEqual(n, 1)
582 sql = sqlalchemy.sql.select(tables.a.columns.name, tables.a.columns.region).select_from(tables.a)
583 self.assertCountEqual(
584 [r._asdict() for r in self.query_list(db, sql)],
585 [{"name": "a1", "region": None}, {"name": "a2", "region": region}],
586 )
588 def testSync(self):
589 """Tests for `Database.sync`."""
590 db = self.makeEmptyDatabase(origin=1)
591 with db.declareStaticTables(create=True) as context:
592 tables = context.addTableTuple(STATIC_TABLE_SPECS)
593 # Insert a row with sync, because it doesn't exist yet.
594 values, inserted = db.sync(tables.b, keys={"name": "b1"}, extra={"value": 10}, returning=["id"])
595 self.assertTrue(inserted)
596 self.assertEqual(
597 [{"id": values["id"], "name": "b1", "value": 10}],
598 [r._asdict() for r in self.query_list(db, tables.b.select())],
599 )
600 # Repeat that operation, which should do nothing but return the
601 # requested values.
602 values, inserted = db.sync(tables.b, keys={"name": "b1"}, extra={"value": 10}, returning=["id"])
603 self.assertFalse(inserted)
604 self.assertEqual(
605 [{"id": values["id"], "name": "b1", "value": 10}],
606 [r._asdict() for r in self.query_list(db, tables.b.select())],
607 )
608 # Repeat the operation without the 'extra' arg, which should also just
609 # return the existing row.
610 values, inserted = db.sync(tables.b, keys={"name": "b1"}, returning=["id"])
611 self.assertFalse(inserted)
612 self.assertEqual(
613 [{"id": values["id"], "name": "b1", "value": 10}],
614 [r._asdict() for r in self.query_list(db, tables.b.select())],
615 )
616 # Repeat the operation with a different value in 'extra'. That still
617 # shouldn't be an error, because 'extra' is only used if we really do
618 # insert. Also drop the 'returning' argument.
619 _, inserted = db.sync(tables.b, keys={"name": "b1"}, extra={"value": 20})
620 self.assertFalse(inserted)
621 self.assertEqual(
622 [{"id": values["id"], "name": "b1", "value": 10}],
623 [r._asdict() for r in self.query_list(db, tables.b.select())],
624 )
625 # Repeat the operation with the correct value in 'compared' instead of
626 # 'extra'.
627 _, inserted = db.sync(tables.b, keys={"name": "b1"}, compared={"value": 10})
628 self.assertFalse(inserted)
629 self.assertEqual(
630 [{"id": values["id"], "name": "b1", "value": 10}],
631 [r._asdict() for r in self.query_list(db, tables.b.select())],
632 )
633 # Repeat the operation with an incorrect value in 'compared'; this
634 # should raise.
635 with self.assertRaises(DatabaseConflictError):
636 db.sync(tables.b, keys={"name": "b1"}, compared={"value": 20})
637 # Try to sync in a read-only database. This should work if and only
638 # if the matching row already exists.
639 with self.asReadOnly(db) as rodb:
640 with rodb.declareStaticTables(create=False) as context:
641 tables = context.addTableTuple(STATIC_TABLE_SPECS)
642 _, inserted = rodb.sync(tables.b, keys={"name": "b1"})
643 self.assertFalse(inserted)
644 self.assertEqual(
645 [{"id": values["id"], "name": "b1", "value": 10}],
646 [r._asdict() for r in self.query_list(rodb, tables.b.select())],
647 )
648 with self.assertRaises(ReadOnlyDatabaseError):
649 rodb.sync(tables.b, keys={"name": "b2"}, extra={"value": 20})
650 # Repeat the operation with a different value in 'compared' and ask to
651 # update.
652 _, updated = db.sync(tables.b, keys={"name": "b1"}, compared={"value": 20}, update=True)
653 self.assertEqual(updated, {"value": 10})
654 self.assertEqual(
655 [{"id": values["id"], "name": "b1", "value": 20}],
656 [r._asdict() for r in self.query_list(db, tables.b.select())],
657 )
659 def testReplace(self):
660 """Tests for `Database.replace`."""
661 db = self.makeEmptyDatabase(origin=1)
662 with db.declareStaticTables(create=True) as context:
663 tables = context.addTableTuple(STATIC_TABLE_SPECS)
664 # Use 'replace' to insert a single row that contains a region and
665 # query to get it back.
666 region = ConvexPolygon((UnitVector3d(1, 0, 0), UnitVector3d(0, 1, 0), UnitVector3d(0, 0, 1)))
667 row1 = {"name": "a1", "region": region}
668 db.replace(tables.a, row1)
669 self.assertEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1])
670 # Insert another row without a region.
671 row2 = {"name": "a2", "region": None}
672 db.replace(tables.a, row2)
673 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2])
674 # Use replace to re-insert both of those rows again, which should do
675 # nothing.
676 db.replace(tables.a, row1, row2)
677 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2])
678 # Replace row1 with a row with no region, while reinserting row2.
679 row1a = {"name": "a1", "region": None}
680 db.replace(tables.a, row1a, row2)
681 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1a, row2])
682 # Replace both rows, returning row1 to its original state, while adding
683 # a new one. Pass them in in a different order.
684 row2a = {"name": "a2", "region": region}
685 row3 = {"name": "a3", "region": None}
686 db.replace(tables.a, row3, row2a, row1)
687 self.assertCountEqual(
688 [r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2a, row3]
689 )
691 def test_replace_pkey_only(self):
692 """Test `Database.replace` on a table that only has primary key."""
693 spec = ddl.TableSpec(
694 [
695 ddl.FieldSpec("a1", dtype=sqlalchemy.BigInteger, primaryKey=True),
696 ddl.FieldSpec("a2", dtype=sqlalchemy.BigInteger, primaryKey=True),
697 ]
698 )
699 db = self.makeEmptyDatabase(origin=1)
700 with db.declareStaticTables(create=True) as context:
701 table = context.addTable("a", spec)
702 row1 = {"a1": 1, "a2": 2}
703 row2 = {"a1": 1, "a2": 3}
704 db.replace(table, row1)
705 db.replace(table, row2)
706 db.replace(table, row1)
707 self.assertCountEqual([r._asdict() for r in self.query_list(db, table.select())], [row1, row2])
709 def testEnsure(self):
710 """Tests for `Database.ensure`."""
711 db = self.makeEmptyDatabase(origin=1)
712 with db.declareStaticTables(create=True) as context:
713 tables = context.addTableTuple(STATIC_TABLE_SPECS)
714 # Use 'ensure' to insert a single row that contains a region and
715 # query to get it back.
716 region = ConvexPolygon((UnitVector3d(1, 0, 0), UnitVector3d(0, 1, 0), UnitVector3d(0, 0, 1)))
717 row1 = {"name": "a1", "region": region}
718 self.assertEqual(db.ensure(tables.a, row1), 1)
719 self.assertEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1])
720 # Insert another row without a region.
721 row2 = {"name": "a2", "region": None}
722 self.assertEqual(db.ensure(tables.a, row2), 1)
723 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2])
724 # Use ensure to re-insert both of those rows again, which should do
725 # nothing.
726 self.assertEqual(db.ensure(tables.a, row1, row2), 0)
727 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2])
728 # Attempt to insert row1's key with no region, while
729 # reinserting row2. This should also do nothing.
730 row1a = {"name": "a1", "region": None}
731 self.assertEqual(db.ensure(tables.a, row1a, row2), 0)
732 self.assertCountEqual([r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2])
733 # Attempt to insert new rows for both existing keys, this time also
734 # adding a new row. Pass them in in a different order. Only the new
735 # row should be added.
736 row2a = {"name": "a2", "region": region}
737 row3 = {"name": "a3", "region": None}
738 self.assertEqual(db.ensure(tables.a, row3, row2a, row1a), 1)
739 self.assertCountEqual(
740 [r._asdict() for r in self.query_list(db, tables.a.select())], [row1, row2, row3]
741 )
742 # Add some data to a table with both a primary key and a different
743 # unique constraint.
744 row_b = {"id": 5, "name": "five", "value": 50}
745 db.insert(tables.b, row_b)
746 # Attempt ensure with primary_key_only=False and a conflict for the
747 # non-PK constraint. This should do nothing.
748 db.ensure(tables.b, {"id": 10, "name": "five", "value": 200})
749 self.assertEqual([r._asdict() for r in self.query_list(db, tables.b.select())], [row_b])
750 # Now use primary_key_only=True with conflict in only the non-PK field.
751 # This should be an integrity error and nothing should change.
752 with self.assertRaises(sqlalchemy.exc.IntegrityError):
753 db.ensure(tables.b, {"id": 10, "name": "five", "value": 200}, primary_key_only=True)
754 self.assertEqual([r._asdict() for r in self.query_list(db, tables.b.select())], [row_b])
755 # With primary_key_only=True a conflict in the primary key is ignored
756 # regardless of whether there is a conflict elsewhere.
757 db.ensure(tables.b, {"id": 5, "name": "ten", "value": 100}, primary_key_only=True)
758 self.assertEqual([r._asdict() for r in self.query_list(db, tables.b.select())], [row_b])
759 db.ensure(tables.b, {"id": 5, "name": "five", "value": 100}, primary_key_only=True)
760 self.assertEqual([r._asdict() for r in self.query_list(db, tables.b.select())], [row_b])
762 def testTransactionNesting(self):
763 """Test that transactions can be nested with the behavior in the
764 presence of exceptions working as documented.
765 """
766 db = self.makeEmptyDatabase(origin=1)
767 with db.declareStaticTables(create=True) as context:
768 tables = context.addTableTuple(STATIC_TABLE_SPECS)
769 # Insert one row so we can trigger integrity errors by trying to insert
770 # a duplicate of it below.
771 db.insert(tables.a, {"name": "a1"})
772 # First test: error recovery via explicit savepoint=True in the inner
773 # transaction.
774 with db.transaction():
775 # This insert should succeed, and should not be rolled back because
776 # the assertRaises context should catch any exception before it
777 # propagates up to the outer transaction.
778 db.insert(tables.a, {"name": "a2"})
779 with self.assertRaises(sqlalchemy.exc.IntegrityError):
780 with db.transaction(savepoint=True):
781 # This insert should succeed, but should be rolled back.
782 db.insert(tables.a, {"name": "a4"})
783 # This insert should fail (duplicate primary key), raising
784 # an exception.
785 db.insert(tables.a, {"name": "a1"})
786 self.assertCountEqual(
787 [r._asdict() for r in self.query_list(db, tables.a.select())],
788 [{"name": "a1", "region": None}, {"name": "a2", "region": None}],
789 )
790 # Second test: error recovery via implicit savepoint=True, when the
791 # innermost transaction is inside a savepoint=True transaction.
792 with db.transaction():
793 # This insert should succeed, and should not be rolled back
794 # because the assertRaises context should catch any
795 # exception before it propagates up to the outer
796 # transaction.
797 db.insert(tables.a, {"name": "a3"})
798 with self.assertRaises(sqlalchemy.exc.IntegrityError):
799 with db.transaction(savepoint=True):
800 # This insert should succeed, but should be rolled back.
801 db.insert(tables.a, {"name": "a4"})
802 with db.transaction():
803 # This insert should succeed, but should be rolled
804 # back.
805 db.insert(tables.a, {"name": "a5"})
806 # This insert should fail (duplicate primary key),
807 # raising an exception.
808 db.insert(tables.a, {"name": "a1"})
809 self.assertCountEqual(
810 [r._asdict() for r in self.query_list(db, tables.a.select())],
811 [{"name": "a1", "region": None}, {"name": "a2", "region": None}, {"name": "a3", "region": None}],
812 )
814 def testTransactionLocking(self):
815 """Test that `Database.transaction` can be used to acquire a lock
816 that prohibits concurrent writes.
817 """
818 db1 = self.makeEmptyDatabase(origin=1)
819 with db1.declareStaticTables(create=True) as context:
820 tables1 = context.addTableTuple(STATIC_TABLE_SPECS)
822 async def _side1(lock: Iterable[str] = ()) -> tuple[set[str], set[str]]:
823 """One side of the concurrent locking test.
825 Parameters
826 ----------
827 lock : `~collections.abc.Iterable` of `str`
828 Locks.
830 Notes
831 -----
832 This optionally locks the table (and maybe the whole database),
833 does a select for its contents, inserts a new row, and then selects
834 again, with some waiting in between to make sure the other side has
835 a chance to _attempt_ to insert in between. If the locking is
836 enabled and works, the difference between the selects should just
837 be the insert done on this thread.
838 """
839 # Give Side2 a chance to create a connection
840 await asyncio.sleep(1.0)
841 with db1.transaction(lock=lock):
842 names1 = {row.name for row in self.query_list(db1, tables1.a.select())}
843 # Give Side2 a chance to insert (which will be blocked if
844 # we've acquired a lock).
845 await asyncio.sleep(2.0)
846 db1.insert(tables1.a, {"name": "a1"})
847 names2 = {row.name for row in self.query_list(db1, tables1.a.select())}
848 return names1, names2
850 async def _side2() -> None:
851 """Other side of the concurrent locking test.
853 Notes
854 -----
855 This side just waits a bit and then tries to insert a row into the
856 table that the other side is trying to lock. Hopefully that
857 waiting is enough to give the other side a chance to acquire the
858 lock and thus make this side block until the lock is released. If
859 this side manages to do the insert before side1 acquires the lock,
860 we'll just warn about not succeeding at testing the locking,
861 because we can only make that unlikely, not impossible.
862 """
864 def _toRunInThread():
865 """Create new SQLite connection for use in thread.
867 SQLite locking isn't asyncio-friendly unless we actually
868 run it in another thread. And SQLite gets very unhappy if
869 we try to use a connection from multiple threads, so we have
870 to create the new connection here instead of out in the main
871 body of the test function.
872 """
873 db2 = self.getNewConnection(db1, writeable=True)
874 with db2.declareStaticTables(create=False) as context:
875 tables2 = context.addTableTuple(STATIC_TABLE_SPECS)
876 with db2.transaction():
877 db2.insert(tables2.a, {"name": "a2"})
879 await asyncio.sleep(2.0)
880 loop = asyncio.get_running_loop()
881 with ThreadPoolExecutor() as pool:
882 await loop.run_in_executor(pool, _toRunInThread)
884 async def _testProblemsWithNoLocking() -> None:
885 """Run side1 and side2 with no locking, attempting to demonstrate
886 the problem that locking is supposed to solve. If we get unlucky
887 with scheduling, side2 will just happen to insert after side1 is
888 done, and we won't have anything definitive. We just warn in that
889 case because we really don't want spurious test failures.
890 """
891 task1 = asyncio.create_task(_side1())
892 task2 = asyncio.create_task(_side2())
894 names1, names2 = await task1
895 await task2
896 if "a2" in names1:
897 warnings.warn(
898 "Unlucky scheduling in no-locking test: concurrent INSERT happened before first SELECT.",
899 stacklevel=1,
900 )
901 self.assertEqual(names1, {"a2"})
902 self.assertEqual(names2, {"a1", "a2"})
903 elif "a2" not in names2:
904 warnings.warn(
905 "Unlucky scheduling in no-locking test: concurrent INSERT "
906 "happened after second SELECT even without locking.",
907 stacklevel=1,
908 )
909 self.assertEqual(names1, set())
910 self.assertEqual(names2, {"a1"})
911 else:
912 # This is the expected case: both INSERTS happen between the
913 # two SELECTS. If we don't get this almost all of the time we
914 # should adjust the sleep amounts.
915 self.assertEqual(names1, set())
916 self.assertEqual(names2, {"a1", "a2"})
918 asyncio.run(_testProblemsWithNoLocking())
920 # Clean up after first test.
921 db1.delete(tables1.a, ["name"], {"name": "a1"}, {"name": "a2"})
923 async def _testSolutionWithLocking() -> None:
924 """Run side1 and side2 with locking, which should make side2 block
925 its insert until side2 releases its lock.
926 """
927 task1 = asyncio.create_task(_side1(lock=[tables1.a]))
928 task2 = asyncio.create_task(_side2())
930 names1, names2 = await task1
931 await task2
932 if "a2" in names1:
933 warnings.warn(
934 "Unlucky scheduling in locking test: concurrent INSERT happened before first SELECT.",
935 stacklevel=1,
936 )
937 self.assertEqual(names1, {"a2"})
938 self.assertEqual(names2, {"a1", "a2"})
939 else:
940 # This is the expected case: the side2 INSERT happens after the
941 # last SELECT on side1. This can also happen due to unlucky
942 # scheduling, and we have no way to detect that here, but the
943 # similar "no-locking" test has at least some chance of being
944 # affected by the same problem and warning about it.
945 self.assertEqual(names1, set())
946 self.assertEqual(names2, {"a1"})
948 asyncio.run(_testSolutionWithLocking())
950 def testTimespanDatabaseRepresentation(self):
951 """Tests for `TimespanDatabaseRepresentation` and the `Database`
952 methods that interact with it.
953 """
954 # Make some test timespans to play with, with the full suite of
955 # topological relationships.
956 start = astropy.time.Time("2020-01-01T00:00:00", format="isot", scale="tai")
957 offset = astropy.time.TimeDelta(60, format="sec")
958 timestamps = [start + offset * n for n in range(3)]
959 aTimespans = [Timespan(begin=None, end=None)]
960 aTimespans.extend(Timespan(begin=None, end=t) for t in timestamps)
961 aTimespans.extend(Timespan(begin=t, end=None) for t in timestamps)
962 aTimespans.extend(Timespan.fromInstant(t) for t in timestamps)
963 aTimespans.append(Timespan.makeEmpty())
964 aTimespans.extend(Timespan(begin=t1, end=t2) for t1, t2 in itertools.combinations(timestamps, 2))
965 # Make another list of timespans that span the full range but don't
966 # overlap. This is a subset of the previous list.
967 bTimespans = [Timespan(begin=None, end=timestamps[0])]
968 bTimespans.extend(
969 Timespan(begin=t1, end=t2) for t1, t2 in zip(timestamps[:-1], timestamps[1:], strict=True)
970 )
971 bTimespans.append(Timespan(begin=timestamps[-1], end=None))
972 # Make a database and create a table with that database's timespan
973 # representation. This one will have no exclusion constraint and
974 # a nullable timespan.
975 db = self.makeEmptyDatabase(origin=1)
976 TimespanReprClass = db.getTimespanRepresentation()
977 aSpec = ddl.TableSpec(
978 fields=[
979 ddl.FieldSpec(name="id", dtype=sqlalchemy.Integer, primaryKey=True),
980 ],
981 )
982 for fieldSpec in TimespanReprClass.makeFieldSpecs(nullable=True):
983 aSpec.fields.add(fieldSpec)
984 with db.declareStaticTables(create=True) as context:
985 aTable = context.addTable("a", aSpec)
986 self.maxDiff = None
988 def _convertRowForInsert(row: dict) -> dict:
989 """Convert a row containing a Timespan instance into one suitable
990 for insertion into the database.
991 """
992 result = row.copy()
993 ts = result.pop(TimespanReprClass.NAME)
994 return TimespanReprClass.update(ts, result=result)
996 def _convertRowFromSelect(row: dict) -> dict:
997 """Convert a row from the database into one containing a Timespan.
999 Parameters
1000 ----------
1001 row : `dict`
1002 Original row.
1004 Returns
1005 -------
1006 row : `dict`
1007 The updated row.
1008 """
1009 result = row.copy()
1010 timespan = TimespanReprClass.extract(result)
1011 for name in TimespanReprClass.getFieldNames():
1012 del result[name]
1013 result[TimespanReprClass.NAME] = timespan
1014 return result
1016 # Insert rows into table A, in chunks just to make things interesting.
1017 # Include one with a NULL timespan.
1018 aRows = [{"id": n, TimespanReprClass.NAME: t} for n, t in enumerate(aTimespans)]
1019 aRows.append({"id": len(aRows), TimespanReprClass.NAME: None})
1020 db.insert(aTable, _convertRowForInsert(aRows[0]))
1021 db.insert(aTable, *[_convertRowForInsert(r) for r in aRows[1:3]])
1022 db.insert(aTable, *[_convertRowForInsert(r) for r in aRows[3:]])
1023 # Add another one with a NULL timespan, but this time by invoking
1024 # the server-side default.
1025 aRows.append({"id": len(aRows)})
1026 db.insert(aTable, aRows[-1])
1027 aRows[-1][TimespanReprClass.NAME] = None
1028 # Test basic round-trip through database.
1029 self.assertEqual(
1030 aRows,
1031 [
1032 _convertRowFromSelect(row._asdict())
1033 for row in self.query_list(db, aTable.select().order_by(aTable.columns.id))
1034 ],
1035 )
1036 # Create another table B with a not-null timespan and (if the database
1037 # supports it), an exclusion constraint. Use ensureTableExists this
1038 # time to check that mode of table creation vs. timespans.
1039 bSpec = ddl.TableSpec(
1040 fields=[
1041 ddl.FieldSpec(name="id", dtype=sqlalchemy.Integer, primaryKey=True),
1042 ddl.FieldSpec(name="key", dtype=sqlalchemy.Integer, nullable=False),
1043 ],
1044 )
1045 for fieldSpec in TimespanReprClass.makeFieldSpecs(nullable=False):
1046 bSpec.fields.add(fieldSpec)
1047 if TimespanReprClass.hasExclusionConstraint(): 1047 ↛ 1048line 1047 didn't jump to line 1048 because the condition on line 1047 was never true
1048 bSpec.exclusion.add(("key", TimespanReprClass))
1049 bTable = db.ensureTableExists("b", bSpec)
1050 # Insert rows into table B, again in chunks. Each Timespan appears
1051 # twice, but with different values for the 'key' field (which should
1052 # still be okay for any exclusion constraint we may have defined).
1053 bRows = [{"id": n, "key": 1, TimespanReprClass.NAME: t} for n, t in enumerate(bTimespans)]
1054 offset = len(bRows)
1055 bRows.extend(
1056 {"id": n + offset, "key": 2, TimespanReprClass.NAME: t} for n, t in enumerate(bTimespans)
1057 )
1058 db.insert(bTable, *[_convertRowForInsert(r) for r in bRows[:2]])
1059 db.insert(bTable, _convertRowForInsert(bRows[2]))
1060 db.insert(bTable, *[_convertRowForInsert(r) for r in bRows[3:]])
1061 # Insert a row with no timespan into table B. This should invoke the
1062 # server-side default, which is a timespan over (-∞, ∞). We set
1063 # key=3 to avoid upsetting an exclusion constraint that might exist.
1064 bRows.append({"id": len(bRows), "key": 3})
1065 db.insert(bTable, bRows[-1])
1066 bRows[-1][TimespanReprClass.NAME] = Timespan(None, None)
1067 # Test basic round-trip through database.
1068 self.assertEqual(
1069 bRows,
1070 [
1071 _convertRowFromSelect(row._asdict())
1072 for row in self.query_list(db, bTable.select().order_by(bTable.columns.id))
1073 ],
1074 )
1075 # Test that we can't insert timespan=None into this table.
1076 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1077 db.insert(
1078 bTable, _convertRowForInsert({"id": len(bRows), "key": 4, TimespanReprClass.NAME: None})
1079 )
1080 # IFF this database supports exclusion constraints, test that they
1081 # also prevent inserts.
1082 if TimespanReprClass.hasExclusionConstraint(): 1082 ↛ 1083line 1082 didn't jump to line 1083 because the condition on line 1082 was never true
1083 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1084 db.insert(
1085 bTable,
1086 _convertRowForInsert(
1087 {"id": len(bRows), "key": 1, TimespanReprClass.NAME: Timespan(None, timestamps[1])}
1088 ),
1089 )
1090 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1091 db.insert(
1092 bTable,
1093 _convertRowForInsert(
1094 {
1095 "id": len(bRows),
1096 "key": 1,
1097 TimespanReprClass.NAME: Timespan(timestamps[0], timestamps[2]),
1098 }
1099 ),
1100 )
1101 with self.assertRaises(sqlalchemy.exc.IntegrityError):
1102 db.insert(
1103 bTable,
1104 _convertRowForInsert(
1105 {"id": len(bRows), "key": 1, TimespanReprClass.NAME: Timespan(timestamps[2], None)}
1106 ),
1107 )
1108 # Test NULL checks in SELECT queries, on both tables.
1109 aRepr = TimespanReprClass.from_columns(aTable.columns)
1110 self.assertEqual(
1111 [row[TimespanReprClass.NAME] is None for row in aRows],
1112 [
1113 row.f
1114 for row in self.query_list(
1115 db, sqlalchemy.sql.select(aRepr.isNull().label("f")).order_by(aTable.columns.id)
1116 )
1117 ],
1118 )
1119 bRepr = TimespanReprClass.from_columns(bTable.columns)
1120 self.assertEqual(
1121 [False for row in bRows],
1122 [
1123 row.f
1124 for row in self.query_list(
1125 db, sqlalchemy.sql.select(bRepr.isNull().label("f")).order_by(bTable.columns.id)
1126 )
1127 ],
1128 )
1129 # Test relationships expressions that relate in-database timespans to
1130 # Python-literal timespans, all from the more complete 'a' set; check
1131 # that this is consistent with Python-only relationship tests.
1132 for rhsRow in aRows:
1133 if rhsRow[TimespanReprClass.NAME] is None:
1134 continue
1135 with self.subTest(rhsRow=repr(rhsRow)):
1136 expected = {}
1137 for lhsRow in aRows:
1138 if lhsRow[TimespanReprClass.NAME] is None:
1139 expected[lhsRow["id"]] = (None, None, None, None)
1140 else:
1141 expected[lhsRow["id"]] = (
1142 lhsRow[TimespanReprClass.NAME].overlaps(rhsRow[TimespanReprClass.NAME]),
1143 lhsRow[TimespanReprClass.NAME].contains(rhsRow[TimespanReprClass.NAME]),
1144 lhsRow[TimespanReprClass.NAME] < rhsRow[TimespanReprClass.NAME],
1145 lhsRow[TimespanReprClass.NAME] > rhsRow[TimespanReprClass.NAME],
1146 )
1147 rhsRepr = TimespanReprClass.fromLiteral(rhsRow[TimespanReprClass.NAME])
1148 sql = sqlalchemy.sql.select(
1149 aTable.columns.id.label("lhs"),
1150 aRepr.overlaps(rhsRepr).label("overlaps"),
1151 aRepr.contains(rhsRepr).label("contains"),
1152 (aRepr < rhsRepr).label("less_than"),
1153 (aRepr > rhsRepr).label("greater_than"),
1154 ).select_from(aTable)
1155 queried = {
1156 row.lhs: (row.overlaps, row.contains, row.less_than, row.greater_than)
1157 for row in self.query_list(db, sql)
1158 }
1159 self.assertEqual(expected, queried)
1160 # Test relationship expressions that relate in-database timespans to
1161 # each other, all from the more complete 'a' set; check that this is
1162 # consistent with Python-only relationship tests.
1163 expected = {}
1164 for lhs, rhs in itertools.product(aRows, aRows):
1165 lhsT = lhs[TimespanReprClass.NAME]
1166 rhsT = rhs[TimespanReprClass.NAME]
1167 if lhsT is not None and rhsT is not None:
1168 expected[lhs["id"], rhs["id"]] = (
1169 lhsT.overlaps(rhsT),
1170 lhsT.contains(rhsT),
1171 lhsT < rhsT,
1172 lhsT > rhsT,
1173 )
1174 else:
1175 expected[lhs["id"], rhs["id"]] = (None, None, None, None)
1176 lhsSubquery = aTable.alias("lhs")
1177 rhsSubquery = aTable.alias("rhs")
1178 lhsRepr = TimespanReprClass.from_columns(lhsSubquery.columns)
1179 rhsRepr = TimespanReprClass.from_columns(rhsSubquery.columns)
1180 sql = sqlalchemy.sql.select(
1181 lhsSubquery.columns.id.label("lhs"),
1182 rhsSubquery.columns.id.label("rhs"),
1183 lhsRepr.overlaps(rhsRepr).label("overlaps"),
1184 lhsRepr.contains(rhsRepr).label("contains"),
1185 (lhsRepr < rhsRepr).label("less_than"),
1186 (lhsRepr > rhsRepr).label("greater_than"),
1187 ).select_from(lhsSubquery.join(rhsSubquery, onclause=sqlalchemy.sql.literal(True)))
1188 queried = {
1189 (row.lhs, row.rhs): (row.overlaps, row.contains, row.less_than, row.greater_than)
1190 for row in self.query_list(db, sql)
1191 }
1192 self.assertEqual(expected, queried)
1193 # Test relationship expressions between in-database timespans and
1194 # Python-literal instantaneous times.
1195 for t in timestamps:
1196 with self.subTest(t=repr(t)):
1197 expected = {}
1198 for lhsRow in aRows:
1199 if lhsRow[TimespanReprClass.NAME] is None:
1200 expected[lhsRow["id"]] = (None, None, None, None)
1201 else:
1202 expected[lhsRow["id"]] = (
1203 lhsRow[TimespanReprClass.NAME].contains(t),
1204 lhsRow[TimespanReprClass.NAME].overlaps(t),
1205 lhsRow[TimespanReprClass.NAME] < t,
1206 lhsRow[TimespanReprClass.NAME] > t,
1207 )
1208 rhs = sqlalchemy.sql.literal(t, type_=ddl.AstropyTimeNsecTai)
1209 sql = sqlalchemy.sql.select(
1210 aTable.columns.id.label("lhs"),
1211 aRepr.contains(rhs).label("contains"),
1212 aRepr.overlaps(rhs).label("overlaps_point"),
1213 (aRepr < rhs).label("less_than"),
1214 (aRepr > rhs).label("greater_than"),
1215 ).select_from(aTable)
1216 queried = {
1217 row.lhs: (row.contains, row.overlaps_point, row.less_than, row.greater_than)
1218 for row in self.query_list(db, sql)
1219 }
1220 self.assertEqual(expected, queried)
1222 def testConstantRows(self):
1223 """Test Database.constant_rows."""
1224 new_db = self.makeEmptyDatabase()
1225 with new_db.declareStaticTables(create=True) as context:
1226 static = context.addTableTuple(STATIC_TABLE_SPECS)
1227 b_ids = new_db.insert(
1228 static.b,
1229 {"name": "b1", "value": 11},
1230 {"name": "b2", "value": 12},
1231 {"name": "b3", "value": 13},
1232 returnIds=True,
1233 )
1234 values_spec = ddl.TableSpec(
1235 [
1236 ddl.FieldSpec(name="b", dtype=sqlalchemy.BigInteger),
1237 ddl.FieldSpec(name="s", dtype=sqlalchemy.String(8)),
1238 ddl.FieldSpec(name="r", dtype=ddl.Base64Region()),
1239 ],
1240 )
1241 values_data = [
1242 {"b": b_ids[0], "s": "b1", "r": None},
1243 {"b": b_ids[2], "s": "b3", "r": Circle.empty()},
1244 ]
1245 values = new_db.constant_rows(values_spec.fields, *values_data)
1246 select_values_alone = sqlalchemy.sql.select(
1247 values.columns["b"], values.columns["s"], values.columns["r"]
1248 )
1249 self.assertCountEqual(
1250 [row._mapping for row in self.query_list(new_db, select_values_alone)],
1251 values_data,
1252 )
1253 select_values_joined = sqlalchemy.sql.select(
1254 values.columns["s"].label("name"),
1255 static.b.columns["value"].label("value"),
1256 values.columns["r"].label("region"),
1257 ).select_from(values.join(static.b, onclause=static.b.columns["id"] == values.columns["b"]))
1258 self.assertCountEqual(
1259 [row._mapping for row in self.query_list(new_db, select_values_joined)],
1260 [
1261 {"value": 11, "name": "b1", "region": None},
1262 {"value": 13, "name": "b3", "region": Circle.empty()},
1263 ],
1264 )
1266 def test_aggregate(self) -> None:
1267 """Test Database.apply_any_aggregate, ddl.Base64Region.union_aggregate,
1268 and TimespanDatabaseRepresetnation.apply_any_aggregate.
1269 """
1270 db = self.makeEmptyDatabase()
1271 with db.declareStaticTables(create=True) as context:
1272 t = context.addTable(
1273 "t",
1274 ddl.TableSpec(
1275 [
1276 ddl.FieldSpec("id", sqlalchemy.BigInteger(), nullable=False),
1277 ddl.FieldSpec("name", sqlalchemy.String(16), nullable=False),
1278 ddl.FieldSpec.for_region(),
1279 ]
1280 + list(db.getTimespanRepresentation().makeFieldSpecs(nullable=True)),
1281 ),
1282 )
1283 pixelization = Mq3cPixelization(10)
1284 start = astropy.time.Time("2020-01-01T00:00:00", format="isot", scale="tai")
1285 offset = astropy.time.TimeDelta(60, format="sec")
1286 timespans = [Timespan(begin=start + offset * n, end=start + offset * (n + 1)) for n in range(3)]
1287 ts_cls = db.getTimespanRepresentation()
1288 ts_col = ts_cls.from_columns(t.columns)
1289 db.insert(
1290 t,
1291 ts_cls.update(timespans[0], result={"id": 1, "name": "a", "region": pixelization.quad(12058870)}),
1292 ts_cls.update(timespans[1], result={"id": 2, "name": "a", "region": pixelization.quad(12058871)}),
1293 ts_cls.update(timespans[2], result={"id": 3, "name": "b", "region": pixelization.quad(12058872)}),
1294 ts_cls.update(timespans[2], result={"id": 3, "name": "b", "region": pixelization.quad(12058873)}),
1295 )
1296 # This should use DISTINCT ON in PostgreSQL and GROUP BY in SQLite.
1297 if db.has_distinct_on: 1297 ↛ 1298line 1297 didn't jump to line 1298 because the condition on line 1297 was never true
1298 sql = sqlalchemy.select(
1299 t.c.id.label("i"),
1300 t.c.name.label("n"),
1301 *ts_col.flatten("t"),
1302 ).select_from(t)
1303 # SQLAlchemy 2.1 deprecates Select.distinct(*columns) in favor of
1304 # Select.ext(dialects.postgresql.distinct_on(*columns)), but that
1305 # does not exist in 2.0.
1306 if (distinct_on := getattr(sqlalchemy.dialects.postgresql, "distinct_on", None)) and hasattr(
1307 sql, "ext"
1308 ):
1309 sql = sql.ext(distinct_on(t.c.id))
1310 else:
1311 sql = sql.distinct(t.c.id)
1312 elif db.has_any_aggregate:
1313 sql = (
1314 sqlalchemy.select(
1315 t.c.id.label("i"),
1316 db.apply_any_aggregate(t.c.name).label("n"),
1317 *ts_col.apply_any_aggregate(db.apply_any_aggregate).flatten("t"),
1318 )
1319 .select_from(t)
1320 .group_by(t.c.id)
1321 )
1322 else:
1323 raise AssertionError(
1324 "PostgreSQL should support DISTINCT ON and SQLite should support no-op any aggregates."
1325 )
1326 self.assertCountEqual(
1327 [(row.i, row.n, ts_cls.extract(row._mapping, "t")) for row in self.query_list(db, sql)],
1328 [(1, "a", timespans[0]), (2, "a", timespans[1]), (3, "b", timespans[2])],
1329 )
1330 # Test union_aggregate in all versions of both database, with a GROUP
1331 # BY that does not need apply_any_aggregate.
1332 self.assertCountEqual(
1333 [
1334 (row.i, row.r.encode())
1335 for row in self.query_list(
1336 db,
1337 sqlalchemy.select(
1338 t.c.id.label("i"), ddl.Base64Region.union_aggregate(t.c.region).label("r")
1339 )
1340 .select_from(t)
1341 .group_by(t.c.id),
1342 )
1343 ],
1344 [
1345 (1, pixelization.quad(12058870).encode()),
1346 (2, pixelization.quad(12058871).encode()),
1347 (3, UnionRegion(pixelization.quad(12058872), pixelization.quad(12058873)).encode()),
1348 ],
1349 )
1350 if db.has_any_aggregate: 1350 ↛ exitline 1350 didn't return from function 'test_aggregate' because the condition on line 1350 was always true
1351 # This should use run in SQLite and PostgreSQL 16+.
1352 self.assertCountEqual(
1353 [
1354 (row.i, row.n, row.r.encode())
1355 for row in self.query_list(
1356 db,
1357 sqlalchemy.select(
1358 t.c.id.label("i"),
1359 db.apply_any_aggregate(t.c.name).label("n"),
1360 ddl.Base64Region.union_aggregate(t.c.region).label("r"),
1361 )
1362 .select_from(t)
1363 .group_by(t.c.id),
1364 )
1365 ],
1366 [
1367 (1, "a", pixelization.quad(12058870).encode()),
1368 (2, "a", pixelization.quad(12058871).encode()),
1369 (3, "b", UnionRegion(pixelization.quad(12058872), pixelization.quad(12058873)).encode()),
1370 ],
1371 )