migration for initial matrix population

shared/workspaces/alembic/versions/84c6626eb6cd_populate_initial_matrix.py

+"""populate initial matrix
+
+Revision ID: 84c6626eb6cd
+Revises: d822cf6844bb
+Create Date: 2024-05-16 15:19:40.127244
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "84c6626eb6cd"
+down_revision = "d822cf6844bb"
+branch_labels = None
+depends_on = None
+
+versions = [
+    ("casa-release-4.7.0-1", False),
+    ("casa-release-4.7.0", False),
+    ("casa-release-4.7.1", False),
+    ("casa-release-4.7.2", False),
+    ("casa-release-5.1.0-68", True),
+    ("casa-release-5.1.0-69", True),
+    ("casa-release-5.1.0-71", True),
+    ("casa-release-5.1.0-73", True),
+    ("casa-release-5.1.0-74", True),
+    ("casa-release-5.1.1-5", True),
+    ("casa-release-5.1.2-4", True),
+    ("casa-release-5.4.0-68", False),
+    ("casa-release-5.4.0-70", False),
+    ("casa-release-5.4.1-31", True),
+    ("casa-release-5.4.1-32", True),
+    ("casa-release-5.4.2-5", True),
+    ("casa-release-5.4.2-8", True),
+    ("casa-pipeline-release-5.6.1-8", True),
+    ("casa-pipeline-release-5.6.2-2", True),
+    ("casa-pipeline-release-5.6.2-3", True),
+    ("casa-pipeline-release-5.6.2-6", True),
+    ("casa-pipeline-release-5.6.3-19", True),
+    ("casa-6.1.1-13-pipeline-2020.1.0.36", False),
+    ("casa-6.1.1-15-pipeline-2020.1.0.40", False),
+    ("casa-6.1.2-7-pipeline-2020.1.0.36", False),
+    ("casa-6.1.2-7-pipeline-2020.1.0.40", False),
+    ("casa-6.2.1-7-pipeline-2021.2.0.128", False),
+    ("casa-6.4.1-12-pipeline-2022.2.0.64", True),
+    ("casa-6.4.1-12-pipeline-2022.2.0.68", True),
+    ("casa-6.5.3-28-pipeline-2023.0.0.36", True),
+    ("casa-6.5.4-9-pipeline-2023.1.0.124", True),
+    ("casa-6.5.4-9-pipeline-2023.1.0.125", True),
+]
+
+
+def upgrade():
+    for version in versions:
+        # insert into casa versions table
+        op.execute(
+            f"""
+        INSERT INTO casa_matrix_casa_versions (casa_version, is_cluster_compatible)
+        VALUES ('{version[0]}', {version[1]})
+        """
+        )
+
+        # if cluster compatible, insert valid capabilities
+        if version[1] is True:
+            # calibrated MS restores
+            op.execute(
+                f"""
+                INSERT INTO casa_matrix_capabilities
+                SELECT 'restore_cms' as capability_name,
+                        cv.matrix_id as matrix_id
+                FROM casa_matrix_casa_versions cv
+                WHERE cv.casa_version = '{version[0]}'            
+                """
+            )
+            # basic MS restores
+            op.execute(
+                f"""
+                INSERT INTO casa_matrix_capabilities
+                SELECT 'restore_basic_ms' as capability_name,
+                        cv.matrix_id as matrix_id
+                FROM casa_matrix_casa_versions cv
+                WHERE cv.casa_version = '{version[0]}'            
+                """
+            )
+
+            # if this is a CASA6 release we can also run WS std calibration or imaging capabilities
+            if "casa-6" in version[0]:
+                # standard calibration
+                op.execute(
+                    f"""
+                    INSERT INTO casa_matrix_capabilities
+                    SELECT 'std_calibration' as capability_name,
+                            cv.matrix_id as matrix_id
+                    FROM casa_matrix_casa_versions cv
+                    WHERE cv.casa_version = '{version[0]}'            
+                    """
+                )
+                # standard cms imaging
+                op.execute(
+                    f"""
+                    INSERT INTO casa_matrix_capabilities
+                    SELECT 'std_cms_imaging' as capability_name,
+                            cv.matrix_id as matrix_id
+                    FROM casa_matrix_casa_versions cv
+                    WHERE cv.casa_version = '{version[0]}'            
+                    """
+                )
+                # standard restore imaging
+                op.execute(
+                    f"""
+                    INSERT INTO casa_matrix_capabilities
+                    SELECT 'std_restore_imaging' as capability_name,
+                            cv.matrix_id as matrix_id
+                    FROM casa_matrix_casa_versions cv
+                    WHERE cv.casa_version = '{version[0]}'            
+                    """
+                )
+
+
+def downgrade():
+    op.execute(
+        """
+        TRUNCATE casa_matrix_capabilities, casa_matrix_casa_versions
+        """
+    )