Merge pull request #514 from lsst/tickets/DM-29273

DM-29273: Fix deprecated sqlalchemy code

python/lsst/daf/butler/registry/attributes.py

@@ -44,12 +44,6 @@
 )
 
 
-class MissingAttributesTableError(RuntimeError):
-    """Exception raised when a database is missing attributes table.
-    """
-    pass
-
-
 # This manager is supposed to have super-stable schema that never changes
 # but there may be cases when we need data migration on this table so we
 # keep version for it as well.
@@ -86,36 +80,15 @@ def initialize(cls, db: Database, context: StaticTablesContext) -> ButlerAttribu
         table = context.addTable(cls._TABLE_NAME, cls._TABLE_SPEC)
         return cls(db=db, table=table)
 
-    def _checkTableExists(self) -> None:
-        """Check that attributes table exists or raise an exception.
-
-        This should not be called on every read operation but instead only
-        when an exception is raised from SELECT query.
-        """
-        # Database metadata is likely not populated at this point yet, so we
-        # have to use low-level connection object to check it.
-        # TODO: Once we have stable gen3 schema everywhere this test can be
-        # dropped (DM-27373).
-        if not self._table.exists(bind=self._db._connection):
-            raise MissingAttributesTableError(
-                f"`{self._table.name}` table is missing from schema, schema has to"
-                " be initialized before use (database is probably outdated)."
-            ) from None
-
     def get(self, name: str, default: Optional[str] = None) -> Optional[str]:
         # Docstring inherited from ButlerAttributeManager.
-        try:
-            sql = sqlalchemy.sql.select([self._table.columns.value]).where(
-                self._table.columns.name == name
-            )
-            row = self._db.query(sql).fetchone()
-            if row is not None:
-                return row[0]
-            return default
-        except sqlalchemy.exc.OperationalError:
-            # if this is due to missing table raise different exception
-            self._checkTableExists()
-            raise
+        sql = sqlalchemy.sql.select([self._table.columns.value]).where(
+            self._table.columns.name == name
+        )
+        row = self._db.query(sql).fetchone()
+        if row is not None:
+            return row[0]
+        return default
 
     def set(self, name: str, value: str, *, force: bool = False) -> None:
         # Docstring inherited from ButlerAttributeManager.
@@ -151,14 +124,9 @@ def items(self) -> Iterable[Tuple[str, str]]:
 
     def empty(self) -> bool:
         # Docstring inherited from ButlerAttributeManager.
-        try:
-            sql = sqlalchemy.sql.select([sqlalchemy.sql.func.count()]).select_from(self._table)
-            row = self._db.query(sql).fetchone()
-            return row[0] == 0
-        except sqlalchemy.exc.OperationalError:
-            # if this is due to missing table raise different exception
-            self._checkTableExists()
-            raise
+        sql = sqlalchemy.sql.select([sqlalchemy.sql.func.count()]).select_from(self._table)
+        row = self._db.query(sql).fetchone()
+        return row[0] == 0
 
     @classmethod
     def currentVersion(cls) -> Optional[VersionTuple]:

python/lsst/daf/butler/registry/interfaces/_database.py

@@ -1303,7 +1303,12 @@ def insert(self, table: sqlalchemy.schema.Table, *rows: dict, returnIds: bool =
         if not returnIds:
             if select is not None:
                 if names is None:
-                    names = select.columns.keys()
+                    # columns() is deprecated since 1.4, but
+                    # selected_columns() method did not exist in 1.3.
+                    if hasattr(select, "selected_columns"):
+                        names = select.selected_columns.keys()
+                    else:
+                        names = select.columns.keys()
                 self._connection.execute(table.insert().from_select(names, select))
             else:
                 self._connection.execute(table.insert(), *rows)

python/lsst/daf/butler/registry/queries/_query.py

@@ -28,6 +28,7 @@
 import itertools
 from typing import (
     Callable,
+    Dict,
     Iterable,
     Iterator,
     Mapping,
@@ -632,14 +633,21 @@ def __init__(self, *,
         self._simpleQuery = simpleQuery
         self._columns = columns
         self._uniqueness = uniqueness
+        self._datasetQueryColumns: Optional[DatasetQueryColumns] = None
+        self._dimensionColumns: Dict[str, sqlalchemy.sql.ColumnElement] = {}
+        self._regionColumns: Dict[str, sqlalchemy.sql.ColumnElement] = {}
 
     def isUnique(self) -> bool:
         # Docstring inherited from Query.
         return self._uniqueness is not DirectQueryUniqueness.NOT_UNIQUE
 
     def getDimensionColumn(self, name: str) -> sqlalchemy.sql.ColumnElement:
         # Docstring inherited from Query.
-        return self._columns.getKeyColumn(name).label(name)
+        column = self._dimensionColumns.get(name)
+        if column is None:
+            column = self._columns.getKeyColumn(name).label(name)
+            self._dimensionColumns[name] = column
+        return column
 
     @property
     def spatial(self) -> Iterator[DimensionElement]:
@@ -648,22 +656,28 @@ def spatial(self) -> Iterator[DimensionElement]:
 
     def getRegionColumn(self, name: str) -> sqlalchemy.sql.ColumnElement:
         # Docstring inherited from Query.
-        return self._columns.regions[name].column.label(f"{name}_region")
+        column = self._regionColumns.get(name)
+        if column is None:
+            column = self._columns.regions[name].column.label(f"{name}_region")
+            self._regionColumns[name] = column
+        return column
 
     def getDatasetColumns(self) -> Optional[DatasetQueryColumns]:
         # Docstring inherited from Query.
-        base = self._columns.datasets
-        if base is None:
-            return None
-        ingestDate = base.ingestDate
-        if ingestDate is not None:
-            ingestDate = ingestDate.label("ingest_date")
-        return DatasetQueryColumns(
-            datasetType=base.datasetType,
-            id=base.id.label("dataset_id"),
-            runKey=base.runKey.label(self.managers.collections.getRunForeignKeyName()),
-            ingestDate=ingestDate,
-        )
+        if self._datasetQueryColumns is None:
+            base = self._columns.datasets
+            if base is None:
+                return None
+            ingestDate = base.ingestDate
+            if ingestDate is not None:
+                ingestDate = ingestDate.label("ingest_date")
+            self._datasetQueryColumns = DatasetQueryColumns(
+                datasetType=base.datasetType,
+                id=base.id.label("dataset_id"),
+                runKey=base.runKey.label(self.managers.collections.getRunForeignKeyName()),
+                ingestDate=ingestDate,
+            )
+        return self._datasetQueryColumns
 
     @property
     def sql(self) -> sqlalchemy.sql.FromClause:

tests/test_sqlite.py

@@ -30,7 +30,6 @@
 
 from lsst.daf.butler import ddl
 from lsst.daf.butler.registry.databases.sqlite import SqliteDatabase
-from lsst.daf.butler.registry.attributes import MissingAttributesTableError
 from lsst.daf.butler.registry.tests import DatabaseTests, RegistryTests
 from lsst.daf.butler.registry import Registry
 from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
@@ -272,7 +271,7 @@ def testMissingAttributes(self):
         # dropped (DM-27373).
         config = self.makeRegistryConfig()
         config["db"] = "sqlite://"
-        with self.assertRaises(MissingAttributesTableError):
+        with self.assertRaises(sqlalchemy.exc.OperationalError):
             Registry.fromConfig(config)