# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.
"""Models used by debusine client."""
import json
from collections.abc import Iterator, Sequence
from datetime import datetime
from enum import StrEnum
from pathlib import Path
from typing import Any, Generic, Literal, NewType, TypeVar
try:
import pydantic.v1 as pydantic
except ImportError:
import pydantic as pydantic # type: ignore
from debusine.assets import (
AssetCategory,
BaseAssetDataModel,
asset_class_from_category,
)
from debusine.utils import calculate_hash
[docs]class StrictBaseModel(pydantic.BaseModel):
"""Stricter pydantic configuration."""
[docs] class Config:
"""Set up stricter pydantic Config."""
validate_assignment = True
[docs]class PaginatedResponse(StrictBaseModel):
"""Paginated response from the API."""
count: int | None
next: pydantic.AnyUrl | None
previous: pydantic.AnyUrl | None
results: list[dict[str, Any]]
[docs]class WorkRequestRequest(StrictBaseModel):
"""Client send a WorkRequest to the server."""
task_name: str
workspace: str | None = None
task_data: dict[str, Any]
event_reactions: dict[str, Any]
[docs]class WorkRequestResponse(StrictBaseModel):
"""Server return a WorkRequest to the client."""
id: int
created_at: datetime
started_at: datetime | None = None
completed_at: datetime | None = None
duration: int | None = None
status: str
result: str
worker: int | None = None
task_type: str
task_name: str
task_data: dict[str, Any]
dynamic_task_data: dict[str, Any] | None
priority_base: int
priority_adjustment: int
artifacts: list[int]
workspace: str
def __str__(self) -> str:
"""Return representation of the object."""
return f'WorkRequest: {self.id}'
[docs]class WorkRequestExternalDebsignRequest(StrictBaseModel):
"""Client sends data from an external `debsign` run to the server."""
signed_artifact: int
[docs]class OnWorkRequestCompleted(StrictBaseModel):
"""
Server return an OnWorkRequestCompleted to the client.
Returned via websocket consumer endpoint.
"""
work_request_id: int
completed_at: datetime
result: str
[docs]class WorkflowTemplateRequest(StrictBaseModel):
"""Client sends a WorkflowTemplate to the server."""
name: str
task_name: str
workspace: str | None = None
task_data: dict[str, Any]
priority: int
[docs]class WorkflowTemplateResponse(StrictBaseModel):
"""Server returns a WorkflowTemplate to the server."""
id: int
task_name: str
workspace: str
task_data: dict[str, Any]
priority: int
[docs]class CreateWorkflowRequest(StrictBaseModel):
"""Client sends a workflow creation request to the server."""
template_name: str
workspace: str | None = None
task_data: dict[str, Any]
# With pydantic >= 2.1.0, use Annotated[str, pydantic.Field(max_length=255)]
# instead. Unfortunately, while pydantic v1 seems to accept this, it
# silently ignores the max_length annotation.
[docs]class StrMaxLength255(pydantic.ConstrainedStr):
"""A string with a maximum length of 255 characters."""
max_length = 255
[docs]class FileRequest(StrictBaseModel):
"""Declare a FileRequest: client sends it to the server."""
size: int = pydantic.Field(ge=0)
checksums: dict[str, StrMaxLength255]
type: Literal["file"]
[docs] @staticmethod
def create_from(path: Path) -> "FileRequest":
"""Return a FileRequest for the file path."""
return FileRequest(
size=path.stat().st_size,
checksums={
"sha256": pydantic.parse_obj_as(
StrMaxLength255, calculate_hash(path, "sha256").hex()
)
},
type="file",
)
[docs]class FileResponse(StrictBaseModel):
"""Declare a FileResponse: server sends it to the client."""
size: int = pydantic.Field(ge=0)
checksums: dict[str, StrMaxLength255]
type: Literal["file"]
url: pydantic.AnyUrl
FilesRequestType = NewType("FilesRequestType", dict[str, FileRequest])
FilesResponseType = NewType("FilesResponseType", dict[str, FileResponse])
[docs]class ArtifactCreateRequest(StrictBaseModel):
"""Declare an ArtifactCreateRequest: client sends it to the server."""
category: str
workspace: str | None = None
files: FilesRequestType = FilesRequestType({})
data: dict[str, Any] = {}
work_request: int | None = None
expire_at: datetime | None = None
[docs]class ArtifactResponse(StrictBaseModel):
"""Declare an ArtifactResponse: server sends it to the client."""
id: int
workspace: str
category: str
created_at: datetime
data: dict[str, Any]
download_tar_gz_url: pydantic.AnyUrl
files_to_upload: list[str]
expire_at: datetime | None = None
files: FilesResponseType = FilesResponseType({})
[docs]class RemoteArtifact(StrictBaseModel):
"""Declare RemoteArtifact."""
id: int
workspace: str
[docs]class AssetCreateRequest(StrictBaseModel):
"""Request for the Asset creation API."""
category: AssetCategory
data: BaseAssetDataModel
work_request: int | None = None
workspace: str
[docs]class AssetResponse(StrictBaseModel):
"""Response from an Asset creation / listing API."""
id: int
category: AssetCategory
data: BaseAssetDataModel | dict[str, Any]
work_request: int | None = None
workspace: str
[docs] @pydantic.root_validator
def parse_data(cls, values: dict[str, Any]) -> dict[str, Any]: # noqa: U100
"""Parse data using the correct data model."""
if isinstance(values.get("data"), dict):
category = values.get("category")
if category: # pragma: no cover
asset_cls = asset_class_from_category(category)
values["data"] = asset_cls.parse_obj(values["data"])
return values
[docs]class AssetsResponse(StrictBaseModel):
"""A response from the server with multiple AssetResponse objects."""
# This model parses a list, not an object/dict.
__root__: Sequence[AssetResponse]
def __iter__(self) -> Iterator[AssetResponse]: # type: ignore[override]
"""Iterate over individual asset responses."""
return iter(self.__root__)
[docs]class RelationType(StrEnum):
"""Possible values for `RelationCreateRequest.type`."""
EXTENDS = "extends"
RELATES_TO = "relates-to"
BUILT_USING = "built-using"
[docs]class RelationCreateRequest(StrictBaseModel):
"""Declare a RelationCreateRequest: client sends it to the server."""
artifact: int
target: int
type: RelationType
[docs]class RelationResponse(RelationCreateRequest):
"""Declare a RelationResponse."""
id: int
[docs]class RelationsResponse(StrictBaseModel):
"""A response from the server with multiple RelationResponse objects."""
# This model parses a list, not an object/dict.
__root__: Sequence[RelationResponse]
def __iter__(self) -> Iterator[RelationResponse]: # type: ignore[override]
"""Iterate over individual relation responses."""
return iter(self.__root__)
[docs]class LookupChildType(StrEnum):
"""Possible values for `LookupDict.child_type` and `expect_type`."""
BARE = "bare"
ARTIFACT = "artifact"
ARTIFACT_OR_PROMISE = "artifact-or-promise"
COLLECTION = "collection"
ANY = "any"
[docs]class LookupResultType(StrEnum):
"""A collection item type returned by a lookup."""
BARE = "b"
ARTIFACT = "a"
COLLECTION = "c"
[docs]class LookupSingleRequest(StrictBaseModel):
"""A request from the client to look up a single collection item."""
lookup: int | str
work_request: int
expect_type: LookupChildType
default_category: str | None = None
[docs]class LookupMultipleRequest(StrictBaseModel):
"""A request from the client to look up multiple collection items."""
lookup: list[int | str | dict[str, Any]]
work_request: int
expect_type: LookupChildType
default_category: str | None = None
[docs]class LookupSingleResponse(StrictBaseModel):
"""A response from the server with a single lookup result."""
result_type: LookupResultType
collection_item: int | None = None
artifact: int | None = None
collection: int | None = None
[docs]class LookupSingleResponseArtifact(LookupSingleResponse):
"""
A response from the server with a single lookup result for an artifact.
Used to assist type annotations.
"""
result_type: Literal[LookupResultType.ARTIFACT]
artifact: int
[docs]class LookupSingleResponseCollection(LookupSingleResponse):
"""
A response from the server with a single lookup result for a collection.
Used to assist type annotations.
"""
result_type: Literal[LookupResultType.COLLECTION]
collection: int
LSR = TypeVar("LSR", bound=LookupSingleResponse, covariant=True)
[docs]class LookupMultipleResponse(StrictBaseModel, Generic[LSR]):
"""A response from the server with multiple lookup results."""
# This model parses a list, not an object/dict.
__root__: Sequence[LSR]
def __iter__(self) -> Iterator[LSR]: # type: ignore[override]
"""Iterate over individual results."""
return iter(self.__root__)
[docs]def model_to_json_serializable_dict(
model: pydantic.BaseModel, exclude_unset: bool = False
) -> dict[Any, Any]:
"""
Similar to model.dict() but the returned dictionary is JSON serializable.
For example, a datetime() is not JSON serializable. Using this method will
return a dictionary with a string instead of a datetime object.
Replace with model_dump() in Pydantic 2.
"""
serializable = json.loads(model.json(exclude_unset=exclude_unset))
assert isinstance(serializable, dict)
return serializable
[docs]def model_to_json_serializable_list(
model: pydantic.BaseModel, exclude_unset: bool = False
) -> list[dict[Any, Any]]:
"""Similar to model_to_json_serializable_dict, but for a list response."""
serializable = json.loads(model.json(exclude_unset=exclude_unset))
assert isinstance(serializable, list)
return serializable