DM-45460: Add migration script for ApdbSql 0.1.1

doc/lsst.dax.apdb_migrate/migrations/sql/ApdbSql.rst

@@ -2,4 +2,14 @@
 Migrations for ``ApdbSql`` tree
 ###############################
 
-There are no migration script for this tree yet.
+Upgrade from 0.1.0 to 0.1.1
+===========================
+
+Migration script: `ApdbSql_0.1.1.py <https://github.com/lsst-dm/dax_apdb_migrate/blob/main/migrations/sql/ApdbSql/ApdbSql_0.1.1.py>`_
+
+This schema upgrade changes all TIMESTAMP columns to TIMESTAMP WITH TIMEZONE.
+This is relevant for Postgres only, SQLite does not know about timezones, so the script is no-op for SQLite.
+
+An example command for applying the schema upgrade::
+
+    $ apdb-migrate-sql upgrade -s SCHEMA_NAME $APDB_URL ApdbSql_0.1.1

migrations/sql/ApdbSql/ApdbSql_0.1.1.py

@@ -0,0 +1,61 @@
+"""Migration script for ApdbSql 0.1.1.
+
+Revision ID: ApdbSql_0.1.1
+Revises: ApdbSql_0.1.0
+Create Date: 2024-07-29 11:02:45.959048
+
+"""
+
+import logging
+
+import sqlalchemy
+from lsst.dax.apdb_migrate.sql.context import Context
+
+# revision identifiers, used by Alembic.
+revision = "ApdbSql_0.1.1"
+down_revision = "ApdbSql_0.1.0"
+branch_labels = None
+depends_on = None
+
+_LOG = logging.getLogger(__name__)
+
+
+def upgrade() -> None:
+    """Change all TIMESTAMP columns to TIMESTAMP WITH TIMEZONE."""
+    version = revision.split("_")[1]
+    _migrate(sqlalchemy.types.TIMESTAMP(timezone=True), version)
+
+
+def downgrade() -> None:
+    """Change all TIMESTAMP WITH TIMEZONE columns to TIMESTAMP."""
+    version = down_revision.split("_")[1]
+    _migrate(sqlalchemy.types.TIMESTAMP(timezone=False), version)
+
+
+def _migrate(type_: type, to_version: str) -> None:
+    ctx = Context()
+
+    # SQLite has no TIMEZONE, this migration is not needed.
+    if not ctx.is_sqlite:
+        with ctx.reflection_bind() as bind:
+            inspector = sqlalchemy.inspect(bind)
+            tables = inspector.get_table_names(schema=ctx.schema)
+            _LOG.info("All table names: %s", tables)
+
+            table_columns = {}
+            for table in tables:
+                columns_to_update = []
+                for column in inspector.get_columns(table, schema=ctx.schema):
+                    if isinstance(column["type"], sqlalchemy.types.TIMESTAMP):
+                        columns_to_update.append(column["name"])
+                if columns_to_update:
+                    table_columns[table] = columns_to_update
+
+        for table, columns in table_columns.items():
+            _LOG.info("Updating table %s, columns: %s", table, columns)
+            with ctx.batch_alter_table(table) as batch_op:
+                for column_name in columns:
+                    batch_op.alter_column(column_name, type_=type_)
+
+    # Update metadata version.
+    ctx.apdb_meta.update_tree_version("ApdbSql", to_version)

migrations/sql/ApdbSqlReplica/ApdbSqlReplica_1.0.0.py

@@ -0,0 +1,23 @@
+"""Migration script for ApdbSqlReplica 1.0.0.
+
+Revision ID: ApdbSqlReplica_1.0.0
+Revises: ApdbSqlReplica_root
+Create Date: 2024-07-29 11:49:49.532074
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "ApdbSqlReplica_1.0.0"
+down_revision = "ApdbSqlReplica_root"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade is not needed as this is the initial version."""
+    raise NotImplementedError()
+
+
+def downgrade() -> None:
+    """Downgrade is not needed as this is the initial version."""
+    raise NotImplementedError()

migrations/sql/ApdbSqlReplica/ApdbSqlReplica_root.py

@@ -0,0 +1,23 @@
+"""An initial pseudo-revision of the 'ApdbSqlReplica' tree.
+
+Revision ID: ApdbSqlReplica_root
+Revises:
+Create Date: 2024-07-29 11:47:20.755789
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "ApdbSqlReplica_root"
+down_revision = None
+branch_labels = ("ApdbSqlReplica",)
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade is not needed for this revision."""
+    raise NotImplementedError()
+
+
+def downgrade() -> None:
+    """Downgrade is not needed for this revision."""
+    raise NotImplementedError()

python/lsst/dax/apdb_migrate/sql/context.py

@@ -79,7 +79,7 @@ def get_table(self, table_name: str, reload: bool = False) -> sqlalchemy.Table:
                 if table.name == table_name:
                     self.metadata.remove(table)
                     break
-        with self._reflection_bind() as bind:
+        with self.reflection_bind() as bind:
             return sqlalchemy.schema.Table(table_name, self.metadata, autoload_with=bind, schema=self.schema)
 
     def get_mig_option(self, option: str) -> str | None:
@@ -113,7 +113,7 @@ def batch_alter_table(
         return alembic.op.batch_alter_table(table, schema=self.schema, **kwargs)
 
     @contextlib.contextmanager
-    def _reflection_bind(self) -> Iterator[sqlalchemy.engine.Connection]:
+    def reflection_bind(self) -> Iterator[sqlalchemy.engine.Connection]:
         """Return database connection to be used for reflection. In online mode
         this returns connection instantiated by Alembic, in offline mode it
         creates new engine using configured URL.