# Copyright 2022 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 Sequence
from datetime import datetime
from enum import StrEnum
from pathlib import Path
from typing import Annotated, Any, Generic, Literal, NewType, TypeVar
try:
import pydantic.v1 as pydantic
except ImportError:
import pydantic as pydantic # type: ignore
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_name: str
task_data: dict[str, Any]
priority_base: int
priority_adjustment: int
artifacts: list[int]
workspace: str
def __str__(self):
"""Return representation of the object."""
return f'WorkRequest: {self.id}'
[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 FileRequest(StrictBaseModel):
"""Declare a FileRequest: client sends it to the server."""
size: int = pydantic.Field(ge=0)
checksums: dict[str, Annotated[str, pydantic.Field(max_length=255)]]
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": 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, Annotated[str, pydantic.Field(max_length=255)]]
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
RelationCreateRequestType = Literal["extends", "relates-to", "built-using"]
[docs]class RelationCreateRequest(StrictBaseModel):
"""Declare a RelationCreateRequest: client sends it to the server."""
artifact: int
target: int
type: RelationCreateRequestType
[docs]class RelationResponse(RelationCreateRequest):
"""Declare a RelationResponse."""
id: int
[docs]class CollectionItemType(StrEnum):
"""A collection item type."""
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: CollectionItemType
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: CollectionItemType
default_category: str | None = None
[docs]class LookupSingleResponse(StrictBaseModel):
"""A response from the server with a single lookup result."""
result_type: CollectionItemType
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[CollectionItemType.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[CollectionItemType.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):
"""Iterate over individual results."""
return iter(self.__root__)
[docs]def model_to_json_serializable_dict(
model: pydantic.BaseModel,
) -> 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.
"""
return json.loads(model.json())