Coverage for python/lsst/dax/apdb/tests/_apdb.py: 99%
741 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-07 09:42 +0000
1# This file is part of dax_apdb.
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/>.
22from __future__ import annotations
24__all__ = ["ApdbSchemaUpdateTest", "ApdbTest", "update_schema_yaml"]
26import contextlib
27import itertools
28import logging.config
29import os
30import tempfile
31from abc import ABC, abstractmethod
32from collections.abc import Generator
33from contextlib import closing
34from tempfile import TemporaryDirectory
35from typing import TYPE_CHECKING, Any
37import astropy.time
38import felis.datamodel
39import numpy
40import pandas
41import yaml
43from lsst.resources import ResourcePath
44from lsst.sphgeom import Angle, Circle, LonLat, Region, UnitVector3d
46from .. import (
47 Apdb,
48 ApdbConfig,
49 ApdbReassignDiaSourceToSSObjectRecord,
50 ApdbReplica,
51 ApdbTableData,
52 ApdbTables,
53 ApdbUpdateRecord,
54 ApdbWithdrawDiaSourceRecord,
55 DiaForcedSourceId,
56 DiaObjectId,
57 DiaSourceId,
58 IncompatibleVersionError,
59 ReplicaChunk,
60 VersionTuple,
61)
62from .data_factory import (
63 makeForcedSourceCatalog,
64 makeObjectCatalog,
65 makeSourceCatalog,
66 makeTimestamp,
67 makeTimestampColumn,
68)
69from .utils import TestCaseMixin
71if TYPE_CHECKING:
72 from ..pixelization import Pixelization
75# Optionally configure logging from a config file.
76if log_config := os.environ.get("DAX_APDB_TEST_LOG_CONFIG"): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true
77 logging.config.fileConfig(log_config)
80def _make_region(xyz: tuple[float, float, float] | LonLat = (1.0, 1.0, -1.0)) -> Region:
81 """Make a region to use in tests"""
82 if isinstance(xyz, LonLat):
83 pointing_v = UnitVector3d(xyz)
84 else:
85 pointing_v = UnitVector3d(*xyz)
86 fov = 0.0013 # radians
87 region = Circle(pointing_v, Angle(fov / 2))
88 return region
91@contextlib.contextmanager
92def update_schema_yaml(
93 schema_file: str,
94 drop_metadata: bool = False,
95 version: str | None = None,
96 add_columns: dict[str, list[felis.datamodel.Column]] | None = None,
97) -> Generator[str]:
98 """Update schema definition and return name of the new schema file.
100 Parameters
101 ----------
102 schema_file : `str`
103 Path for the existing YAML file with APDB schema.
104 drop_metadata : `bool`
105 If `True` then remove metadata table from the list of tables.
106 version : `str` or `None`
107 If non-empty string then set schema version to this string, if empty
108 string then remove schema version from config, if `None` - don't change
109 the version in config.
111 Yields
112 ------
113 Path for the updated configuration file.
114 """
115 rp = ResourcePath(schema_file, forceAbsolute=False, forceDirectory=False)
116 schema = yaml.safe_load(rp.read())
118 # Edit YAML contents.
119 # Optionally drop metadata table.
120 if drop_metadata:
121 schema["tables"] = [table for table in schema["tables"] if table["name"] != "metadata"]
122 if version is not None:
123 if version == "":
124 del schema["version"]
125 else:
126 schema["version"] = version
127 if add_columns:
128 for table, columns in add_columns.items():
129 for table_dict in schema["tables"]: 129 ↛ 137line 129 didn't jump to line 137 because the loop on line 129 didn't complete
130 if table_dict["name"] == table:
131 table_dict["columns"] += [
132 column.model_dump(by_alias=True, exclude_none=True, exclude_defaults=True)
133 for column in columns
134 ]
135 break
136 else:
137 raise ValueError(f"Unknown table {table}")
139 with TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
140 output_path = os.path.join(tmpdir, "schema.yaml")
141 with open(output_path, "w") as yaml_stream:
142 yaml.dump(schema, stream=yaml_stream)
143 yield output_path
146class ApdbTest(TestCaseMixin, ABC):
147 """Base class for Apdb tests that can be specialized for concrete
148 implementation.
150 This can only be used as a mixin class for a unittest.TestCase and it
151 calls various assert methods.
152 """
154 visit_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
156 processing_time = astropy.time.Time("2021-01-01T12:00:00", format="isot", scale="tai")
158 fsrc_requires_id_list = False
159 """Should be set to True if getDiaForcedSources requires object IDs"""
161 enable_replica: bool = False
162 """Set to true when support for replication is configured"""
164 use_mjd: bool = True
165 """If True then timestamp columns are MJD TAI."""
167 extra_chunk_columns = 1
168 """Number of additional columns in chunk tables."""
170 meta_row_count = 3
171 """Initial row count in metadata table."""
173 # number of columns as defined in tests/config/schema.yaml
174 table_column_count = {
175 ApdbTables.DiaObject: 8,
176 ApdbTables.DiaObjectLast: 6,
177 ApdbTables.DiaSource: 13,
178 ApdbTables.DiaForcedSource: 9,
179 ApdbTables.SSObject: 3,
180 }
182 @abstractmethod
183 def make_instance(self, **kwargs: Any) -> ApdbConfig:
184 """Make database instance and return configuration for it."""
185 raise NotImplementedError()
187 @abstractmethod
188 def getDiaObjects_table(self) -> ApdbTables:
189 """Return type of table returned from getDiaObjects method."""
190 raise NotImplementedError()
192 @abstractmethod
193 def pixelization(self, config: ApdbConfig) -> Pixelization:
194 """Return pixelization used by implementation."""
195 raise NotImplementedError()
197 def assert_catalog(self, catalog: Any, rows: int, table: ApdbTables) -> None:
198 """Validate catalog type and size
200 Parameters
201 ----------
202 catalog : `object`
203 Expected type of this is ``pandas.DataFrame``.
204 rows : `int`
205 Expected number of rows in a catalog.
206 table : `ApdbTables`
207 APDB table type.
208 """
209 self.assertIsInstance(catalog, pandas.DataFrame)
210 self.assertEqual(catalog.shape[0], rows)
211 self.assertEqual(catalog.shape[1], self.table_column_count[table])
213 def assert_table_data(self, catalog: Any, rows: int, table: ApdbTables) -> None:
214 """Validate catalog type and size
216 Parameters
217 ----------
218 catalog : `object`
219 Expected type of this is `ApdbTableData`.
220 rows : `int`
221 Expected number of rows in a catalog.
222 table : `ApdbTables`
223 APDB table type.
224 extra_columns : `int`
225 Count of additional columns expected in ``catalog``.
226 """
227 self.assertIsInstance(catalog, ApdbTableData)
228 n_rows = sum(1 for row in catalog.rows())
229 self.assertEqual(n_rows, rows)
230 # One extra column for replica chunk id
231 self.assertEqual(
232 len(catalog.column_names()), self.table_column_count[table] + self.extra_chunk_columns
233 )
235 def assert_column_types(self, catalog: Any, types: dict[str, felis.datamodel.DataType]) -> None:
236 column_defs = dict(catalog.column_defs())
237 for column, datatype in types.items():
238 self.assertEqual(column_defs[column], datatype)
240 def make_region(self, xyz: tuple[float, float, float] | LonLat = (1.0, 1.0, -1.0)) -> Region:
241 """Make a region to use in tests"""
242 return _make_region(xyz)
244 def test_makeSchema(self) -> None:
245 """Test for making APDB schema."""
246 config = self.make_instance()
247 with closing(Apdb.from_config(config)) as apdb:
248 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject))
249 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast))
250 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource))
251 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource))
252 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata))
253 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject))
254 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource))
255 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match))
257 # Test from_uri factory method with the same config.
258 with tempfile.NamedTemporaryFile() as tmpfile:
259 config.save(tmpfile.name)
260 with closing(Apdb.from_uri(tmpfile.name)) as apdb:
261 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject))
262 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObjectLast))
263 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaSource))
264 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaForcedSource))
265 self.assertIsNotNone(apdb.tableDef(ApdbTables.metadata))
266 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSObject))
267 self.assertIsNotNone(apdb.tableDef(ApdbTables.SSSource))
268 self.assertIsNotNone(apdb.tableDef(ApdbTables.DiaObject_To_Object_Match))
270 def test_empty_gets(self) -> None:
271 """Test for getting data from empty database.
273 All get() methods should return empty results, only useful for
274 checking that code is not broken.
275 """
276 # use non-zero months for Forced/Source fetching
277 config = self.make_instance()
278 with closing(Apdb.from_config(config)) as apdb:
279 region = self.make_region()
280 visit_time = self.visit_time
282 res: pandas.DataFrame | None
284 # get objects by region
285 res = apdb.getDiaObjects(region)
286 self.assert_catalog(res, 0, self.getDiaObjects_table())
288 # get sources by region
289 res = apdb.getDiaSources(region, None, visit_time)
290 self.assert_catalog(res, 0, ApdbTables.DiaSource)
292 res = apdb.getDiaSources(region, [], visit_time)
293 self.assert_catalog(res, 0, ApdbTables.DiaSource)
295 # get sources by object ID, non-empty object list
296 res = apdb.getDiaSources(region, [1, 2, 3], visit_time)
297 self.assert_catalog(res, 0, ApdbTables.DiaSource)
299 # get forced sources by object ID, empty object list
300 res = apdb.getDiaForcedSources(region, [], visit_time)
301 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
303 # get sources by object ID, non-empty object list
304 res = apdb.getDiaForcedSources(region, [1, 2, 3], visit_time)
305 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
307 # data_factory's ccdVisitId generation corresponds to (1, 1)
308 res = apdb.containsVisitDetector(visit=1, detector=1)
309 self.assertFalse(res)
311 # get sources by region
312 if self.fsrc_requires_id_list: 312 ↛ 316line 312 didn't jump to line 316 because the condition on line 312 was always true
313 with self.assertRaises(NotImplementedError):
314 apdb.getDiaForcedSources(region, None, visit_time)
315 else:
316 res = apdb.getDiaForcedSources(region, None, visit_time)
317 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
319 def test_empty_gets_0months(self) -> None:
320 """Test for getting data from empty database.
322 All get() methods should return empty DataFrame or None.
323 """
324 # set read_sources_months to 0 so that Forced/Sources are None
325 config = self.make_instance(read_sources_months=0, read_forced_sources_months=0)
326 with closing(Apdb.from_config(config)) as apdb:
327 region = self.make_region()
328 visit_time = self.visit_time
330 res: pandas.DataFrame | None
332 # get objects by region
333 res = apdb.getDiaObjects(region)
334 self.assert_catalog(res, 0, self.getDiaObjects_table())
336 # get sources by region
337 res = apdb.getDiaSources(region, None, visit_time)
338 self.assertIs(res, None)
340 # get sources by object ID, empty object list
341 res = apdb.getDiaSources(region, [], visit_time)
342 self.assertIs(res, None)
344 # get forced sources by object ID, empty object list
345 res = apdb.getDiaForcedSources(region, [], visit_time)
346 self.assertIs(res, None)
348 # Database is empty, no images exist.
349 res = apdb.containsVisitDetector(visit=1, detector=1)
350 self.assertFalse(res)
352 def test_storeObjects(self) -> None:
353 """Store and retrieve DiaObjects."""
354 # don't care about sources.
355 config = self.make_instance()
356 with closing(Apdb.from_config(config)) as apdb:
357 region = self.make_region()
358 visit_time = self.visit_time
360 # make catalog with Objects
361 catalog = makeObjectCatalog(region, 100)
363 # store catalog
364 apdb.store(visit_time, catalog)
366 # read it back and check sizes
367 res = apdb.getDiaObjects(region)
368 self.assert_catalog(res, len(catalog), self.getDiaObjects_table())
370 # TODO: test apdb.contains with generic implementation in DM-41671
372 def test_storeObjects_empty(self) -> None:
373 """Test calling storeObject when there are no objects: see DM-43270."""
374 config = self.make_instance()
375 with closing(Apdb.from_config(config)) as apdb:
376 region = self.make_region()
377 visit_time = self.visit_time
378 # make catalog with no Objects
379 catalog = makeObjectCatalog(region, 0)
381 with self.assertLogs("lsst.dax.apdb", level="DEBUG") as cm:
382 apdb.store(visit_time, catalog)
383 self.assertIn("No objects", "\n".join(cm.output))
385 def test_storeMovingObject(self) -> None:
386 """Store and retrieve DiaObject which changes its position."""
387 # don't care about sources.
388 config = self.make_instance()
389 with closing(Apdb.from_config(config)) as apdb:
390 pixelization = self.pixelization(config)
392 lon_deg, lat_deg = 0.0, 0.0
393 lonlat1 = LonLat.fromDegrees(lon_deg - 1.0, lat_deg)
394 lonlat2 = LonLat.fromDegrees(lon_deg + 1.0, lat_deg)
395 uv1 = UnitVector3d(lonlat1)
396 uv2 = UnitVector3d(lonlat2)
398 # Check that they fall into different pixels.
399 self.assertNotEqual(pixelization.pixel(uv1), pixelization.pixel(uv2))
401 # Store one object at two different positions.
402 visit_time1 = self.visit_time
403 catalog1 = makeObjectCatalog(lonlat1, 1)
404 apdb.store(visit_time1, catalog1)
406 visit_time2 = visit_time1 + astropy.time.TimeDelta(120.0, format="sec")
407 catalog1 = makeObjectCatalog(lonlat2, 1)
408 apdb.store(visit_time2, catalog1)
410 # Make region covering both points.
411 region = Circle(UnitVector3d(LonLat.fromDegrees(lon_deg, lat_deg)), Angle.fromDegrees(1.1))
412 self.assertTrue(region.contains(uv1))
413 self.assertTrue(region.contains(uv2))
415 # Read it back, must return the latest one.
416 res = apdb.getDiaObjects(region)
417 self.assert_catalog(res, 1, self.getDiaObjects_table())
419 def test_storeSources(self) -> None:
420 """Store and retrieve DiaSources."""
421 config = self.make_instance()
422 with closing(Apdb.from_config(config)) as apdb:
423 region = self.make_region()
424 visit_time = self.visit_time
426 # have to store Objects first
427 objects = makeObjectCatalog(region, 100)
428 oids = list(objects["diaObjectId"])
429 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)
431 # save the objects and sources
432 apdb.store(visit_time, objects, sources)
434 # read it back, no ID filtering
435 res = apdb.getDiaSources(region, None, visit_time)
436 self.assert_catalog(res, len(sources), ApdbTables.DiaSource)
438 # read it back and filter by ID
439 res = apdb.getDiaSources(region, oids, visit_time)
440 self.assert_catalog(res, len(sources), ApdbTables.DiaSource)
442 # read it back to get schema
443 res = apdb.getDiaSources(region, [], visit_time)
444 self.assert_catalog(res, 0, ApdbTables.DiaSource)
446 # test if a visit is present
447 # data_factory's ccdVisitId generation corresponds to (1, 1)
448 res = apdb.containsVisitDetector(visit=1, detector=1)
449 self.assertTrue(res)
450 # non-existent image
451 res = apdb.containsVisitDetector(visit=2, detector=42)
452 self.assertFalse(res)
454 def test_storeForcedSources(self) -> None:
455 """Store and retrieve DiaForcedSources."""
456 config = self.make_instance()
457 with closing(Apdb.from_config(config)) as apdb:
458 region = self.make_region()
459 visit_time = self.visit_time
461 # have to store Objects first
462 objects = makeObjectCatalog(region, 100)
463 oids = list(objects["diaObjectId"])
464 catalog = makeForcedSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)
466 apdb.store(visit_time, objects, forced_sources=catalog)
468 # read it back and check sizes
469 res = apdb.getDiaForcedSources(region, oids, visit_time)
470 self.assert_catalog(res, len(catalog), ApdbTables.DiaForcedSource)
472 # read it back to get schema
473 res = apdb.getDiaForcedSources(region, [], visit_time)
474 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
476 # data_factory's ccdVisitId generation corresponds to (1, 1)
477 res = apdb.containsVisitDetector(visit=1, detector=1)
478 self.assertTrue(res)
479 # non-existent image
480 res = apdb.containsVisitDetector(visit=2, detector=42)
481 self.assertFalse(res)
483 def test_null_integer_type(self) -> None:
484 """Test that integer column with NULLs correct type on select."""
485 config = self.make_instance()
486 with closing(Apdb.from_config(config)) as apdb:
487 region = self.make_region()
488 visit_time = self.visit_time
490 # have to store Objects first
491 objects = makeObjectCatalog(region, 100)
492 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)
493 # Reset some diaObjectIds to NULL.
494 sources.loc[0:10, "diaObjectId"] = None
496 # save the objects and sources
497 apdb.store(visit_time, objects, sources)
499 # read it back, no ID filtering
500 res = apdb.getDiaSources(region, None, visit_time)
501 self.assert_catalog(res, len(sources), ApdbTables.DiaSource)
502 assert res is not None, "Expecting catalog, not None"
503 self.assertEqual(res.dtypes["diaObjectId"], pandas.Int64Dtype())
505 def test_timestamps(self) -> None:
506 """Check that timestamp return type is as expected."""
507 config = self.make_instance()
508 with closing(Apdb.from_config(config)) as apdb:
509 region = self.make_region()
510 visit_time = self.visit_time
512 # Cassandra has a millisecond precision, so subtract 1ms to allow
513 # for truncated returned values.
514 time_before = makeTimestamp(self.processing_time, self.use_mjd, -1)
515 objects = makeObjectCatalog(region, 100)
516 oids = list(objects["diaObjectId"])
517 catalog = makeForcedSourceCatalog(
518 objects, visit_time, processing_time=self.processing_time, use_mjd=self.use_mjd
519 )
520 time_after = makeTimestamp(self.processing_time, self.use_mjd)
522 apdb.store(visit_time, objects, forced_sources=catalog)
524 # read it back and check sizes
525 res = apdb.getDiaForcedSources(region, oids, visit_time)
526 assert res is not None
527 self.assert_catalog(res, len(catalog), ApdbTables.DiaForcedSource)
529 time_processed_column = makeTimestampColumn("time_processed", self.use_mjd)
530 self.assertIn(time_processed_column, res.dtypes)
531 dtype = res.dtypes[time_processed_column]
532 timestamp_type_names = (
533 ("float64",) if self.use_mjd else ("datetime64[ms]", "datetime64[us]", "datetime64[ns]")
534 )
535 self.assertIn(dtype.name, timestamp_type_names)
536 # Verify that returned time is sensible.
537 self.assertTrue(all(time_before <= dt <= time_after for dt in res[time_processed_column]))
539 def test_getDiaObjectsForDedup(self) -> None:
540 """Test getDiaObjectsForDedup() method."""
541 config = self.make_instance()
542 with closing(Apdb.from_config(config)) as apdb:
543 region1 = self.make_region((1.0, 1.0, -1.0))
544 region2 = self.make_region((-1.0, 1.0, -1.0))
545 region3 = self.make_region((-1.0, -1.0, -1.0))
546 nobj = 100
547 objects1 = makeObjectCatalog(region1, nobj)
548 objects2 = makeObjectCatalog(region2, nobj, start_id=nobj * 2)
549 objects3 = makeObjectCatalog(region3, nobj, start_id=nobj * 4)
551 visits = [
552 (astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai"), objects1),
553 (astropy.time.Time("2021-01-01T00:10:00", format="isot", scale="tai"), objects2),
554 (astropy.time.Time("2021-01-01T00:20:00", format="isot", scale="tai"), objects3),
555 ]
557 for visit_time, objects in visits:
558 apdb.store(visit_time, objects)
560 catalog = apdb.getDiaObjectsForDedup()
561 self.assertEqual(len(catalog), 300)
563 catalog = apdb.getDiaObjectsForDedup(visits[0][0])
564 self.assertEqual(len(catalog), 300)
566 catalog = apdb.getDiaObjectsForDedup(visits[1][0])
567 self.assertEqual(len(catalog), 200)
569 catalog = apdb.getDiaObjectsForDedup(visits[2][0])
570 self.assertEqual(len(catalog), 100)
572 time = astropy.time.Time("2021-01-01T00:30:00", format="isot", scale="tai")
573 catalog = apdb.getDiaObjectsForDedup(time)
574 self.assertEqual(len(catalog), 0)
576 def test_getDiaSourcesForDiaObjects(self) -> None:
577 """Test getDiaSourcesForDiaObjects() method."""
578 config = self.make_instance()
579 with closing(Apdb.from_config(config)) as apdb:
580 # Monkey-patch APDB instance to set current time.
581 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]
583 region1 = self.make_region((1.0, 1.0, -1.0))
584 region2 = self.make_region((-1.0, 1.0, -1.0))
585 region3 = self.make_region((-1.0, -1.0, -1.0))
586 nobj = 100
587 objects1 = makeObjectCatalog(region1, nobj)
588 objects2 = makeObjectCatalog(region2, nobj, start_id=nobj * 2)
589 objects3 = makeObjectCatalog(region3, nobj, start_id=nobj * 4)
591 visits = [
592 (astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai"), objects1),
593 (astropy.time.Time("2021-01-01T00:10:00", format="isot", scale="tai"), objects2),
594 (astropy.time.Time("2021-01-01T00:20:00", format="isot", scale="tai"), objects3),
595 ]
597 start_id = 1_000_000
598 for visit_time, objects in visits:
599 sources = makeSourceCatalog(objects, visit_time, start_id=start_id, use_mjd=self.use_mjd)
600 apdb.store(visit_time, objects, sources)
601 start_id += 1_000_000
603 # Take a small number of objects from different regions.
604 object_ids = [
605 DiaObjectId.from_named_tuple(next(objects1.itertuples())),
606 DiaObjectId.from_named_tuple(next(objects2.itertuples())),
607 DiaObjectId.from_named_tuple(next(objects3.itertuples())),
608 ]
610 catalog = apdb.getDiaSourcesForDiaObjects(object_ids, visits[0][0])
611 self.assertEqual(len(catalog), 3)
612 self.assertEqual(set(catalog["diaObjectId"]), {1, 200, 400})
613 self.assertEqual(set(catalog["diaSourceId"]), {1_000_000, 2_000_000, 3_000_000})
615 catalog = apdb.getDiaSourcesForDiaObjects(object_ids, visits[2][0])
616 self.assertEqual(len(catalog), 1)
617 self.assertEqual(set(catalog["diaObjectId"]), {400})
618 self.assertEqual(set(catalog["diaSourceId"]), {3_000_000})
620 def test_reassignDiaSourcesToDiaObjects(self) -> None:
621 """Test reassignDiaSourcesToDiaObjects() method."""
622 config = self.make_instance()
623 with (
624 closing(Apdb.from_config(config)) as apdb,
625 closing(ApdbReplica.from_config(config)) as apdb_replica,
626 ):
627 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]
629 visit_time = self.visit_time
630 lonlat1 = LonLat.fromDegrees(0.0, 0.0)
631 lonlat2 = LonLat.fromDegrees(180.0, 0.0)
632 # regons around lonlat1/2
633 region1 = self.make_region(lonlat1)
634 region2 = self.make_region(lonlat2)
636 # Store 3 objects and sources at the same position in each region.
637 objects = makeObjectCatalog(lonlat1, 3, start_id=100)
638 sources = makeSourceCatalog(objects, visit_time, start_id=1000, use_mjd=self.use_mjd)
639 apdb.store(visit_time, objects, sources)
641 objects = makeObjectCatalog(lonlat2, 3, start_id=200)
642 sources = makeSourceCatalog(objects, visit_time, start_id=2000, use_mjd=self.use_mjd)
643 apdb.store(visit_time, objects, sources)
645 # check that everything as we think it is.
646 objects = apdb.getDiaObjects(region1)
647 self.assertEqual(set(objects["diaObjectId"]), {100, 101, 102})
648 self.assertEqual(list(objects["nDiaSources"]), [1, 1, 1])
649 sources = apdb.getDiaSources(region1, [100, 101, 102], visit_time)
650 assert sources is not None
651 self.assertEqual(set(sources["diaSourceId"]), {1000, 1001, 1002})
652 self.assertEqual(set(sources["diaObjectId"]), {100, 101, 102})
654 dia_source_ids = [DiaSourceId.from_named_tuple(row) for row in sources.itertuples()]
656 # Reassign sources in region1 and increment/decrement nDiaSources.
657 reassign = {
658 dia_source_id: 100
659 for dia_source_id in dia_source_ids
660 if dia_source_id.diaSourceId in (1001, 1002)
661 }
662 apdb.reassignDiaSourcesToDiaObjects(reassign)
664 objects = apdb.getDiaObjects(region1)
665 self.assertEqual(set(objects["nDiaSources"]), {0, 3})
666 sources = apdb.getDiaSources(region1, [100], visit_time)
667 assert sources is not None
668 self.assertEqual(set(sources["diaSourceId"]), {1000, 1001, 1002})
669 self.assertEqual(set(sources["diaObjectId"]), {100})
671 sources = apdb.getDiaSources(region2, [201, 202], visit_time)
672 assert sources is not None
673 self.assertEqual(set(sources["diaSourceId"]), {2001, 2002})
674 dia_source_ids = [DiaSourceId.from_named_tuple(row) for row in sources.itertuples()]
676 # Reassign but do not increment/decrement nDiaSources.
677 reassign = {
678 dia_source_id: 200
679 for dia_source_id in dia_source_ids
680 if dia_source_id.diaSourceId in (2001, 2002)
681 }
682 apdb.reassignDiaSourcesToDiaObjects(
683 reassign, increment_nDiaSources=False, decrement_nDiaSources=False
684 )
686 objects = apdb.getDiaObjects(region2)
687 self.assertEqual(set(objects["nDiaSources"]), {1})
688 sources = apdb.getDiaSources(region2, [200], visit_time)
689 assert sources is not None
690 self.assertEqual(set(sources["diaSourceId"]), {2000, 2001, 2002})
691 self.assertEqual(set(sources["diaObjectId"]), {200})
693 replica_chunks = apdb_replica.getReplicaChunks()
694 if not self.enable_replica:
695 self.assertIsNone(replica_chunks)
696 else:
697 assert replica_chunks is not None
699 # There could be one or two chunks.
700 self.assertTrue(1 <= len(replica_chunks) <= 2)
702 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])
703 # Two reassignments for region1, three increments/decrements
704 # for that region, plus two reassignments for region2 without
705 # increments/decrements.
706 self.assertEqual(len(update_records), 2 + 3 + 2)
708 def test_setValidityEnd(self) -> None:
709 """Store DiaObjects and truncate validity for some."""
710 # don't care about sources.
711 config = self.make_instance()
712 with (
713 closing(Apdb.from_config(config)) as apdb,
714 closing(ApdbReplica.from_config(config)) as apdb_replica,
715 ):
716 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]
718 region = self.make_region()
719 visit_time = self.visit_time
721 # make catalog with Objects
722 catalog = makeObjectCatalog(region, 100)
724 # store catalog
725 apdb.store(visit_time, catalog)
727 # read it back and check sizes
728 res = apdb.getDiaObjects(region)
729 self.assert_catalog(res, 100, self.getDiaObjects_table())
731 # Select first 10 objects.
732 object_ids = [DiaObjectId.from_named_tuple(row) for row in catalog.iloc[:10].itertuples()]
733 count = apdb.setValidityEnd(object_ids, self.processing_time)
734 self.assertEqual(count, 10)
736 res = apdb.getDiaObjects(region)
737 self.assert_catalog(res, 90, self.getDiaObjects_table())
739 replica_chunks = apdb_replica.getReplicaChunks()
740 if not self.enable_replica:
741 self.assertIsNone(replica_chunks)
742 else:
743 # Check that there are 10 update records in replica tables.
744 assert replica_chunks is not None
746 # There could be one or two chunks.
747 self.assertTrue(1 <= len(replica_chunks) <= 2)
749 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])
750 self.assertEqual(len(update_records), 10)
752 # Check that empty list works.
753 count = apdb.setValidityEnd(object_ids, self.processing_time)
754 self.assertEqual(count, 0)
756 # Try with non-existing object.
757 object_ids = [DiaObjectId.from_named_tuple(row) for row in catalog.iloc[10:12].itertuples()]
758 object_ids += [DiaObjectId(diaObjectId=1_000_000, ra=0.0, dec=0.0)]
759 with self.assertRaises(LookupError):
760 apdb.setValidityEnd(object_ids, self.processing_time, raise_on_missing_id=True)
762 count = apdb.setValidityEnd(object_ids, self.processing_time)
763 self.assertEqual(count, 2)
765 def test_resetDedup(self) -> None:
766 """Test resetDedup method."""
767 # don't care about sources.
768 config = self.make_instance()
769 with closing(Apdb.from_config(config)) as apdb:
770 region = self.make_region()
772 # make catalog with Objects
773 objects = makeObjectCatalog(region, 100)
775 visit_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
776 dedup_time1 = astropy.time.Time("2021-01-01T12:00:00", format="isot", scale="tai")
777 visit_time2 = astropy.time.Time("2021-01-02T00:00:00", format="isot", scale="tai")
778 dedup_time2 = astropy.time.Time("2021-01-02T12:00:00", format="isot", scale="tai")
780 # store catalog
781 apdb.store(visit_time1, objects)
783 catalog = apdb.getDiaObjectsForDedup()
784 self.assertEqual(len(catalog), 100)
786 catalog = apdb.getDiaObjectsForDedup(visit_time1)
787 self.assertEqual(len(catalog), 100)
789 apdb.resetDedup(dedup_time1)
791 catalog = apdb.getDiaObjectsForDedup(visit_time1)
792 self.assertEqual(len(catalog), self._count_after_reset_dedup(100))
794 apdb.store(visit_time2, objects)
796 catalog = apdb.getDiaObjectsForDedup()
797 self.assertEqual(len(catalog), 100)
799 catalog = apdb.getDiaObjectsForDedup(dedup_time1)
800 self.assertEqual(len(catalog), 100)
802 apdb.resetDedup(dedup_time2)
804 catalog = apdb.getDiaObjectsForDedup(dedup_time1)
805 self.assertEqual(len(catalog), self._count_after_reset_dedup(100))
807 catalog = apdb.getDiaObjectsForDedup()
808 self.assertEqual(len(catalog), 0)
810 def _count_after_reset_dedup(self, count_before: int) -> int:
811 """Return the number of rows that will be returned by
812 getDiaObjectsForDedup() after resetDedup() was called. For SQL backend
813 deduplication data comes from a regular table, and it is not removed
814 by resetDedup().
815 """
816 raise NotImplementedError()
818 def test_withdraw_sources(self) -> None:
819 """Test withdrawDiaSources() method."""
820 config = self.make_instance()
821 with (
822 closing(Apdb.from_config(config)) as apdb,
823 closing(ApdbReplica.from_config(config)) as apdb_replica,
824 ):
825 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]
827 lonlat1 = LonLat.fromDegrees(0.0, 0.0)
828 region1 = self.make_region(lonlat1)
829 lonlat2 = LonLat.fromDegrees(45.0, 0.0)
830 region2 = self.make_region(lonlat2)
832 # Store 3 objects and sources at the same position in each region.
833 # The code originally updated nDiaSources so there are objects with
834 # nDiaSources > 1, but we dropped that option.
835 visit_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
836 objects1 = makeObjectCatalog(lonlat1, 3, start_id=100)
837 sources1 = makeSourceCatalog(objects1, visit_time1, start_id=1000, use_mjd=self.use_mjd)
838 apdb.store(visit_time1, objects1, sources1)
840 visit_time2 = astropy.time.Time("2021-01-01T00:01:00", format="isot", scale="tai")
841 objects2 = makeObjectCatalog(lonlat2, 3, start_id=200)
842 sources2 = makeSourceCatalog(objects2, visit_time2, start_id=2000, use_mjd=self.use_mjd)
843 apdb.store(visit_time2, objects2, sources2)
845 # Fetch everything and verify.
846 objects1 = apdb.getDiaObjects(region1)
847 objects2 = apdb.getDiaObjects(region2)
848 self.assertEqual({row.diaObjectId for row in objects1.itertuples()}, {100, 101, 102})
849 self.assertEqual({row.diaObjectId for row in objects2.itertuples()}, {200, 201, 202})
851 sources1 = apdb.getDiaSources(region1, None, visit_time2)
852 sources2 = apdb.getDiaSources(region2, None, visit_time2)
853 assert sources1 is not None and sources2 is not None
854 source_ids = [
855 DiaSourceId.from_named_tuple(row)
856 for row in itertools.chain(sources1.itertuples(), sources2.itertuples())
857 ]
858 sources_by_id = {source_id.diaSourceId: source_id for source_id in source_ids}
860 if self.use_mjd:
861 self.assertTrue(all(pandas.isna(sources1["timeWithdrawnMjdTai"])))
862 self.assertTrue(all(pandas.isna(sources2["timeWithdrawnMjdTai"])))
863 else:
864 self.assertTrue(all(pandas.isnull(sources1["time_withdrawn"])))
865 self.assertTrue(all(pandas.isnull(sources2["time_withdrawn"])))
867 # Withdraw a bunch of sources.
868 withdraw_time1 = astropy.time.Time("2021-01-01T10:00:00", format="isot", scale="tai")
869 withdraw_time2 = astropy.time.Time("2021-01-02T10:00:00", format="isot", scale="tai")
870 apdb.withdrawDiaSources([sources_by_id[i] for i in (1000, 2001)], timeWithdrawn=withdraw_time1)
871 # Withdraw diaSourceId=1000 second time, should not change it.
872 apdb.withdrawDiaSources(
873 [sources_by_id[i] for i in (1000, 2000, 1002)], timeWithdrawn=withdraw_time2
874 )
876 sources1 = apdb.getDiaSources(region1, None, visit_time2)
877 sources2 = apdb.getDiaSources(region2, None, visit_time2)
878 assert sources1 is not None and sources2 is not None
879 sources1.set_index("diaSourceId", inplace=True)
880 sources2.set_index("diaSourceId", inplace=True)
881 if self.use_mjd:
882 self.assertEqual(sources1.loc[1000, "timeWithdrawnMjdTai"], withdraw_time1.mjd)
883 self.assertTrue(numpy.isnan(sources1.loc[1001, "timeWithdrawnMjdTai"]))
884 self.assertEqual(sources1.loc[1002, "timeWithdrawnMjdTai"], withdraw_time2.mjd)
885 self.assertEqual(sources2.loc[2000, "timeWithdrawnMjdTai"], withdraw_time2.mjd)
886 self.assertEqual(sources2.loc[2001, "timeWithdrawnMjdTai"], withdraw_time1.mjd)
887 self.assertTrue(numpy.isnan(sources2.loc[2002, "timeWithdrawnMjdTai"]))
888 else:
889 # Exact type and values depend on backend, I don't want to
890 # overcomplicate it, just check for NaT.
891 self.assertEqual(
892 dict(numpy.isnat(sources1["time_withdrawn"])), {1000: False, 1001: True, 1002: False}
893 )
894 self.assertEqual(
895 dict(numpy.isnat(sources2["time_withdrawn"])), {2000: False, 2001: False, 2002: True}
896 )
898 # Check replication update tables.
899 replica_chunks = apdb_replica.getReplicaChunks()
900 if not self.enable_replica:
901 self.assertIsNone(replica_chunks)
902 else:
903 # Check that there are 4 update records in replica tables.
904 assert replica_chunks is not None
906 # There could be one or two chunks.
907 self.assertTrue(1 <= len(replica_chunks) <= 2)
909 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])
910 self.assertEqual(len(update_records), 4)
912 def test_withdraw_forced_sources(self) -> None:
913 """Test withdrawDiaForcedSources() method."""
914 config = self.make_instance()
915 with (
916 closing(Apdb.from_config(config)) as apdb,
917 closing(ApdbReplica.from_config(config)) as apdb_replica,
918 ):
919 apdb._current_time = lambda: self.processing_time # type: ignore[method-assign]
921 lonlat = LonLat.fromDegrees(0.0, 0.0)
922 region = self.make_region(lonlat)
924 # Store 3 objects and sources at the same position in each region.
925 objects = makeObjectCatalog(lonlat, 3, start_id=100)
926 fsources = makeForcedSourceCatalog(objects, self.visit_time, use_mjd=self.use_mjd)
927 apdb.store(self.visit_time, objects, None, fsources)
929 fsources = apdb.getDiaForcedSources(region, [100, 101, 102], self.visit_time)
930 assert fsources is not None
931 source_ids = [DiaForcedSourceId.from_named_tuple(row) for row in fsources.itertuples()]
932 sources_by_id = {source_id.diaObjectId: source_id for source_id in source_ids}
934 if self.use_mjd:
935 self.assertTrue(all(pandas.isna(fsources["timeWithdrawnMjdTai"])))
936 else:
937 self.assertTrue(all(pandas.isnull(fsources["time_withdrawn"])))
939 # Withdraw sources.
940 withdraw_time1 = astropy.time.Time("2021-01-01T10:00:00", format="isot", scale="tai")
941 apdb.withdrawDiaForcedSources(
942 [sources_by_id[i] for i in (100, 102)], timeWithdrawn=withdraw_time1
943 )
944 withdraw_time2 = astropy.time.Time("2021-01-02T10:00:00", format="isot", scale="tai")
945 # DiaSourceId=102 withdrawn second time, it has no effect.
946 apdb.withdrawDiaForcedSources(
947 [sources_by_id[i] for i in (101, 102)], timeWithdrawn=withdraw_time2
948 )
950 fsources = apdb.getDiaForcedSources(region, [100, 101, 102], self.visit_time)
951 assert fsources is not None
952 fsources.set_index("diaObjectId", inplace=True)
953 if self.use_mjd:
954 self.assertEqual(fsources.loc[100, "timeWithdrawnMjdTai"], withdraw_time1.mjd)
955 self.assertEqual(fsources.loc[101, "timeWithdrawnMjdTai"], withdraw_time2.mjd)
956 self.assertEqual(fsources.loc[102, "timeWithdrawnMjdTai"], withdraw_time1.mjd)
957 else:
958 # Exact type and values depend on backend, I don't want to
959 # overcomplicate it, just check that all of them are not NULL.
960 self.assertFalse(any(pandas.isnull(fsources["time_withdrawn"])))
962 # Check replication update tables.
963 replica_chunks = apdb_replica.getReplicaChunks()
964 if not self.enable_replica:
965 self.assertIsNone(replica_chunks)
966 else:
967 # Check that there are 3 update records in replica tables.
968 assert replica_chunks is not None
970 # There could be one or two chunks.
971 self.assertTrue(1 <= len(replica_chunks) <= 2)
973 update_records = apdb_replica.getUpdateRecordChunks([chunk.id for chunk in replica_chunks])
974 self.assertEqual(len(update_records), 3)
976 def test_getChunks(self) -> None:
977 """Store and retrieve replica chunks."""
978 # don't care about sources.
979 config = self.make_instance()
980 with (
981 closing(Apdb.from_config(config)) as apdb,
982 closing(ApdbReplica.from_config(config)) as apdb_replica,
983 ):
984 visit_time = self.visit_time
986 region1 = self.make_region((1.0, 1.0, -1.0))
987 region2 = self.make_region((-1.0, -1.0, -1.0))
988 nobj = 100
989 objects1 = makeObjectCatalog(region1, nobj)
990 objects2 = makeObjectCatalog(region2, nobj, start_id=nobj * 2)
992 # With the default 10 minutes replica chunk window we should have 4
993 # records.
994 visits = [
995 (astropy.time.Time("2021-01-01T00:01:00", format="isot", scale="tai"), objects1),
996 (astropy.time.Time("2021-01-01T00:02:00", format="isot", scale="tai"), objects2),
997 (astropy.time.Time("2021-01-01T00:11:00", format="isot", scale="tai"), objects1),
998 (astropy.time.Time("2021-01-01T00:12:00", format="isot", scale="tai"), objects2),
999 (astropy.time.Time("2021-01-01T00:45:00", format="isot", scale="tai"), objects1),
1000 (astropy.time.Time("2021-01-01T00:46:00", format="isot", scale="tai"), objects2),
1001 (astropy.time.Time("2021-03-01T00:01:00", format="isot", scale="tai"), objects1),
1002 (astropy.time.Time("2021-03-01T00:02:00", format="isot", scale="tai"), objects2),
1003 ]
1005 start_id = 0
1006 for visit_time, objects in visits:
1007 sources = makeSourceCatalog(objects, visit_time, start_id=start_id, use_mjd=self.use_mjd)
1008 fsources = makeForcedSourceCatalog(objects, visit_time, visit=start_id, use_mjd=self.use_mjd)
1009 apdb.store(visit_time, objects, sources, fsources)
1010 start_id += nobj
1012 replica_chunks = apdb_replica.getReplicaChunks()
1013 if not self.enable_replica:
1014 self.assertIsNone(replica_chunks)
1016 with self.assertRaisesRegex(ValueError, "APDB is not configured for replication"):
1017 apdb_replica.getTableDataChunks(ApdbTables.DiaObject, [])
1019 else:
1020 assert replica_chunks is not None
1021 self.assertEqual(len(replica_chunks), 4)
1023 with self.assertRaisesRegex(ValueError, "does not support replica chunks"):
1024 apdb_replica.getTableDataChunks(ApdbTables.SSObject, [])
1026 def _check_chunks(replica_chunks: list[ReplicaChunk], n_records: int | None = None) -> None:
1027 if n_records is None:
1028 n_records = len(replica_chunks) * nobj
1029 res = apdb_replica.getTableDataChunks(
1030 ApdbTables.DiaObject, (chunk.id for chunk in replica_chunks)
1031 )
1032 self.assert_table_data(res, n_records, ApdbTables.DiaObject)
1033 validityStartColumn = "validityStartMjdTai" if self.use_mjd else "validityStart"
1034 validityStartType = (
1035 felis.datamodel.DataType.double
1036 if self.use_mjd
1037 else felis.datamodel.DataType.timestamp
1038 )
1039 self.assert_column_types(
1040 res,
1041 {
1042 "apdb_replica_chunk": felis.datamodel.DataType.long,
1043 "diaObjectId": felis.datamodel.DataType.long,
1044 validityStartColumn: validityStartType,
1045 "ra": felis.datamodel.DataType.double,
1046 "dec": felis.datamodel.DataType.double,
1047 "parallax": felis.datamodel.DataType.float,
1048 "nDiaSources": felis.datamodel.DataType.int,
1049 },
1050 )
1052 res = apdb_replica.getTableDataChunks(
1053 ApdbTables.DiaSource, (chunk.id for chunk in replica_chunks)
1054 )
1055 self.assert_table_data(res, n_records, ApdbTables.DiaSource)
1056 self.assert_column_types(
1057 res,
1058 {
1059 "apdb_replica_chunk": felis.datamodel.DataType.long,
1060 "diaSourceId": felis.datamodel.DataType.long,
1061 "visit": felis.datamodel.DataType.long,
1062 "detector": felis.datamodel.DataType.short,
1063 },
1064 )
1066 res = apdb_replica.getTableDataChunks(
1067 ApdbTables.DiaForcedSource, (chunk.id for chunk in replica_chunks)
1068 )
1069 self.assert_table_data(res, n_records, ApdbTables.DiaForcedSource)
1070 self.assert_column_types(
1071 res,
1072 {
1073 "apdb_replica_chunk": felis.datamodel.DataType.long,
1074 "diaObjectId": felis.datamodel.DataType.long,
1075 "visit": felis.datamodel.DataType.long,
1076 "detector": felis.datamodel.DataType.short,
1077 },
1078 )
1080 # read it back and check sizes
1081 _check_chunks(replica_chunks, 800)
1082 _check_chunks(replica_chunks[1:], 600)
1083 _check_chunks(replica_chunks[1:-1], 400)
1084 _check_chunks(replica_chunks[2:3], 200)
1085 _check_chunks([])
1087 # try to remove some of those
1088 deleted_chunks = replica_chunks[:1]
1089 apdb_replica.deleteReplicaChunks(chunk.id for chunk in deleted_chunks)
1091 # All queries on deleted ids should return empty set.
1092 _check_chunks(deleted_chunks, 0)
1094 replica_chunks = apdb_replica.getReplicaChunks()
1095 assert replica_chunks is not None
1096 self.assertEqual(len(replica_chunks), 3)
1098 _check_chunks(replica_chunks, 600)
1100 def test_reassignObjects(self) -> None:
1101 """Reassign DiaObjects."""
1102 # don't care about sources.
1103 config = self.make_instance()
1104 with closing(Apdb.from_config(config)) as apdb:
1105 region = self.make_region()
1106 visit_time = self.visit_time
1107 objects = makeObjectCatalog(region, 100)
1108 oids = list(objects["diaObjectId"])
1109 sources = makeSourceCatalog(objects, visit_time, use_mjd=self.use_mjd)
1110 apdb.store(visit_time, objects, sources)
1112 # read it back and filter by ID
1113 res = apdb.getDiaSources(region, oids, visit_time)
1114 self.assert_catalog(res, len(sources), ApdbTables.DiaSource)
1116 apdb.reassignDiaSources({1: 1, 2: 2, 5: 5})
1117 res = apdb.getDiaSources(region, oids, visit_time)
1118 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource)
1120 with self.assertRaisesRegex(ValueError, r"do not exist.*\D1000"):
1121 apdb.reassignDiaSources(
1122 {
1123 1000: 1,
1124 7: 3,
1125 }
1126 )
1127 self.assert_catalog(res, len(sources) - 3, ApdbTables.DiaSource)
1129 def test_storeUpdateRecord(self) -> None:
1130 """Test _storeUpdateRecord() method."""
1131 config = self.make_instance()
1132 with closing(Apdb.from_config(config)) as apdb:
1133 # Times are totally arbitrary.
1134 update_time_ns1 = 2_000_000_000_000_000_000
1135 update_time_ns2 = 2_000_000_001_000_000_000
1136 records = [
1137 ApdbReassignDiaSourceToSSObjectRecord(
1138 update_time_ns=update_time_ns1,
1139 update_order=0,
1140 diaSourceId=1,
1141 ssObjectId=1,
1142 ssObjectReassocTimeMjdTai=60000.0,
1143 ra=45.0,
1144 dec=-45.0,
1145 midpointMjdTai=60000.0,
1146 ),
1147 ApdbWithdrawDiaSourceRecord(
1148 update_time_ns=update_time_ns1,
1149 update_order=1,
1150 diaSourceId=123456,
1151 timeWithdrawnMjdTai=61000.0,
1152 ra=45.0,
1153 dec=-45.0,
1154 midpointMjdTai=60000.0,
1155 ),
1156 ApdbReassignDiaSourceToSSObjectRecord(
1157 update_time_ns=update_time_ns1,
1158 update_order=3,
1159 diaSourceId=2,
1160 ssObjectId=3,
1161 ssObjectReassocTimeMjdTai=60000.0,
1162 ra=45.0,
1163 dec=-45.0,
1164 midpointMjdTai=60000.0,
1165 ),
1166 ApdbWithdrawDiaSourceRecord(
1167 update_time_ns=update_time_ns2,
1168 update_order=0,
1169 diaSourceId=123456,
1170 timeWithdrawnMjdTai=61000.0,
1171 ra=45.0,
1172 dec=-45.0,
1173 midpointMjdTai=60000.0,
1174 ),
1175 ]
1177 update_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
1178 chunk = ReplicaChunk.make_replica_chunk(update_time, 600)
1180 if not self.enable_replica:
1181 with self.assertRaises(TypeError):
1182 self.store_update_records(apdb, records, chunk)
1183 else:
1184 self.store_update_records(apdb, records, chunk)
1186 with closing(ApdbReplica.from_config(config)) as apdb_replica:
1187 records_returned = apdb_replica.getUpdateRecordChunks([chunk.id])
1189 # Input records are ordered, output will be ordered too.
1190 self.assertEqual(records_returned, records)
1192 @abstractmethod
1193 def store_update_records(self, apdb: Apdb, records: list[ApdbUpdateRecord], chunk: ReplicaChunk) -> None:
1194 """Store update records in database, must be overriden in subclass."""
1195 raise NotImplementedError()
1197 def test_midpointMjdTai_src(self) -> None:
1198 """Test for time filtering of DiaSources."""
1199 config = self.make_instance()
1200 with closing(Apdb.from_config(config)) as apdb:
1201 region = self.make_region()
1202 # 2021-01-01 plus 360 days is 2021-12-27
1203 src_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
1204 src_time2 = astropy.time.Time("2021-01-01T00:00:02", format="isot", scale="tai")
1205 visit_time0 = astropy.time.Time("2021-12-26T23:59:59", format="isot", scale="tai")
1206 visit_time1 = astropy.time.Time("2021-12-27T00:00:01", format="isot", scale="tai")
1207 visit_time2 = astropy.time.Time("2021-12-27T00:00:03", format="isot", scale="tai")
1208 one_sec = astropy.time.TimeDelta(1.0, format="sec")
1210 objects = makeObjectCatalog(region, 100)
1211 oids = list(objects["diaObjectId"])
1212 sources = makeSourceCatalog(objects, src_time1, 0, use_mjd=self.use_mjd)
1213 apdb.store(src_time1, objects, sources)
1215 sources = makeSourceCatalog(objects, src_time2, 100, use_mjd=self.use_mjd)
1216 apdb.store(src_time2, objects, sources)
1218 # reading at time of last save should read all
1219 res = apdb.getDiaSources(region, oids, src_time2)
1220 self.assert_catalog(res, 200, ApdbTables.DiaSource)
1222 # one second before 12 months
1223 res = apdb.getDiaSources(region, oids, visit_time0)
1224 self.assert_catalog(res, 200, ApdbTables.DiaSource)
1226 # reading at later time of last save should only read a subset
1227 res = apdb.getDiaSources(region, oids, visit_time1)
1228 self.assert_catalog(res, 100, ApdbTables.DiaSource)
1230 # reading at later time of last save should only read a subset
1231 res = apdb.getDiaSources(region, oids, visit_time2)
1232 self.assert_catalog(res, 0, ApdbTables.DiaSource)
1234 # Use explicit start time argument instead of 12 month window,
1235 # visit time does not matter in this case, set it to before all
1236 # data.
1237 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time1 - one_sec)
1238 self.assert_catalog(res, 200, ApdbTables.DiaSource)
1240 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time2 - one_sec)
1241 self.assert_catalog(res, 100, ApdbTables.DiaSource)
1243 res = apdb.getDiaSources(region, oids, src_time1 - one_sec, src_time2 + one_sec)
1244 self.assert_catalog(res, 0, ApdbTables.DiaSource)
1246 def test_midpointMjdTai_fsrc(self) -> None:
1247 """Test for time filtering of DiaForcedSources."""
1248 config = self.make_instance()
1249 with closing(Apdb.from_config(config)) as apdb:
1250 region = self.make_region()
1251 src_time1 = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
1252 src_time2 = astropy.time.Time("2021-01-01T00:00:02", format="isot", scale="tai")
1253 visit_time0 = astropy.time.Time("2021-12-26T23:59:59", format="isot", scale="tai")
1254 visit_time1 = astropy.time.Time("2021-12-27T00:00:01", format="isot", scale="tai")
1255 visit_time2 = astropy.time.Time("2021-12-27T00:00:03", format="isot", scale="tai")
1256 one_sec = astropy.time.TimeDelta(1.0, format="sec")
1258 objects = makeObjectCatalog(region, 100)
1259 oids = list(objects["diaObjectId"])
1260 sources = makeForcedSourceCatalog(objects, src_time1, 1, use_mjd=self.use_mjd)
1261 apdb.store(src_time1, objects, forced_sources=sources)
1263 sources = makeForcedSourceCatalog(objects, src_time2, 2, use_mjd=self.use_mjd)
1264 apdb.store(src_time2, objects, forced_sources=sources)
1266 # reading at time of last save should read all
1267 res = apdb.getDiaForcedSources(region, oids, src_time2)
1268 self.assert_catalog(res, 200, ApdbTables.DiaForcedSource)
1270 # one second before 12 months
1271 res = apdb.getDiaForcedSources(region, oids, visit_time0)
1272 self.assert_catalog(res, 200, ApdbTables.DiaForcedSource)
1274 # reading at later time of last save should only read a subset
1275 res = apdb.getDiaForcedSources(region, oids, visit_time1)
1276 self.assert_catalog(res, 100, ApdbTables.DiaForcedSource)
1278 # reading at later time of last save should only read a subset
1279 res = apdb.getDiaForcedSources(region, oids, visit_time2)
1280 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
1282 # Use explicit start time argument instead of 12 month window,
1283 # visit time does not matter in this case, set it to before all
1284 # data.
1285 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time1 - one_sec)
1286 self.assert_catalog(res, 200, ApdbTables.DiaForcedSource)
1288 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time2 - one_sec)
1289 self.assert_catalog(res, 100, ApdbTables.DiaForcedSource)
1291 res = apdb.getDiaForcedSources(region, oids, src_time1 - one_sec, src_time2 + one_sec)
1292 self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
1294 def test_metadata(self) -> None:
1295 """Simple test for writing/reading metadata table"""
1296 config = self.make_instance()
1297 with closing(Apdb.from_config(config)) as apdb:
1298 metadata = apdb.metadata
1300 # APDB should write two or three metadata items with version
1301 # numbers and a frozen JSON config.
1302 self.assertFalse(metadata.empty())
1303 self.assertEqual(len(list(metadata.items())), self.meta_row_count)
1305 metadata.set("meta", "data")
1306 metadata.set("data", "meta")
1308 self.assertFalse(metadata.empty())
1309 self.assertTrue(set(metadata.items()) >= {("meta", "data"), ("data", "meta")})
1311 with self.assertRaisesRegex(KeyError, "Metadata key 'meta' already exists"):
1312 metadata.set("meta", "data1")
1314 metadata.set("meta", "data2", force=True)
1315 self.assertTrue(set(metadata.items()) >= {("meta", "data2"), ("data", "meta")})
1317 self.assertTrue(metadata.delete("meta"))
1318 self.assertIsNone(metadata.get("meta"))
1319 self.assertFalse(metadata.delete("meta"))
1321 self.assertEqual(metadata.get("data"), "meta")
1322 self.assertEqual(metadata.get("meta", "meta"), "meta")
1324 def test_schemaVersionFromYaml(self) -> None:
1325 """Check version number handling for reading schema from YAML."""
1326 config = self.make_instance()
1327 default_schema = config.schema_file
1328 with closing(Apdb.from_config(config)) as apdb:
1329 self.assertEqual(apdb.schema.schemaVersion(), VersionTuple(0, 1, 1))
1331 with update_schema_yaml(default_schema, version="") as schema_file:
1332 config = self.make_instance(schema_file=schema_file)
1333 with closing(Apdb.from_config(config)) as apdb:
1334 self.assertEqual(
1335 apdb.schema.schemaVersion(),
1336 VersionTuple(0, 1, 0),
1337 )
1339 with update_schema_yaml(default_schema, version="99.0.0") as schema_file:
1340 config = self.make_instance(schema_file=schema_file)
1341 with closing(Apdb.from_config(config)) as apdb:
1342 self.assertEqual(
1343 apdb.schema.schemaVersion(),
1344 VersionTuple(99, 0, 0),
1345 )
1347 def test_config_freeze(self) -> None:
1348 """Test that some config fields are correctly frozen in database."""
1349 config = self.make_instance()
1351 # `enable_replica` is the only parameter that is frozen in all
1352 # implementations.
1353 config.enable_replica = not self.enable_replica
1354 with closing(Apdb.from_config(config)) as apdb:
1355 frozen_config = apdb.getConfig()
1356 self.assertEqual(frozen_config.enable_replica, self.enable_replica)
1359class ApdbSchemaUpdateTest(TestCaseMixin, ABC):
1360 """Base class for unit tests that verify how schema changes work."""
1362 visit_time = astropy.time.Time("2021-01-01T00:00:00", format="isot", scale="tai")
1364 @abstractmethod
1365 def make_instance(self, **kwargs: Any) -> ApdbConfig:
1366 """Make config class instance used in all tests.
1368 This method should return configuration that point to the identical
1369 database instance on each call (i.e. ``db_url`` must be the same,
1370 which also means for sqlite it has to use on-disk storage).
1371 """
1372 raise NotImplementedError()
1374 @abstractmethod
1375 def add_column(self, apdb: Apdb, table: ApdbTables, column: felis.datamodel.Column) -> None:
1376 """Add a column to an existing table."""
1377 raise NotImplementedError()
1379 def make_region(self, xyz: tuple[float, float, float] | LonLat = (1.0, 1.0, -1.0)) -> Region:
1380 """Make a region to use in tests"""
1381 return _make_region(xyz)
1383 def test_schema_add_replica(self) -> None:
1384 """Check that new code can work with old schema without replica
1385 tables.
1386 """
1387 # Make schema without replica tables.
1388 config = self.make_instance(enable_replica=False)
1389 # Make APDB instance configured for replication.
1390 config.enable_replica = True
1391 with (
1392 closing(Apdb.from_config(config)) as apdb,
1393 closing(ApdbReplica.from_config(config)) as apdb_replica,
1394 ):
1395 # Try to insert something, should work OK.
1396 region = self.make_region()
1397 visit_time = self.visit_time
1399 # have to store Objects first
1400 objects = makeObjectCatalog(region, 100)
1401 sources = makeSourceCatalog(objects, visit_time)
1402 fsources = makeForcedSourceCatalog(objects, visit_time)
1403 apdb.store(visit_time, objects, sources, fsources)
1405 # There should be no replica chunks.
1406 replica_chunks = apdb_replica.getReplicaChunks()
1407 self.assertIsNone(replica_chunks)
1409 def test_schemaVersionCheck(self) -> None:
1410 """Check version number compatibility."""
1411 config = self.make_instance()
1412 with closing(Apdb.from_config(config)) as apdb:
1413 self.assertEqual(apdb.schema.schemaVersion(), VersionTuple(0, 1, 1))
1414 # Claim that schema version is now 99.0.0, must raise an exception.
1415 with update_schema_yaml(config.schema_file, version="99.0.0") as schema_file:
1416 config.schema_file = schema_file
1417 with self.assertRaises(IncompatibleVersionError):
1418 with closing(Apdb.from_config(config)) as apdb:
1419 # Version is checked only when we try to do connect.
1420 apdb.metadata.items()
1422 def test_add_non_null_column(self) -> None:
1423 """Test for adding a non-null column to a table with existing data."""
1424 # Want replica tables in the schema too.
1425 config = self.make_instance(enable_replica=True)
1426 with closing(Apdb.from_config(config)) as apdb:
1427 region = self.make_region()
1428 visit_time = self.visit_time
1430 # Store some data.
1431 objects = makeObjectCatalog(region, 100)
1432 sources = makeSourceCatalog(objects, visit_time)
1433 fsources = makeForcedSourceCatalog(objects, visit_time)
1434 apdb.store(visit_time, objects, sources, fsources)
1436 new_column = felis.datamodel.Column.model_validate(
1437 {
1438 "@id": "#DiaSource.non_null_column",
1439 "name": "non_null_column",
1440 "datatype": felis.datamodel.DataType.int,
1441 "nullable": False,
1442 "value": 123,
1443 },
1444 )
1445 self.add_column(apdb, ApdbTables.DiaSource, new_column)
1447 # Update configured schema.
1448 columns_to_add = {"DiaSource": [new_column]}
1449 with update_schema_yaml(config.schema_file, add_columns=columns_to_add) as schema_file:
1450 config.schema_file = schema_file
1451 with (
1452 closing(Apdb.from_config(config)) as apdb,
1453 closing(ApdbReplica.from_config(config)) as apdb_replica,
1454 ):
1455 # Read sources.
1456 sources = apdb.getDiaSources(region, None, visit_time)
1457 assert sources is not None
1458 self.assertFalse((sources["non_null_column"].isna()).any())
1459 self.assertTrue((sources["non_null_column"] == 123).all())
1461 replica_chunks = apdb_replica.getReplicaChunks()
1462 assert replica_chunks is not None, "Replication must be enabled"
1464 source_data = apdb_replica.getTableDataChunks(
1465 ApdbTables.DiaSource, [chunk.id for chunk in replica_chunks]
1466 )
1467 idx = source_data.column_names().index("non_null_column")
1468 self.assertTrue(all(source[idx] == 123 for source in source_data.rows()))