# 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 datetime import datetime
from pathlib import Path
from typing import Any, Literal, NewType, Optional
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: Optional[int]
next: Optional[pydantic.AnyUrl]
previous: Optional[pydantic.AnyUrl]
results: list[dict[str, Any]]
[docs]
class WorkRequestRequest(StrictBaseModel):
"""Client send a WorkRequest to the server."""
task_name: str
workspace: Optional[str] = None
task_data: dict[str, Any]
[docs]
class WorkRequestResponse(StrictBaseModel):
"""Server return a WorkRequest to the client."""
id: int
created_at: datetime
started_at: Optional[datetime] = None
completed_at: Optional[datetime] = None
duration: Optional[int] = None
status: str
result: str
worker: Optional[int] = None
task_name: str
task_data: dict[str, Any]
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
# Backport of pydantic.NonNegativeInt from 1.8.
[docs]
class NonNegativeInt(pydantic.ConstrainedInt):
"""A non-negative integer."""
ge = 0
# With pydantic >= 2.1.0, use Annotated[str,
# pydantic.StringConstraints(max_length=255)] instead. Unfortunately there
# doesn't seem to be a less ugly way to write this with earlier versions of
# pydantic that mypy is happy with.
[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: NonNegativeInt
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=pydantic.parse_obj_as(NonNegativeInt, 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: NonNegativeInt
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: Optional[str] = None
files: FilesRequestType = FilesRequestType({})
data: dict[str, Any] = {}
work_request: Optional[int] = None
expire_at: Optional[datetime] = 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: Optional[datetime] = 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]
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())