DM-41639: Allow Apdb queries for visit completion

python/lsst/dax/apdb/apdb.py

@@ -252,6 +252,23 @@ def getDiaForcedSources(
         """
         raise NotImplementedError()
 
+    @abstractmethod
+    def containsVisitDetector(self, visit: int, detector: int) -> bool:
+        """Test whether data for a given visit-detector is present in the APDB.
+
+        Parameters
+        ----------
+        visit, detector : `int`
+            The ID of the visit-detector to search for.
+
+        Returns
+        -------
+        present : `bool`
+            `True` if some DiaObject, DiaSource, or DiaForcedSource records
+            exist for the specified observation, `False` otherwise.
+        """
+        raise NotImplementedError()
+
     @abstractmethod
     def getInsertIds(self) -> list[ApdbInsertId] | None:
         """Return collection of insert identifiers known to the database.

python/lsst/dax/apdb/apdbCassandra.py

@@ -366,6 +366,10 @@ def getDiaForcedSources(
 
         return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
 
+    def containsVisitDetector(self, visit: int, detector: int) -> bool:
+        # docstring is inherited from a base class
+        raise NotImplementedError()
+
     def getInsertIds(self) -> list[ApdbInsertId] | None:
         # docstring is inherited from a base class
         if not self._schema.has_insert_id:

python/lsst/dax/apdb/apdbSql.py

@@ -387,6 +387,44 @@ def getDiaForcedSources(
         _LOG.debug("found %s DiaForcedSources", len(sources))
         return sources
 
+    def containsVisitDetector(self, visit: int, detector: int) -> bool:
+        # docstring is inherited from a base class
+        raise NotImplementedError()
+
+    def containsCcdVisit(self, ccdVisitId: int) -> bool:
+        """Test whether data for a given visit-detector is present in the APDB.
+
+        This method is a placeholder until `Apdb.containsVisitDetector` can
+        be implemented.
+
+        Parameters
+        ----------
+        ccdVisitId : `int`
+            The packed ID of the visit-detector to search for.
+
+        Returns
+        -------
+        present : `bool`
+            `True` if some DiaSource records exist for the specified
+            observation, `False` otherwise.
+        """
+        # TODO: remove this method in favor of containsVisitDetector on either
+        # DM-41671 or a ticket that removes ccdVisitId from these tables
+        src_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaSource)
+        frcsrc_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaForcedSource)
+        # Query should load only one leaf page of the index
+        query1 = sql.select(src_table.c.ccdVisitId).filter_by(ccdVisitId=ccdVisitId).limit(1)
+        # Backup query in case an image was processed but had no diaSources
+        query2 = sql.select(frcsrc_table.c.ccdVisitId).filter_by(ccdVisitId=ccdVisitId).limit(1)
+
+        with self._engine.begin() as conn:
+            result = conn.execute(query1).scalar_one_or_none()
+            if result is not None:
+                return True
+            else:
+                result = conn.execute(query2).scalar_one_or_none()
+                return result is not None
+
     def getInsertIds(self) -> list[ApdbInsertId] | None:
         # docstring is inherited from a base class
         if not self._schema.has_insert_id:

python/lsst/dax/apdb/tests/_apdb.py

@@ -28,7 +28,7 @@
 
 import pandas
 from lsst.daf.base import DateTime
-from lsst.dax.apdb import ApdbConfig, ApdbInsertId, ApdbTableData, ApdbTables, make_apdb
+from lsst.dax.apdb import ApdbConfig, ApdbInsertId, ApdbSql, ApdbTableData, ApdbTables, make_apdb
 from lsst.sphgeom import Angle, Circle, Region, UnitVector3d
 
 from .data_factory import makeForcedSourceCatalog, makeObjectCatalog, makeSourceCatalog, makeSSObjectCatalog
@@ -70,6 +70,9 @@ class ApdbTest(TestCaseMixin, ABC):
     use_insert_id: bool = False
     """Set to true when support for Insert IDs is configured"""
 
+    allow_visit_query: bool = True
+    """Set to true when contains is implemented"""
+
     # number of columns as defined in tests/config/schema.yaml
     table_column_count = {
         ApdbTables.DiaObject: 8,
@@ -126,7 +129,7 @@ def assert_table_data(self, catalog: Any, rows: int, table: ApdbTables) -> None:
         self.assertEqual(len(catalog.column_names()), self.table_column_count[table] + 1)
 
     def test_makeSchema(self) -> None:
-        """Test for makeing APDB schema."""
+        """Test for making APDB schema."""
         config = self.make_config()
         apdb = make_apdb(config)
 
@@ -175,6 +178,19 @@ def test_empty_gets(self) -> None:
         res = apdb.getDiaForcedSources(region, [1, 2, 3], visit_time)
         self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
 
+        # test if a visit has objects/sources
+        if self.allow_visit_query:
+            res = apdb.containsVisitDetector(visit=0, detector=0)
+            self.assertFalse(res)
+        else:
+            with self.assertRaises(NotImplementedError):
+                apdb.containsVisitDetector(visit=0, detector=0)
+
+        # alternative method not part of the Apdb API
+        if isinstance(apdb, ApdbSql):
+            res = apdb.containsCcdVisit(1)
+            self.assertFalse(res)
+
         # get sources by region
         if self.fsrc_requires_id_list:
             with self.assertRaises(NotImplementedError):
@@ -214,6 +230,19 @@ def test_empty_gets_0months(self) -> None:
         res = apdb.getDiaForcedSources(region, [], visit_time)
         self.assertIs(res, None)
 
+        # test if a visit has objects/sources
+        if self.allow_visit_query:
+            res = apdb.containsVisitDetector(visit=0, detector=0)
+            self.assertFalse(res)
+        else:
+            with self.assertRaises(NotImplementedError):
+                apdb.containsVisitDetector(visit=0, detector=0)
+
+        # alternative method not part of the Apdb API
+        if isinstance(apdb, ApdbSql):
+            res = apdb.containsCcdVisit(1)
+            self.assertFalse(res)
+
     def test_storeObjects(self) -> None:
         """Store and retrieve DiaObjects."""
         # don't care about sources.
@@ -234,6 +263,8 @@ def test_storeObjects(self) -> None:
         res = apdb.getDiaObjects(region)
         self.assert_catalog(res, len(catalog), self.getDiaObjects_table())
 
+        # TODO: test apdb.contains with generic implementation from DM-41671
+
     def test_storeSources(self) -> None:
         """Store and retrieve DiaSources."""
         config = self.make_config()
@@ -263,6 +294,22 @@ def test_storeSources(self) -> None:
         res = apdb.getDiaSources(region, [], visit_time)
         self.assert_catalog(res, 0, ApdbTables.DiaSource)
 
+        # test if a visit is present
+        # data_factory's ccdVisitId generation corresponds to (0, 0)
+        if self.allow_visit_query:
+            res = apdb.containsVisitDetector(visit=0, detector=0)
+            self.assertTrue(res)
+        else:
+            with self.assertRaises(NotImplementedError):
+                apdb.containsVisitDetector(visit=0, detector=0)
+
+        # alternative method not part of the Apdb API
+        if isinstance(apdb, ApdbSql):
+            res = apdb.containsCcdVisit(1)
+            self.assertTrue(res)
+            res = apdb.containsCcdVisit(42)
+            self.assertFalse(res)
+
     def test_storeForcedSources(self) -> None:
         """Store and retrieve DiaForcedSources."""
         config = self.make_config()
@@ -287,6 +334,15 @@ def test_storeForcedSources(self) -> None:
         res = apdb.getDiaForcedSources(region, [], visit_time)
         self.assert_catalog(res, 0, ApdbTables.DiaForcedSource)
 
+        # TODO: test apdb.contains with generic implementation from DM-41671
+
+        # alternative method not part of the Apdb API
+        if isinstance(apdb, ApdbSql):
+            res = apdb.containsCcdVisit(1)
+            self.assertTrue(res)
+            res = apdb.containsCcdVisit(42)
+            self.assertFalse(res)
+
     def test_getHistory(self) -> None:
         """Store and retrieve catalog history."""
         # don't care about sources.

tests/test_apdbCassandra.py

@@ -97,6 +97,7 @@ def make_config(self, **kwargs: Any) -> ApdbCassandraConfig:
 class ApdbCassandraTestCase(ApdbCassandraMixin, unittest.TestCase, ApdbTest):
     """A test case for ApdbCassandra class"""
 
+    allow_visit_query = False
     time_partition_tables = False
     time_partition_start: str | None = None
     time_partition_end: str | None = None

tests/test_apdbSql.py

@@ -46,6 +46,7 @@ class ApdbSQLiteTestCase(unittest.TestCase, ApdbTest):
 
     fsrc_requires_id_list = True
     dia_object_index = "baseline"
+    allow_visit_query = False
 
     def make_config(self, **kwargs: Any) -> ApdbConfig:
         """Make config class instance used in all tests."""
@@ -97,6 +98,7 @@ class ApdbPostgresTestCase(unittest.TestCase, ApdbTest):
     dia_object_index = "last_object_table"
     postgresql: Any
     use_insert_id = True
+    allow_visit_query = False
 
     @classmethod
     def setUpClass(cls) -> None: