Skip to content

Commit

Permalink
Fix deprecated sqlalchemy code (DM-29273)
Browse files Browse the repository at this point in the history
Few fixes for the cases that were triggered by unit tests. There may be
cases that are not covered by unit tests, we'll discover them eventually.
  • Loading branch information
andy-slac committed May 5, 2021
1 parent be73124 commit 2865ab6
Show file tree
Hide file tree
Showing 4 changed files with 45 additions and 59 deletions.
52 changes: 10 additions & 42 deletions python/lsst/daf/butler/registry/attributes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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]:
Expand Down
7 changes: 6 additions & 1 deletion python/lsst/daf/butler/registry/interfaces/_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
42 changes: 28 additions & 14 deletions python/lsst/daf/butler/registry/queries/_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import itertools
from typing import (
Callable,
Dict,
Iterable,
Iterator,
Mapping,
Expand Down Expand Up @@ -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]:
Expand All @@ -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:
Expand Down
3 changes: 1 addition & 2 deletions tests/test_sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)


Expand Down

0 comments on commit 2865ab6

Please sign in to comment.