From 061573001339050f36d8976536d2813a4294a384 Mon Sep 17 00:00:00 2001 From: "David H. Irving" Date: Thu, 28 Mar 2024 14:38:57 -0700 Subject: [PATCH] Add migration script for universe 7 --- .../migrations/dimensions-config.rst | 9 +- migrations/dimensions-config/352c30854bb0.py | 118 ++++++++++++++++++ .../daf/butler_migrate/migration_context.py | 12 ++ 3 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 migrations/dimensions-config/352c30854bb0.py diff --git a/doc/lsst.daf.butler_migrate/migrations/dimensions-config.rst b/doc/lsst.daf.butler_migrate/migrations/dimensions-config.rst index 779e08a..f608c07 100644 --- a/doc/lsst.daf.butler_migrate/migrations/dimensions-config.rst +++ b/doc/lsst.daf.butler_migrate/migrations/dimensions-config.rst @@ -85,4 +85,11 @@ Supports group and day_obs as dimensions. - Rename ``group_name`` in the exposure table to ``group``. - Update the ``exposure`` table so ``group`` and ``day_obs`` are foreign keys to the new tables. - Remove ``group_id`` from ``exposure`` table. -- Update ``config:dimensions.json`` to universe 6. \ No newline at end of file +- Update ``config:dimensions.json`` to universe 6. + +daf_butler 6 to 7 +================= + +Migration script: `352c30854bb0.py `_ + +Adds ``can_see_sky`` column to the ``exposure`` table. \ No newline at end of file diff --git a/migrations/dimensions-config/352c30854bb0.py b/migrations/dimensions-config/352c30854bb0.py new file mode 100644 index 0000000..e3d8184 --- /dev/null +++ b/migrations/dimensions-config/352c30854bb0.py @@ -0,0 +1,118 @@ +"""Migration script for dimensions.yaml namespace=daf_butler version=7. + +Revision ID: 352c30854bb0 +Revises: 1fae088c80b6 +Create Date: 2024-03-28 12:14:53.021101 + +""" + +import logging + +import alembic +import sqlalchemy as sa +from alembic import op +from lsst.daf.butler_migrate.migration_context import MigrationContext + +# revision identifiers, used by Alembic. +revision = "352c30854bb0" +down_revision = "1fae088c80b6" +branch_labels = None +depends_on = None + +# Logger name should start with lsst to work with butler logging option. +_LOG = logging.getLogger(f"lsst.{__name__}") + +_NEW_COLUMN = "can_see_sky" +_TABLE = "exposure" +_OBSERVATION_TYPE_COLUMN = "observation_type" + + +def upgrade() -> None: + """Upgrade from daf_butler universe version 6 to version 7 following update + of dimensions.yaml in DM-43101. + + Adds ``can_see_sky`` column to the exposure table, and sets its initial + values based on the the ``observation_type`` column. + """ + + ctx = MigrationContext() + + _LOG.info("Checking that this is an unmodified daf_butler universe 6 repo") + ctx.attributes.validate_dimensions_json(6) + + _LOG.info("Adding can_see_sky column to exposure table") + op.add_column(_TABLE, sa.Column(_NEW_COLUMN, sa.Boolean, nullable=True), schema=ctx.schema) + + # Set values for existing data based on the exposure's observation_type, + # which is closely correlated with whether the sky is visible in the + # exposure. + # + # Any exposure types not in these lists are left null. + _LOG.info("Populating can_see_sky column") + table = ctx.get_table(_TABLE) + _populate_values( + table, + True, + [ + "science", + "object", + "standard", + "sky flat", + "standard_star", + "skyflat", + "focus", + "focusing", + "exp", + "skyexp", + ], + ) + _populate_values(table, False, ["dark", "bias", "agexp", "domeflat", "dome flat", "zero", "spot"]) + + unhandled_observation_types = _find_unhandled_observation_types(ctx) + if unhandled_observation_types: + _LOG.info( + "WARNING: No default value for can_see_sky is known for the following observation types:\n" + f"{unhandled_observation_types}\n" + "Exposure records with these observation types will have a NULL can_see_sky." + ) + else: + _LOG.info("...can_see_sky values were set for all exposure records.") + + _LOG.info("Updating dimensions.json in ButlerAttributes") + ctx.attributes.replace_dimensions_json(7) + + +def downgrade() -> None: + """Perform schema downgrade.""" + ctx = MigrationContext() + + _LOG.info("Checking that this is an unmodified daf_butler universe 7 repo") + ctx.attributes.validate_dimensions_json(7) + + _LOG.info("dropping can_see_sky column") + op.drop_column(_TABLE, _NEW_COLUMN, schema=ctx.schema) + + _LOG.info("Updating dimensions.json in ButlerAttributes") + ctx.attributes.replace_dimensions_json(6) + + +def _populate_values(table: sa.Table, can_see_sky: bool, observation_types: list[str]) -> None: + op.execute( + table.update() + .values({_NEW_COLUMN: can_see_sky}) + .where(table.columns[_OBSERVATION_TYPE_COLUMN].in_(observation_types)) + .where(table.columns[_NEW_COLUMN].is_(None)) + ) + + +def _find_unhandled_observation_types(ctx: MigrationContext) -> list[str]: + table = ctx.get_table(_TABLE) + return list( + ctx.bind.execute( + sa.select(table.columns[_OBSERVATION_TYPE_COLUMN]) + .distinct() + .where(table.columns[_NEW_COLUMN].is_(None)) + ) + .scalars() + .all() + ) diff --git a/python/lsst/daf/butler_migrate/migration_context.py b/python/lsst/daf/butler_migrate/migration_context.py index 39a3ef2..b46a248 100644 --- a/python/lsst/daf/butler_migrate/migration_context.py +++ b/python/lsst/daf/butler_migrate/migration_context.py @@ -46,4 +46,16 @@ def __init__(self) -> None: self.attributes = ButlerAttributes(self.bind, self.schema) def get_table(self, table_name: str) -> sa.Table: + """Create a SQLAlchemy table object for the current database. + + Parameters + ---------- + table_name : `str` + Name of the table. + + Returns + ------- + table : ``sqlalchemy.Table`` + Table object. + """ return sa.schema.Table(table_name, self.metadata, autoload_with=self.bind, schema=self.schema)