Index

class AccountingRecord(Base):
    """
    AccountingRecord table.
    """

    __tablename__ = "accountingrecord"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("user_oauth.id"), nullable=False
    )
    timestamp: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    num_tasks: Mapped[int]
    num_new_images: Mapped[int]

class AccountingRecordSlurm(Base):
    """
    AccountingRecordSlurm table.
    """

    __tablename__ = "accountingrecordslurm"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("user_oauth.id"), nullable=False
    )
    timestamp: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    slurm_job_ids: Mapped[list[int]] = mapped_column(
        ARRAY(Integer), nullable=True, default=list
    )
    fractal_job_id: Mapped[int | None] = mapped_column(
        ForeignKey("jobv2.id"), nullable=True, default=lambda: None
    )
    resource_id: Mapped[int | None] = mapped_column(
        ForeignKey("resource.id", ondelete="SET NULL"),
        nullable=True,
        default=lambda: None,
    )

class DatasetV2(Base):
    """
    Dataset table.
    """

    __tablename__ = "datasetv2"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    project_id: Mapped[int] = mapped_column(
        ForeignKey("projectv2.id", ondelete="CASCADE")
    )
    project: Mapped["ProjectV2"] = relationship(lazy="selectin")
    is_starred: Mapped[bool] = mapped_column(
        BOOLEAN,
        server_default="false",
        nullable=False,
    )

    timestamp_created: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )

    zarr_dir: Mapped[str]
    images: Mapped[list[dict[str, Any]]] = mapped_column(
        JSONB, server_default="[]", nullable=False
    )

    @property
    def image_zarr_urls(self) -> list[str]:
        return [image["zarr_url"] for image in self.images]

class HistoryImageCache(Base):
    """
    HistoryImageCache table.
    """

    __tablename__ = "historyimagecache"

    zarr_url: Mapped[str] = mapped_column(primary_key=True)
    dataset_id: Mapped[int] = mapped_column(
        ForeignKey("datasetv2.id", ondelete="CASCADE"),
        primary_key=True,
        index=True,
    )
    workflowtask_id: Mapped[int] = mapped_column(
        ForeignKey("workflowtaskv2.id", ondelete="CASCADE"),
        primary_key=True,
        index=True,
    )

    latest_history_unit_id: Mapped[int] = mapped_column(
        ForeignKey("historyunit.id", ondelete="CASCADE"),
        index=True,
    )

class HistoryRun(Base):
    """
    HistoryRun table.
    """

    __tablename__ = "historyrun"

    id: Mapped[int] = mapped_column(primary_key=True)
    dataset_id: Mapped[int] = mapped_column(
        ForeignKey("datasetv2.id", ondelete="CASCADE"),
    )
    workflowtask_id: Mapped[int | None] = mapped_column(
        ForeignKey("workflowtaskv2.id", ondelete="SET NULL"),
        default=lambda: None,
    )
    job_id: Mapped[int] = mapped_column(ForeignKey("jobv2.id"))
    task_id: Mapped[int | None] = mapped_column(
        ForeignKey("taskv2.id", ondelete="SET NULL")
    )

    workflowtask_dump: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False
    )
    task_group_dump: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False
    )

    timestamp_started: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    status: Mapped[str]
    num_available_images: Mapped[int]

class HistoryUnit(Base):
    """
    HistoryUnit table.
    """

    __tablename__ = "historyunit"

    id: Mapped[int] = mapped_column(primary_key=True)
    history_run_id: Mapped[int] = mapped_column(
        ForeignKey("historyrun.id", ondelete="CASCADE"),
        index=True,
    )

    logfile: Mapped[str]
    has_warnings: Mapped[bool] = mapped_column(default=False)
    status: Mapped[str]
    zarr_urls: Mapped[list[str]] = mapped_column(ARRAY(String), default=list)

class JobV2(Base):
    """
    Job table.
    """

    __tablename__ = "jobv2"

    id: Mapped[int] = mapped_column(primary_key=True)
    project_id: Mapped[int | None] = mapped_column(
        ForeignKey("projectv2.id", ondelete="SET NULL"), default=lambda: None
    )
    workflow_id: Mapped[int | None] = mapped_column(
        ForeignKey("workflowv2.id", ondelete="SET NULL"), default=lambda: None
    )
    dataset_id: Mapped[int | None] = mapped_column(
        ForeignKey("datasetv2.id", ondelete="SET NULL"), default=lambda: None
    )

    user_email: Mapped[str] = mapped_column(nullable=False)
    slurm_account: Mapped[str | None] = mapped_column(default=lambda: None)

    dataset_dump: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    workflow_dump: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    project_dump: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
    fractal_server_version: Mapped[str] = mapped_column(
        String, server_default="pre-2.19.0", nullable=False
    )

    worker_init: Mapped[str | None] = mapped_column(default=lambda: None)
    working_dir: Mapped[str | None] = mapped_column(default=lambda: None)
    working_dir_user: Mapped[str | None] = mapped_column(default=lambda: None)
    first_task_index: Mapped[int]
    last_task_index: Mapped[int]

    start_timestamp: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    end_timestamp: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), default=lambda: None
    )
    status: Mapped[str] = mapped_column(default=JobStatusType.SUBMITTED)
    log: Mapped[str | None] = mapped_column(default=lambda: None)
    executor_error_log: Mapped[str | None] = mapped_column(default=lambda: None)

    attribute_filters: Mapped[dict[str, list[int | float | str | bool]]] = (
        mapped_column(JSONB, nullable=False, server_default="{}")
    )
    type_filters: Mapped[dict[str, bool]] = mapped_column(
        JSONB, nullable=False, server_default="{}"
    )

    __table_args__ = (
        Index(
            "ix_jobv2_one_submitted_job_per_dataset",
            "dataset_id",
            unique=True,
            postgresql_where=text(f"status = '{JobStatusType.SUBMITTED}'"),
        ),
    )

class OAuthAccount(Base):
    """
    ORM model for OAuth accounts (`oauthaccount` database table).

    This class is a modification of
    [`SQLAlchemyBaseOAuthAccountTable`](https://github.com/fastapi-users/fastapi-users-db-sqlalchemy/blob/main/fastapi_users_db_sqlalchemy/__init__.py).
    Original Copyright: 2021 François Voron, released under MIT licence.

    Attributes:
        id:
        user_id:
        user:
        oauth_name:
        access_token:
        expires_at:
        refresh_token:
        account_id:
        account_email:
    """

    __tablename__ = "oauthaccount"

    id: Mapped[int] = mapped_column(primary_key=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("user_oauth.id"), nullable=False
    )
    user: Mapped["UserOAuth | None"] = relationship(
        back_populates="oauth_accounts"
    )
    oauth_name: Mapped[str] = mapped_column(index=True, nullable=False)
    access_token: Mapped[str] = mapped_column(nullable=False)
    expires_at: Mapped[int | None] = mapped_column(
        nullable=True, default=lambda: None
    )
    refresh_token: Mapped[str | None] = mapped_column(
        nullable=True, default=lambda: None
    )
    account_id: Mapped[str] = mapped_column(index=True, nullable=False)
    account_email: Mapped[str] = mapped_column(nullable=False)

class Profile(Base):
    """
    Profile table.
    """

    __tablename__ = "profile"

    id: Mapped[int] = mapped_column(primary_key=True)

    resource_id: Mapped[int] = mapped_column(
        ForeignKey("resource.id", ondelete="RESTRICT")
    )

    resource_type: Mapped[str]
    """
    Type of resource (either `local`, `slurm_sudo` or `slurm_ssh`).
    """

    name: Mapped[str] = mapped_column(unique=True)
    """
    Profile name.
    """

    username: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Username to be impersonated, either via `sudo -u` or via `ssh`.
    """

    ssh_key_path: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Path to the private SSH key of user `username` (only relevant if
    `resource_type="slurm_ssh"`).
    """

    jobs_remote_dir: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Remote path of the job folder (only relevant if
    `resource_type="slurm_ssh"`).
    """

    tasks_remote_dir: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Remote path of the task folder (only relevant if
    `resource_type="slurm_ssh"`).
    """

    pixi_cache_dir: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Override for the `PIXI_CACHE_DIR` variable, which would otherwise default
    to the `cache` subfolder of `PIXI_HOME` (only relevant if
    `resource_type="slurm_ssh"`)
    """

class ProjectV2(Base):
    """
    Project table.
    """

    __tablename__ = "projectv2"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    is_starred: Mapped[bool] = mapped_column(
        BOOLEAN,
        server_default="false",
        nullable=False,
    )
    description: Mapped[str | None] = mapped_column(default=lambda: None)

    resource_id: Mapped[int] = mapped_column(
        ForeignKey("resource.id", ondelete="RESTRICT")
    )
    timestamp_created: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )

class Resource(Base):
    """
    Resource table.
    """

    __tablename__ = "resource"

    id: Mapped[int] = mapped_column(primary_key=True)

    type: Mapped[str]
    """
    One of `local`, `slurm_sudo` or `slurm_ssh` - matching with
    `settings.FRACTAL_RUNNER_BACKEND`.
    """

    name: Mapped[str] = mapped_column(unique=True)
    """
    Resource name.
    """

    timestamp_created: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    """
    Creation timestamp (autogenerated).
    """

    host: Mapped[str | None] = mapped_column(default=lambda: None)
    """
    Address for ssh connections, when `type="slurm_ssh"`.
    """

    prevent_new_submissions: Mapped[bool] = mapped_column(
        BOOLEAN,
        server_default="false",
        nullable=False,
    )
    """
    When set to true: Prevent new job submissions and stop execution of
    ongoing jobs as soon as the current task is complete.
    """

    jobs_local_dir: Mapped[str]
    """
    Base local folder for job subfolders (containing artifacts and logs).
    """

    jobs_runner_config: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False, server_default="{}"
    )
    """
    Runner configuration, matching one of `JobRunnerConfigLocal` or
    `JobRunnerConfigSLURM` schemas.
    """

    jobs_slurm_python_worker: Mapped[str | None] = mapped_column(
        default=lambda: None
    )
    """
    On SLURM deloyments, this is the Python interpreter that runs the
    `fractal-server` worker from within the SLURM jobs.
    """

    jobs_poll_interval: Mapped[int]
    """
    On SLURM resources: the interval to wait before new `squeue` calls.
    On local resources: ignored.
    """

    # task_settings
    tasks_local_dir: Mapped[str]
    """
    Base local folder for task-package subfolders.
    """

    tasks_python_config: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False, server_default="{}"
    )
    """
    Python configuration for task collection. Example:
    ```json
    {
      "default_version": "3.10",
      "versions:{
        "3.10": "/xxx/venv-3.10/bin/python",
        "3.11": "/xxx/venv-3.11/bin/python",
        "3.12": "/xxx/venv-3.12/bin/python"
       }
    }
    ```
    """

    tasks_pixi_config: Mapped[dict[str, Any]] = mapped_column(
        JSONB, nullable=False, server_default="{}"
    )
    """
    Pixi configuration for task collection. Basic example:
    ```json
    {
        "default_version": "0.41.0",
        "versions": {
            "0.40.0": "/xxx/pixi/0.40.0/",
            "0.41.0": "/xxx/pixi/0.41.0/"
        },
    }
    ```
    """

    @property
    def pip_cache_dir_arg(self: Self) -> str:
        """
        If `pip_cache_dir` is set (in `self.tasks_python_config`), then
        return `--cache_dir /something`; else return `--no-cache-dir`.
        """
        _pip_cache_dir = self.tasks_python_config.get("pip_cache_dir", None)
        if _pip_cache_dir is not None:
            return f"--cache-dir {_pip_cache_dir}"
        else:
            return "--no-cache-dir"

    # Check constraints
    __table_args__ = (
        # `type` column must be one of "local", "slurm_sudo" or "slurm_ssh"
        CheckConstraint(
            "type IN ('local', 'slurm_sudo', 'slurm_ssh')",
            name="correct_type",
        ),
        # If `type` is not "local", `jobs_slurm_python_worker` must be set
        CheckConstraint(
            "(type = 'local') OR (jobs_slurm_python_worker IS NOT NULL)",
            name="jobs_slurm_python_worker_set",
        ),
    )

class TaskGroupV2(Base):
    __tablename__ = "taskgroupv2"

    id: Mapped[int] = mapped_column(primary_key=True)
    task_list: Mapped[list["TaskV2"]] = relationship(
        lazy="selectin", cascade="all, delete-orphan"
    )

    user_id: Mapped[int] = mapped_column(ForeignKey("user_oauth.id"))
    user_group_id: Mapped[int | None] = mapped_column(
        ForeignKey("usergroup.id", ondelete="SET NULL"), default=lambda: None
    )
    resource_id: Mapped[int] = mapped_column(
        ForeignKey("resource.id", ondelete="RESTRICT")
    )

    origin: Mapped[str]
    pkg_name: Mapped[str]
    version: Mapped[str]
    python_version: Mapped[str | None] = mapped_column(default=lambda: None)
    pixi_version: Mapped[str | None] = mapped_column(default=lambda: None)
    path: Mapped[str | None] = mapped_column(default=lambda: None)
    archive_path: Mapped[str | None] = mapped_column(default=lambda: None)
    pip_extras: Mapped[str | None] = mapped_column(default=lambda: None)
    pinned_package_versions_pre: Mapped[dict[str, str]] = mapped_column(
        JSONB,
        server_default="{}",
        default={},
        nullable=True,
    )
    pinned_package_versions_post: Mapped[dict[str, str]] = mapped_column(
        JSONB,
        server_default="{}",
        default={},
        nullable=True,
    )
    env_info: Mapped[str | None] = mapped_column(default=lambda: None)
    venv_path: Mapped[str | None] = mapped_column(default=lambda: None)

    active: Mapped[bool] = mapped_column(default=True)
    timestamp_created: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )
    timestamp_last_used: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=(
            datetime(2024, 11, 20, tzinfo=timezone.utc).isoformat()
        ),
        default=get_timestamp,
    )

    __table_args__ = (
        Index(
            "ix_taskgroupv2_path_unique_constraint",
            "path",
            "resource_id",
            unique=True,
        ),
    )

    @property
    def pip_install_string(self) -> str:
        """
        Prepare string to be used in `python -m pip install`.
        """
        _check_origin_not_pixi(self.origin)

        extras = (
            f"[{self.pip_extras}]"
            if (self.pip_extras is not None and self.pip_extras != "")
            else ""
        )

        if self.archive_path is not None:
            return f"{self.archive_path}{extras}"
        else:
            return f"{self.pkg_name}{extras}=={self.version}"

    @property
    def pinned_package_versions_pre_string(self) -> str:
        """
        Prepare string to be used in `python -m pip install`.
        """
        _check_origin_not_pixi(self.origin)

        if self.pinned_package_versions_pre is None:
            return ""
        output = _create_dependency_string(self.pinned_package_versions_pre)
        return output

    @property
    def pinned_package_versions_post_string(self) -> str:
        """
        Prepare string to be used in `python -m pip install`.
        """
        _check_origin_not_pixi(self.origin)

        if self.pinned_package_versions_post is None:
            return ""
        output = _create_dependency_string(self.pinned_package_versions_post)
        return output

class TaskV2(Base):
    """
    Model for the `taskv2` database table.

    Attributes:
        id:
        name:
        type:
        version:
        command_non_parallel:
        command_parallel:
        meta_non_parallel:
        meta_parallel:
        input_types:
        output_types:
        taskgroupv2_id:
        args_schema_version:
        args_schema_non_parallel:
        args_schema_parallel:
        docs_info:
        docs_link:
        category:
        modality:
        authors:
        tags:
    """

    __tablename__ = "taskv2"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    type: Mapped[str]
    command_non_parallel: Mapped[str | None] = mapped_column(
        default=lambda: None
    )
    command_parallel: Mapped[str | None] = mapped_column(default=lambda: None)

    meta_non_parallel: Mapped[dict[str, Any]] = mapped_column(
        JSON, server_default="{}", default={}, nullable=False
    )
    meta_parallel: Mapped[dict[str, Any]] = mapped_column(
        JSON, server_default="{}", default={}, nullable=False
    )

    version: Mapped[str]
    args_schema_non_parallel: Mapped[dict[str, Any] | None] = mapped_column(
        JSON, default=lambda: None
    )
    args_schema_parallel: Mapped[dict[str, Any] | None] = mapped_column(
        JSON, default=lambda: None
    )
    args_schema_version: Mapped[str | None] = mapped_column(
        default=lambda: None
    )
    docs_info: Mapped[str | None] = mapped_column(default=lambda: None)
    docs_link: Mapped[str | None] = mapped_column(default=lambda: None)

    input_types: Mapped[dict[str, bool]] = mapped_column(
        JSONB, nullable=True, default={}
    )
    output_types: Mapped[dict[str, bool]] = mapped_column(
        JSONB, nullable=True, default={}
    )

    taskgroupv2_id: Mapped[int] = mapped_column(ForeignKey("taskgroupv2.id"))

    category: Mapped[str | None] = mapped_column(default=lambda: None)
    modality: Mapped[str | None] = mapped_column(default=lambda: None)
    authors: Mapped[str | None] = mapped_column(default=lambda: None)
    tags: Mapped[list[str]] = mapped_column(
        JSONB, server_default="[]", nullable=False
    )
    is_core: Mapped[bool] = mapped_column(
        BOOLEAN,
        server_default="false",
        nullable=False,
    )
    __table_args__ = (
        Index(
            "ix_taskv2_one_task_name_per_task_group",
            "name",
            "taskgroupv2_id",
            unique=True,
        ),
    )

class UserGroup(Base):
    """
    ORM model for the `usergroup` database table.

    Attributes:
        id: ID of the group
        name: Name of the group
        timestamp_created: Time of creation
    """

    __tablename__ = "usergroup"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(unique=True)
    timestamp_created: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, default=get_timestamp
    )