fractal_server.app.models¶

class AccountingRecord(SQLModel, table=True):
    """
    AccountingRecord table.
    """

    id: int | None = Field(default=None, primary_key=True)
    user_id: int = Field(foreign_key="user_oauth.id", nullable=False)
    timestamp: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )
    num_tasks: int
    num_new_images: int

class AccountingRecordSlurm(SQLModel, table=True):
    """
    AccountingRecordSlurm table.
    """

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

class DatasetV2(SQLModel, table=True):
    """
    Dataset table.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: int | None = Field(default=None, primary_key=True)
    name: str

    project_id: int = Field(foreign_key="projectv2.id", ondelete="CASCADE")
    project: ProjectV2 = Relationship(
        sa_relationship_kwargs=dict(lazy="selectin"),
    )
    is_starred: bool = Field(
        sa_column=Column(
            BOOLEAN,
            server_default="false",
            nullable=False,
        ),
    )

    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )

    zarr_dir: str
    images: list[dict[str, Any]] = Field(
        sa_column=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(SQLModel, table=True):
    """
    HistoryImageCache table.
    """

    zarr_url: str = Field(primary_key=True)
    dataset_id: int = Field(
        primary_key=True,
        foreign_key="datasetv2.id",
        ondelete="CASCADE",
        index=True,
    )
    workflowtask_id: int = Field(
        primary_key=True,
        foreign_key="workflowtaskv2.id",
        ondelete="CASCADE",
        index=True,
    )

    latest_history_unit_id: int = Field(
        foreign_key="historyunit.id",
        ondelete="CASCADE",
        index=True,
    )

class HistoryRun(SQLModel, table=True):
    """
    HistoryRun table.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: int | None = Field(default=None, primary_key=True)
    dataset_id: int = Field(
        foreign_key="datasetv2.id",
        ondelete="CASCADE",
    )
    workflowtask_id: int | None = Field(
        foreign_key="workflowtaskv2.id",
        default=None,
        ondelete="SET NULL",
    )
    job_id: int = Field(foreign_key="jobv2.id")
    task_id: int | None = Field(foreign_key="taskv2.id", ondelete="SET NULL")

    workflowtask_dump: dict[str, Any] = Field(
        sa_column=Column(JSONB, nullable=False),
    )
    task_group_dump: dict[str, Any] = Field(
        sa_column=Column(JSONB, nullable=False),
    )

    timestamp_started: datetime = Field(
        sa_column=Column(DateTime(timezone=True), nullable=False),
        default_factory=get_timestamp,
    )
    status: str
    num_available_images: int

class HistoryUnit(SQLModel, table=True):
    """
    HistoryUnit table.
    """

    id: int | None = Field(default=None, primary_key=True)
    history_run_id: int = Field(
        foreign_key="historyrun.id",
        ondelete="CASCADE",
        index=True,
    )

    logfile: str
    has_warnings: bool = Field(
        sa_column=Column(
            BOOLEAN,
            server_default="false",
            nullable=False,
        ),
    )
    status: str
    zarr_urls: list[str] = Field(
        sa_column=Column(ARRAY(String)),
        default_factory=list,
    )

class JobV2(SQLModel, table=True):
    """
    Job table.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    id: int | None = Field(default=None, primary_key=True)
    project_id: int | None = Field(
        foreign_key="projectv2.id", default=None, ondelete="SET NULL"
    )
    workflow_id: int | None = Field(
        foreign_key="workflowv2.id", default=None, ondelete="SET NULL"
    )
    dataset_id: int | None = Field(
        foreign_key="datasetv2.id", default=None, ondelete="SET NULL"
    )

    user_email: str = Field(nullable=False)
    slurm_account: str | None = None

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

    worker_init: str | None = None
    working_dir: str | None = None
    working_dir_user: str | None = None
    first_task_index: int
    last_task_index: int

    start_timestamp: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )
    end_timestamp: datetime | None = Field(
        default=None, sa_column=Column(DateTime(timezone=True))
    )
    status: str = JobStatusType.SUBMITTED
    log: str | None = None
    executor_error_log: str | None = None

    attribute_filters: dict[str, list[int | float | str | bool]] = Field(
        sa_column=Column(JSONB, nullable=False, server_default="{}")
    )
    type_filters: dict[str, bool] = Field(
        sa_column=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(SQLModel, table=True):
    """
    ORM model for OAuth accounts (`oauthaccount` database table).

    This class is based on fastapi_users_db_sqlmodel::SQLModelBaseOAuthAccount.
    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: int | None = Field(default=None, primary_key=True)
    user_id: int = Field(foreign_key="user_oauth.id", nullable=False)
    user: Optional["UserOAuth"] = Relationship(back_populates="oauth_accounts")
    oauth_name: str = Field(index=True, nullable=False)
    access_token: str = Field(nullable=False)
    expires_at: int | None = Field(nullable=True, default=None)
    refresh_token: str | None = Field(nullable=True, default=None)
    account_id: str = Field(index=True, nullable=False)
    account_email: str = Field(nullable=False)
    model_config = ConfigDict(from_attributes=True)

class Profile(SQLModel, table=True):
    """
    Profile table.
    """

    id: int | None = Field(default=None, primary_key=True)

    resource_id: int = Field(foreign_key="resource.id", ondelete="RESTRICT")

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

    name: str = Field(unique=True)
    """
    Profile name.
    """

    username: str | None = None
    """
    Username to be impersonated, either via `sudo -u` or via `ssh`.
    """

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

    jobs_remote_dir: str | None = None
    """
    Remote path of the job folder (only relevant if
    `resource_type="slurm_ssh"`).
    """

    tasks_remote_dir: str | None = None
    """
    Remote path of the task folder (only relevant if
    `resource_type="slurm_ssh"`).
    """

    pixi_cache_dir: str | None = 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(SQLModel, table=True):
    """
    Project table.
    """

    id: int | None = Field(default=None, primary_key=True)
    name: str

    is_starred: bool = Field(
        sa_column=Column(
            BOOLEAN,
            server_default="false",
            nullable=False,
        ),
    )
    description: str | None = None

    resource_id: int = Field(foreign_key="resource.id", ondelete="RESTRICT")
    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )

class Resource(SQLModel, table=True):
    """
    Resource table.
    """

    id: int | None = Field(default=None, primary_key=True)

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

    name: str = Field(unique=True)
    """
    Resource name.
    """

    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )
    """
    Creation timestamp (autogenerated).
    """

    host: str | None = None
    """
    Address for ssh connections, when `type="slurm_ssh"`.
    """

    prevent_new_submissions: bool = Field(
        sa_column=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: str
    """
    Base local folder for job subfolders (containing artifacts and logs).
    """

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

    jobs_slurm_python_worker: str | None = None
    """
    On SLURM deloyments, this is the Python interpreter that runs the
    `fractal-server` worker from within the SLURM jobs.
    """

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

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

    tasks_python_config: dict[str, Any] = Field(
        sa_column=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: dict[str, Any] = Field(
        sa_column=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(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    task_list: list[TaskV2] = Relationship(
        sa_relationship_kwargs=dict(
            lazy="selectin", cascade="all, delete-orphan"
        ),
    )

    user_id: int = Field(foreign_key="user_oauth.id")
    user_group_id: int | None = Field(
        foreign_key="usergroup.id", default=None, ondelete="SET NULL"
    )
    resource_id: int = Field(foreign_key="resource.id", ondelete="RESTRICT")

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

    active: bool = True
    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )
    timestamp_last_used: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(
            DateTime(timezone=True),
            nullable=False,
            server_default=(
                datetime(2024, 11, 20, tzinfo=timezone.utc).isoformat()
            ),
        ),
    )

    __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(SQLModel, table=True):
    """
    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:
    """

    id: int | None = Field(default=None, primary_key=True)
    name: str

    type: str
    command_non_parallel: str | None = None
    command_parallel: str | None = None

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

    version: str
    args_schema_non_parallel: dict[str, Any] | None = Field(
        sa_column=Column(JSON), default=None
    )
    args_schema_parallel: dict[str, Any] | None = Field(
        sa_column=Column(JSON), default=None
    )
    args_schema_version: str | None = None
    docs_info: str | None = None
    docs_link: str | None = None

    input_types: dict[str, bool] = Field(sa_column=Column(JSONB), default={})
    output_types: dict[str, bool] = Field(sa_column=Column(JSONB), default={})

    taskgroupv2_id: int = Field(foreign_key="taskgroupv2.id")

    category: str | None = None
    modality: str | None = None
    authors: str | None = None
    tags: list[str] = Field(
        sa_column=Column(JSONB, server_default="[]", nullable=False)
    )
    is_core: bool = Field(
        sa_column=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(SQLModel, table=True):
    """
    ORM model for the `usergroup` database table.

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

    id: int | None = Field(default=None, primary_key=True)
    name: str = Field(unique=True)
    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )

class UserOAuth(SQLModel, table=True):
    """
    ORM model for the `user_oauth` database table.

    This class is a modification of
    [`SQLModelBaseUserDB`](https://github.com/fastapi-users/fastapi-users-db-sqlmodel/blob/83980d7f20886120f4636a102ab1822b4c366f63/fastapi_users_db_sqlmodel/__init__.py#L15-L32)
    from `fastapi_users_db_sqlmodel`.
    Original Copyright: 2022 François Voron, released under MIT licence.

    Note that several class attributes are
    [the default ones from `fastapi-users`
    ](https://fastapi-users.github.io/fastapi-users/latest/configuration/schemas/).

    Attributes:
        id:
        email:
        hashed_password:
        is_active:
            If this is `False`, the user has no access to the `/api/v2/`
            endpoints.
        is_superuser:
        is_verified:
            If this is `False`, the user has no access to the `/api/v2/`
            endpoints.
        oauth_accounts:
        profile_id:
            Foreign key linking the user to a `Profile`. If this is unset,
            the user has no access to the `/api/v2/` endpoints.
        project_dirs:
            Absolute paths of the user's project directory. This is used (A) as
            a default base folder for the `zarr_dir` of new datasets (where
            the output Zarr are located), and (B) as a folder which is included
            by default in the paths that a user is allowed to stream (if the
            `fractal-data` integration is set up).
            two goals:
        slurm_accounts:
            List of SLURM accounts that the user can select upon running a job.
    """

    model_config = ConfigDict(from_attributes=True)

    __tablename__ = "user_oauth"

    id: int | None = Field(default=None, primary_key=True)

    email: EmailStr = Field(
        sa_column_kwargs={"unique": True, "index": True},
        nullable=False,
    )
    hashed_password: str
    is_active: bool = Field(default=True, nullable=False)
    is_superuser: bool = Field(default=False, nullable=False)
    is_verified: bool = Field(default=False, nullable=False)
    is_guest: bool = Field(
        sa_column=Column(
            BOOLEAN,
            server_default="false",
            nullable=False,
        ),
    )

    oauth_accounts: list["OAuthAccount"] = Relationship(
        back_populates="user",
        sa_relationship_kwargs={"lazy": "joined", "cascade": "all, delete"},
    )

    profile_id: int | None = Field(
        foreign_key="profile.id",
        default=None,
        ondelete="RESTRICT",
    )

    project_dirs: list[str] = Field(
        sa_column=Column(ARRAY(String), nullable=False),
    )

    slurm_accounts: list[str] = Field(
        sa_column=Column(ARRAY(String), server_default="{}"),
    )

    __table_args__ = (
        CheckConstraint(
            "NOT (is_superuser AND is_guest)",
            name="superuser_is_not_guest",
        ),
    )