climate-ref 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

climate_ref/executor/pbs_scheduler.py

@@ -0,0 +1,152 @@
+import re
+import shutil
+import subprocess
+import textwrap
+from typing import Any
+
+from parsl.launchers import SimpleLauncher
+from parsl.providers import PBSProProvider
+
+
+class SmartPBSProvider(PBSProProvider):
+    """
+    A PBSProProvider subclass that adapts to systems where `-l select` is not supported.
+
+    Falls back to individual resource requests (ncpus, mem, jobfs, storage) if needed.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        account: str | None = None,
+        queue: str | None = None,
+        scheduler_options: str = "",
+        worker_init: str = "",
+        nodes_per_block: int | None = 1,
+        cpus_per_node: int | None = 1,
+        ncpus: int | None = None,
+        mem: str = "4GB",
+        jobfs: str = "10GB",
+        storage: str = "",
+        init_blocks: int = 1,
+        min_blocks: int = 0,
+        max_blocks: int = 1,
+        parallelism: int = 1,
+        launcher: SimpleLauncher = SimpleLauncher(),
+        walltime: str = "00:20:00",
+        cmd_timeout: int = 120,
+    ) -> None:
+        self.ncpus = ncpus
+        self.mem = mem
+        self.jobfs = jobfs
+        self.storage = storage
+        self._select_supported = self._detect_select_support()
+
+        # Prepare fallback resource dictionary
+        self._fallback_resources = {"mem": mem, "jobfs": jobfs, "storage": storage}
+
+        # Parse and strip select if present in scheduler_options
+        if not self._select_supported and "-l select=" in scheduler_options:
+            scheduler_options = self._parse_select_from_scheduler_options(scheduler_options)
+
+        # Determine fallback ncpus
+        if "ncpus" not in self._fallback_resources:
+            self._fallback_resources["ncpus"] = str(ncpus if ncpus is not None else (cpus_per_node or 1))
+
+        # Map ncpus to cpus_per_node if needed (select mode only)
+        if self._select_supported:
+            if not ncpus and cpus_per_node:
+                cpus_per_node = ncpus
+            elif ncpus and cpus_per_node and int(ncpus) != int(cpus_per_node):
+                print(f"Warning: ncpus={ncpus} and cpus_per_node={cpus_per_node} differ.")
+                print(f"Using cpus_per_node={cpus_per_node}.")
+        else:
+            cpus_per_node = int(self._fallback_resources["ncpus"])
+
+        super().__init__(
+            account=account,
+            queue=queue,
+            scheduler_options=scheduler_options,
+            select_options="",  # Not used; we handle resources ourselves
+            worker_init=worker_init,
+            nodes_per_block=nodes_per_block,
+            cpus_per_node=cpus_per_node,
+            init_blocks=init_blocks,
+            min_blocks=min_blocks,
+            max_blocks=max_blocks,
+            parallelism=parallelism,
+            launcher=launcher,
+            walltime=walltime,
+            cmd_timeout=cmd_timeout,
+        )  # type: ignore
+
+        if not self._select_supported:
+            self.template_string = self._fallback_template()
+
+    def _detect_select_support(self) -> bool:
+        """Detect whether `-l select` is supported by the underlying PBS system."""
+        qsub_path = shutil.which("qsub")
+        if qsub_path is None:
+            raise RuntimeError("qsub command not found. Ensure PBS is installed and in PATH.")
+
+        try:
+            result = subprocess.run(  # noqa: S603
+                [qsub_path, "-l", "wd,select=1:ncpus=1", "--version"],
+                capture_output=True,
+                timeout=5,
+                check=False,
+            )
+            stderr = result.stderr.decode().lower()
+            return "unknown" not in stderr and result.returncode == 0
+        except Exception:
+            return False
+
+    def _parse_select_from_scheduler_options(self, scheduler_options: str) -> str:
+        """
+        Parse `-l select=...` from scheduler_options and update fallback resources.
+
+        Removes the select line from scheduler_options.
+        """
+        select_pattern = r"-l\s+select=([^\s]+)"
+        match = re.search(select_pattern, scheduler_options)
+        if match:
+            select_string = match.group(1)
+            scheduler_options = re.sub(select_pattern, "", scheduler_options).strip()
+
+            parts = select_string.split(":")[1:]  # skip the initial `select=1`
+            for part in parts:
+                if "=" in part:
+                    key, val = part.split("=")
+                    self._fallback_resources[key.strip()] = val.strip()
+        return scheduler_options
+
+    def _fallback_template(self) -> str:
+        """Submit script template used if `select` is not supported."""
+        return textwrap.dedent("""\
+            #!/bin/bash
+            #PBS -N ${jobname}
+            #PBS -l ncpus=${ncpus}
+            #PBS -l mem=${mem}
+            #PBS -l jobfs=${jobfs}
+            #PBS -l walltime=${walltime}
+            #PBS -l storage=${storage}
+            #PBS -o ${job_stdout_path}
+            #PBS -e ${job_stderr_path}
+            ${scheduler_options}
+
+            ${worker_init}
+
+            export JOBNAME="${jobname}"
+            ${user_script}
+
+        """)
+
+    def _write_submit_script(
+        self, template: str, script_filename: str, job_name: str, configs: dict[str, Any]
+    ) -> str:
+        """Inject fallback values into the submit script if `select` is not supported."""
+        if not self._select_supported:
+            configs.setdefault("ncpus", self._fallback_resources.get("ncpus", "1"))
+            configs.setdefault("mem", self._fallback_resources.get("mem", "4GB"))
+            configs.setdefault("jobfs", self._fallback_resources.get("jobfs", "10GB"))
+            configs.setdefault("storage", self._fallback_resources.get("storage", "gdata1"))
+        return super()._write_submit_script(template, script_filename, job_name, configs)  # type: ignore

climate_ref/executor/result_handling.py

@@ -17,11 +17,12 @@ from loguru import logger
 from sqlalchemy import insert
 
 from climate_ref.database import Database
-from climate_ref.models import ScalarMetricValue
+from climate_ref.models import ScalarMetricValue, SeriesMetricValue
 from climate_ref.models.execution import Execution, ExecutionOutput, ResultOutputType
 from climate_ref_core.diagnostics import ExecutionResult, ensure_relative_path
 from climate_ref_core.exceptions import ResultValidationError
 from climate_ref_core.logging import EXECUTION_LOG_FILENAME
+from climate_ref_core.metric_values import SeriesMetricValue as TSeries
 from climate_ref_core.pycmec.controlled_vocabulary import CV
 from climate_ref_core.pycmec.metric import CMECMetric
 from climate_ref_core.pycmec.output import CMECOutput, OutputDict
@@ -65,6 +66,111 @@ def _copy_file_to_results(
     shutil.copy(input_directory / filename, output_filename)
 
 
+def _process_execution_scalar(
+    database: Database,
+    result: "ExecutionResult",
+    execution: Execution,
+    cv: CV,
+) -> None:
+    """
+    Process the scalar values from the execution result and store them in the database
+
+    This also validates the scalar values against the controlled vocabulary
+    """
+    # Load the metric bundle from the file
+    cmec_metric_bundle = CMECMetric.load_from_json(result.to_output_path(result.metric_bundle_filename))
+
+    # Check that the diagnostic values conform with the controlled vocabulary
+    try:
+        cv.validate_metrics(cmec_metric_bundle)
+    except (ResultValidationError, AssertionError):
+        # TODO: Remove once we have settled on a controlled vocabulary
+        logger.exception("Diagnostic values do not conform with the controlled vocabulary")
+        # execution.mark_failed()
+
+    # Perform a bulk insert of scalar values
+    # The current implementation will swallow the exception, but display a log message
+    try:
+        scalar_values = [
+            {
+                "execution_id": execution.id,
+                "value": result.value,
+                "attributes": result.attributes,
+                **result.dimensions,
+            }
+            for result in cmec_metric_bundle.iter_results()
+        ]
+        logger.debug(f"Ingesting {len(scalar_values)} scalar values for execution {execution.id}")
+        if scalar_values:
+            # Perform this in a nested transaction to rollback if something goes wrong
+            # We will lose the metric values for a given execution, but not the whole execution
+            with database.session.begin_nested():
+                database.session.execute(
+                    insert(ScalarMetricValue),
+                    scalar_values,
+                )
+    # This is a broad exception catch to ensure we log any issues
+    except Exception:
+        logger.exception("Something went wrong when ingesting diagnostic values")
+
+
+def _process_execution_series(
+    config: "Config",
+    database: Database,
+    result: "ExecutionResult",
+    execution: Execution,
+    cv: CV,
+) -> None:
+    """
+    Process the series values from the execution result and store them in the database
+
+    This also copies the series values file from the scratch directory to the results directory
+    and validates the series values against the controlled vocabulary.
+    """
+    assert result.series_filename, "Series filename must be set in the result"
+
+    _copy_file_to_results(
+        config.paths.scratch,
+        config.paths.results,
+        execution.output_fragment,
+        result.series_filename,
+    )
+
+    # Load the series values from the file
+    series_values_path = result.to_output_path(result.series_filename)
+    series_values = TSeries.load_from_json(series_values_path)
+
+    try:
+        cv.validate_metrics(series_values)
+    except (ResultValidationError, AssertionError):
+        # TODO: Remove once we have settled on a controlled vocabulary
+        logger.exception("Diagnostic values do not conform with the controlled vocabulary")
+        # execution.mark_failed()
+
+    # Perform a bulk insert of series values
+    try:
+        series_values_content = [
+            {
+                "execution_id": execution.id,
+                "values": series_result.values,
+                "attributes": series_result.attributes,
+                **series_result.dimensions,
+            }
+            for series_result in series_values
+        ]
+        logger.debug(f"Ingesting {len(series_values)} series values for execution {execution.id}")
+        if series_values:
+            # Perform this in a nested transaction to rollback if something goes wrong
+            # We will lose the metric values for a given execution, but not the whole execution
+            with database.session.begin_nested():
+                database.session.execute(
+                    insert(SeriesMetricValue),
+                    series_values_content,
+                )
+    except Exception:
+        logger.exception("Something went wrong when ingesting diagnostic series values")
+
+
 def handle_execution_result(
     config: "Config",
     database: Database,
@@ -88,7 +194,7 @@ def handle_execution_result(
     result
         The result of the diagnostic execution, either successful or failed
     """
-    # Always copy log data
+    # Always copy log data to the results directory
     _copy_file_to_results(
         config.paths.scratch,
         config.paths.results,
@@ -96,74 +202,52 @@ def handle_execution_result(
         EXECUTION_LOG_FILENAME,
     )
 
-    if result.successful and result.metric_bundle_filename is not None:
-        logger.info(f"{execution} successful")
+    if not result.successful or result.metric_bundle_filename is None:
+        logger.error(f"{execution} failed")
+        execution.mark_failed()
+        return
+
+    logger.info(f"{execution} successful")
+
+    _copy_file_to_results(
+        config.paths.scratch,
+        config.paths.results,
+        execution.output_fragment,
+        result.metric_bundle_filename,
+    )
 
+    if result.output_bundle_filename:
         _copy_file_to_results(
             config.paths.scratch,
             config.paths.results,
             execution.output_fragment,
-            result.metric_bundle_filename,
+            result.output_bundle_filename,
+        )
+        _handle_output_bundle(
+            config,
+            database,
+            execution,
+            result.to_output_path(result.output_bundle_filename),
         )
-        execution.mark_successful(result.as_relative_path(result.metric_bundle_filename))
-
-        if result.output_bundle_filename:
-            _copy_file_to_results(
-                config.paths.scratch,
-                config.paths.results,
-                execution.output_fragment,
-                result.output_bundle_filename,
-            )
-            _handle_output_bundle(
-                config,
-                database,
-                execution,
-                result.to_output_path(result.output_bundle_filename),
-            )
 
-        cmec_metric_bundle = CMECMetric.load_from_json(result.to_output_path(result.metric_bundle_filename))
-
-        # Check that the diagnostic values conform with the controlled vocabulary
-        try:
-            cv = CV.load_from_file(config.paths.dimensions_cv)
-            cv.validate_metrics(cmec_metric_bundle)
-        except (ResultValidationError, AssertionError):
-            logger.exception("Diagnostic values do not conform with the controlled vocabulary")
-            # execution.mark_failed()
-
-        # Perform a bulk insert of scalar values
-        # The current implementation will swallow the exception, but display a log message
-        try:
-            scalar_values = [
-                {
-                    "execution_id": execution.id,
-                    "value": result.value,
-                    "attributes": result.attributes,
-                    **result.dimensions,
-                }
-                for result in cmec_metric_bundle.iter_results()
-            ]
-            if scalar_values:
-                # Perform this in a nested transaction to rollback if something goes wrong
-                # We will lose the metric values for a given execution, but not the whole execution
-                with database.session.begin_nested():
-                    database.session.execute(
-                        insert(ScalarMetricValue),
-                        scalar_values,
-                    )
-        except Exception:
-            # TODO: Remove once we have settled on a controlled vocabulary
-            logger.exception("Something went wrong when ingesting diagnostic values")
-
-        # TODO Ingest the series values
-
-        # TODO: This should check if the result is the most recent for the execution,
-        # if so then update the dirty fields
-        # i.e. if there are outstanding executions don't make as clean
-        execution.execution_group.dirty = False
-    else:
-        logger.error(f"{execution} failed")
-        execution.mark_failed()
+    cv = CV.load_from_file(config.paths.dimensions_cv)
+
+    if result.series_filename:
+        # Process the series values if they are present
+        # This will ingest the series values into the database
+        _process_execution_series(config=config, database=database, result=result, execution=execution, cv=cv)
+
+    # Process the scalar values
+    # This will ingest the scalar values into the database
+    _process_execution_scalar(database=database, result=result, execution=execution, cv=cv)
+
+    # TODO: This should check if the result is the most recent for the execution,
+    # if so then update the dirty fields
+    # i.e. if there are outstanding executions don't make as clean
+    execution.execution_group.dirty = False
+
+    # Finally, mark the execution as successful
+    execution.mark_successful(result.as_relative_path(result.metric_bundle_filename))
 
 
 def _handle_output_bundle(

climate_ref/migrations/versions/2025-07-20T1521_94beace57a9c_cmip6_finalised.py

@@ -0,0 +1,57 @@
+"""cmip6-finalised
+
+Revision ID: 94beace57a9c
+Revises: 795c1e6cf496
+Create Date: 2025-07-20 15:21:17.132458
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "94beace57a9c"
+down_revision: Union[str, None] = "795c1e6cf496"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("finalised", sa.Boolean(), nullable=True))
+        batch_op.alter_column("experiment", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("frequency", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("grid", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("nominal_resolution", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("realm", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("product", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("standard_name", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("source_type", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("sub_experiment", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("sub_experiment_id", existing_type=sa.VARCHAR(), nullable=True)
+        batch_op.alter_column("units", existing_type=sa.VARCHAR(), nullable=True)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+        batch_op.alter_column("units", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("sub_experiment_id", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("sub_experiment", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("source_type", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("standard_name", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("product", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("realm", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("nominal_resolution", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("grid", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("frequency", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.alter_column("experiment", existing_type=sa.VARCHAR(), nullable=False)
+        batch_op.drop_column("finalised")
+
+    # ### end Alembic commands ###

climate_ref/migrations/versions/2025-08-05T0327_a1b2c3d4e5f6_finalised_on_base_dataset.py

@@ -0,0 +1,57 @@
+"""finalised-on-base-dataset
+
+Move finalised from cmip6_dataset to base dataset table and default all existing rows to True.
+
+Revision ID: a1b2c3d4e5f6
+Revises: 94beace57a9c
+Create Date: 2025-08-05 03:27:00
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "ba5e"
+down_revision: Union[str, None] = "94beace57a9c"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Add finalised to base dataset with default True, non-null
+    with op.batch_alter_table("dataset", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("finalised", sa.Boolean(), nullable=True, server_default=sa.text("true"))
+        )
+
+    # Backfill: ensure all existing rows are True
+    op.execute("UPDATE dataset SET finalised = TRUE WHERE finalised IS NULL")
+
+    # Enforce NOT NULL after backfill
+    with op.batch_alter_table("dataset", schema=None) as batch_op:
+        batch_op.alter_column("finalised", nullable=False)
+
+    # Drop column from cmip6_dataset if it exists
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
+    cmip6_cols = {col["name"] for col in inspector.get_columns("cmip6_dataset")}
+    if "finalised" in cmip6_cols:
+        with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+            batch_op.drop_column("finalised")
+
+
+def downgrade() -> None:
+    # Re-create cmip6_dataset.finalised as non-nullable boolean default False
+    # Note: Original migration 94beace57a9c added cmip6_dataset.finalised NOT NULL, with no default.
+    with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("finalised", sa.Boolean(), nullable=True, server_default=sa.text("false"))
+        )
+
+    # Drop base dataset finalised
+    with op.batch_alter_table("dataset", schema=None) as batch_op:
+        batch_op.drop_column("finalised")

climate_ref/migrations/versions/2025-09-05T2019_8d28e5e0f9c3_add_indexes.py

@@ -0,0 +1,108 @@
+"""add indexes
+
+Revision ID: 8d28e5e0f9c3
+Revises: ba5e
+Create Date: 2025-09-05 20:19:18.311472
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "8d28e5e0f9c3"
+down_revision: Union[str, None] = "ba5e"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_cmip6_dataset_experiment_id"), ["experiment_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_cmip6_dataset_instance_id"), ["instance_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_cmip6_dataset_member_id"), ["member_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_cmip6_dataset_source_id"), ["source_id"], unique=False)
+
+    with op.batch_alter_table("dataset", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_dataset_dataset_type"), ["dataset_type"], unique=False)
+
+    with op.batch_alter_table("dataset_file", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_dataset_file_dataset_id"), ["dataset_id"], unique=False)
+
+    with op.batch_alter_table("diagnostic", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_diagnostic_updated_at"), ["updated_at"], unique=False)
+
+    with op.batch_alter_table("execution", schema=None) as batch_op:
+        batch_op.create_index(
+            batch_op.f("ix_execution_execution_group_id"), ["execution_group_id"], unique=False
+        )
+        batch_op.create_index(batch_op.f("ix_execution_successful"), ["successful"], unique=False)
+        batch_op.create_index(batch_op.f("ix_execution_updated_at"), ["updated_at"], unique=False)
+
+    with op.batch_alter_table("execution_dataset", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_execution_dataset_dataset_id"), ["dataset_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_execution_dataset_execution_id"), ["execution_id"], unique=False)
+
+    with op.batch_alter_table("execution_group", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_execution_group_diagnostic_id"), ["diagnostic_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_execution_group_updated_at"), ["updated_at"], unique=False)
+
+    with op.batch_alter_table("execution_output", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_execution_output_updated_at"), ["updated_at"], unique=False)
+
+    with op.batch_alter_table("metric_value", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_metric_value_execution_id"), ["execution_id"], unique=False)
+        batch_op.create_index(batch_op.f("ix_metric_value_type"), ["type"], unique=False)
+        batch_op.create_index(batch_op.f("ix_metric_value_updated_at"), ["updated_at"], unique=False)
+
+    with op.batch_alter_table("provider", schema=None) as batch_op:
+        batch_op.create_index(batch_op.f("ix_provider_updated_at"), ["updated_at"], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("provider", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_provider_updated_at"))
+
+    with op.batch_alter_table("metric_value", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_metric_value_updated_at"))
+        batch_op.drop_index(batch_op.f("ix_metric_value_type"))
+        batch_op.drop_index(batch_op.f("ix_metric_value_execution_id"))
+
+    with op.batch_alter_table("execution_output", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_execution_output_updated_at"))
+
+    with op.batch_alter_table("execution_group", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_execution_group_updated_at"))
+        batch_op.drop_index(batch_op.f("ix_execution_group_diagnostic_id"))
+
+    with op.batch_alter_table("execution_dataset", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_execution_dataset_execution_id"))
+        batch_op.drop_index(batch_op.f("ix_execution_dataset_dataset_id"))
+
+    with op.batch_alter_table("execution", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_execution_updated_at"))
+        batch_op.drop_index(batch_op.f("ix_execution_successful"))
+        batch_op.drop_index(batch_op.f("ix_execution_execution_group_id"))
+
+    with op.batch_alter_table("diagnostic", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_diagnostic_updated_at"))
+
+    with op.batch_alter_table("dataset_file", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_dataset_file_dataset_id"))
+
+    with op.batch_alter_table("dataset", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_dataset_dataset_type"))
+
+    with op.batch_alter_table("cmip6_dataset", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_cmip6_dataset_source_id"))
+        batch_op.drop_index(batch_op.f("ix_cmip6_dataset_member_id"))
+        batch_op.drop_index(batch_op.f("ix_cmip6_dataset_instance_id"))
+        batch_op.drop_index(batch_op.f("ix_cmip6_dataset_experiment_id"))
+
+    # ### end Alembic commands ###

climate_ref/models/base.py

@@ -38,7 +38,9 @@ class CreatedUpdatedMixin:
     When the dataset was added to the database
     """
 
-    updated_at: Mapped[datetime.datetime] = mapped_column(server_default=func.now(), onupdate=func.now())
+    updated_at: Mapped[datetime.datetime] = mapped_column(
+        server_default=func.now(), onupdate=func.now(), index=True
+    )
     """
     When the dataset was updated.
     """