backend/src/models/order.py - master - FlysAt/DataFoundry

backend/src/models/order.py
151 строка4 KB
AlexanderShmygol
add annotation_executor
03 май 2026, 17:02
03 май 2026, 17:029eaf48f
import enum
from datetime import datetime
from typing import Any, Callable

from sqlalchemy import DateTime, ForeignKey, Integer, Numeric, String, Text
from sqlalchemy import Enum as SqlEnum
from sqlalchemy.dialects.postgresql import ARRAY
from sqlalchemy.orm import Mapped, mapped_column, relationship
from sqlalchemy.sql import func

from src.models.base import Base


def enum_values(enum_cls: type[enum.Enum]) -> list[Callable[[], Any]]:
    return [item.value for item in enum_cls]


class OperationType(str, enum.Enum):
    COLLECTION = "collection"
    LABELING = "labeling"


class OrderStatus(str, enum.Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    REVIEW = "review"
    COMPLETED = "completed"
    CANCELLED = "cancelled"


class VisibilityType(str, enum.Enum):
    PUBLIC = "public"
    REQUEST = "request"
    PRIVATE = "private"


class Order(Base):
    __tablename__ = "order"

    id: Mapped[int] = mapped_column(primary_key=True, index=True)
    title: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str | None] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        nullable=False,
    )
    deadline: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    data_volume: Mapped[str | None] = mapped_column(String(255), nullable=True)
    operation_type: Mapped[OperationType] = mapped_column(
        SqlEnum(
            OperationType,
            name="operation_type_enum",
            values_callable=enum_values,
        ),
        nullable=False,
    )
    budget: Mapped[float | None] = mapped_column(Numeric(12, 2), nullable=True)
    operation_tags: Mapped[list[str]] = mapped_column(
        ARRAY(String),
        nullable=False,
        default=list,
        server_default="{}",
    )
    description_tags: Mapped[list[str]] = mapped_column(
        ARRAY(String),
        nullable=False,
        default=list,
        server_default="{}",
    )
    status: Mapped[OrderStatus] = mapped_column(
        SqlEnum(
            OrderStatus,
            name="order_status_enum",
            values_callable=enum_values,
        ),
        nullable=False,
        default=OrderStatus.PENDING,
        server_default=OrderStatus.PENDING.value,
    )
    quality_requirement: Mapped[int | None] = mapped_column(Integer, nullable=True)
    dataset_link: Mapped[int | None] = mapped_column(
        ForeignKey("order_datasets.id", ondelete="SET NULL"),
        nullable=True,
    )
    license_name: Mapped[str | None] = mapped_column(String(255), nullable=True)
    visibility: Mapped[VisibilityType] = mapped_column(
        SqlEnum(
            VisibilityType,
            name="visibility_type_enum",
            values_callable=enum_values,
        ),
        nullable=False,
        default=VisibilityType.PUBLIC,
        server_default=VisibilityType.PUBLIC.value,
    )
    customer_id: Mapped[int] = mapped_column(
        ForeignKey("user.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    current_executor_id: Mapped[int | None] = mapped_column(
        ForeignKey("user.id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )

    customer = relationship(
        "User",
        foreign_keys=[customer_id],
        back_populates="created_orders",
    )
    current_executor = relationship(
        "User",
        foreign_keys=[current_executor_id],
        back_populates="executing_orders",
    )
    assignments = relationship(
        "OrderAssignment",
        back_populates="order",
        cascade="all, delete-orphan",
    )
    messages = relationship(
        "Message",
        back_populates="order",
        cascade="all, delete-orphan",
    )
    datasets = relationship(
        "OrderDataset",
        foreign_keys="OrderDataset.order_id",
        back_populates="order",
        cascade="all, delete-orphan",
    )
    label_studio_projects = relationship(
        "LabelStudioProject",
        back_populates="order",
        cascade="all, delete-orphan",
    )
    annotation_types = relationship(
        "OrderAnnotationType",
        back_populates="order",
        cascade="all, delete-orphan",
    )
    executor_annotation_types = relationship(
        "OrderAnnotationTypeExecutor",
        back_populates="order",
        cascade="all, delete-orphan",
    )