"""Pydantic models for the SimDB remote API."""
import base64
from datetime import datetime as dt
from datetime import timezone
from pathlib import Path
from typing import (
Annotated,
Any,
Dict,
Generic,
List,
Literal,
Optional,
TypeVar,
Union,
)
from urllib.parse import urlencode
from uuid import UUID, uuid1
import numpy as np
from pydantic import (
BaseModel as _BaseModel,
)
from pydantic import (
BeforeValidator,
ConfigDict,
Field,
PlainSerializer,
model_validator,
)
from pydantic import (
RootModel as _RootModel,
)
from simdb.cli.manifest import DataObject
HexUUID = Annotated[UUID, PlainSerializer(lambda x: x.hex, return_type=str)]
"""UUID serialized as a hex string."""
[docs]
class BaseModel(_BaseModel):
model_config = ConfigDict(use_attribute_docstrings=True)
[docs]
class RootModel(_RootModel):
model_config = ConfigDict(use_attribute_docstrings=True)
def _deserialize_custom_uuid(v: Any) -> UUID:
"""Deserialize CustomUUID format back to UUID."""
if isinstance(v, UUID):
return v
if isinstance(v, dict) and "hex" in v:
return UUID(hex=v["hex"])
raise ValueError(f"Cannot deserialize {v} to UUID")
CustomUUID = Annotated[
UUID,
BeforeValidator(_deserialize_custom_uuid),
PlainSerializer(lambda x: {"_type": "uuid.UUID", "hex": x.hex}),
]
"""UUID with custom serialization format."""
StatusLiteral = Literal[
"not validated", "accepted", "failed", "passed", "deprecated", "deleted"
]
"""String representation of a simulation status"""
def _array_to_range(value: Any) -> Any:
"""Convert a numpy array or list to a RangeValue if it contains numeric data."""
if value is None:
return None
if (
isinstance(value, dict)
and "dtype" in value
and "shape" in value
and "bytes" in value
):
np_bytes = base64.decodebytes(value["bytes"].encode())
return _array_to_range(np.frombuffer(np_bytes, dtype=value["dtype"]))
if isinstance(value, np.ndarray):
value = _array_to_range(value.tolist())
if isinstance(value, (list, tuple)):
if len(value) == 0:
return value
try:
float_values = [float(v) for v in value]
return RangeValue(min=min(float_values), max=max(float_values))
except (TypeError, ValueError):
return value
return value
[docs]
class RangeValue(BaseModel):
"""A numeric range with min and max bounds."""
model_config = ConfigDict(extra="forbid")
min: float
max: float
MetadataValue = Union[
CustomUUID,
str,
int,
float,
bool,
list,
RangeValue,
dict[str, Any],
None,
]
"""Supported types for simulation metadata values. Numpy arrays and regular arrays
containing numeric data are automatically converted to RangeValue."""
[docs]
class StatusPatchData(BaseModel):
"""Post data for updating simulation status."""
status: StatusLiteral
"""New simulation status."""
[docs]
class DeletedSimulation(BaseModel):
"""Reference to a deleted simulation."""
simulation: UUID
"""UUID of the deleted simulation."""
files: List[str]
"""List of deleted file paths."""
[docs]
class SimulationDeleteResponse(BaseModel):
"""Response from DELETE v1.2/simulations/{uuid}."""
deleted: DeletedSimulation
"""Reference to the deleted simulation."""
[docs]
class FileData(BaseModel):
"""Model representing a file in the system."""
type: Literal["UNKNOWN", "UUID", "FILE", "IMAS"]
"""File type."""
uri: str
"""URI to the file location."""
uuid: CustomUUID = Field(default_factory=lambda: uuid1())
"""Unique identifier for the file."""
checksum: str
"""Checksum of the file."""
datetime: dt
"""Timestamp of the file."""
usage: Optional[str] = None
"""File usage description."""
purpose: Optional[str] = None
"""Purpose of the file."""
sensitivity: Optional[str] = None
"""Sensitivity level of the file."""
access: Optional[str] = None
"""Access permissions."""
embargo: Optional[str] = None
"""Embargo information."""
[docs]
class FileDataList(RootModel):
"""List of FileData items."""
root: List[FileData] = []
def __getitem__(self, item) -> FileData:
"""Allow indexing on the list."""
return self.root[item]
[docs]
class SimulationReference(BaseModel):
"""Reference to a simulation."""
uuid: CustomUUID
"""UUID of the simulation."""
alias: Optional[str] = None
"""Alias of the simulation."""
[docs]
class SimulationData(BaseModel):
"""Core simulation data."""
uuid: CustomUUID = Field(default_factory=lambda: uuid1())
"""Unique identifier of the simulation."""
alias: Optional[str] = None
"""Human-readable alias."""
datetime: dt = Field(default_factory=lambda: dt.now(timezone.utc))
"""Creation timestamp."""
inputs: FileDataList = FileDataList()
"""List of input files."""
outputs: FileDataList = FileDataList()
"""List of output files."""
metadata: MetadataDataList = MetadataDataList()
"""Simulation metadata."""
[docs]
class SimulationDataResponse(SimulationData):
"""Simulation data with parent/child references."""
parents: List[SimulationReference]
"""Parent simulations."""
children: List[SimulationReference]
"""Child simulations."""
[docs]
class SimulationPatchResponse(BaseModel):
pass
[docs]
class SimulationPostData(BaseModel):
"""Data for creating a new simulation."""
simulation: SimulationData
"""The simulation data to create."""
add_watcher: bool
"""Whether to add a watcher for this simulation."""
uploaded_by: Optional[str] = None
"""User who uploaded the simulation."""
[docs]
class ValidationResult(BaseModel):
"""Result of simulation validation."""
passed: bool
"""Whether validation passed."""
error: Optional[str] = None
"""Error message if validation failed."""
[docs]
class SimulationPostResponse(BaseModel):
"""Response from creating a simulation."""
ingested: HexUUID
"""UUID of the ingested simulation."""
error: Optional[str] = None
"""Error message if ingestion failed."""
validation: Optional[ValidationResult] = None
"""Validation result."""
[docs]
class SimulationListItem(BaseModel):
"""Summary of a simulation for list views."""
uuid: CustomUUID
"""UUID of the simulation."""
alias: Optional[str] = None
"""Alias of the simulation."""
datetime: str
"""Creation timestamp."""
metadata: MetadataDataList = MetadataDataList()
"""Simulation metadata."""
T = TypeVar("T")
"""Type variable for generic paginated responses."""
[docs]
class PaginatedResponse(BaseModel, Generic[T]):
"""Generic paginated response wrapper."""
count: int
"""Total number of items."""
page: int
"""Current page number."""
limit: int
"""Number of items per page."""
results: List[T]
"""List of results for this page."""
[docs]
class SimulationTraceData(SimulationData):
"""Simulation data with status history."""
status: Optional[StatusLiteral] = None
"""Current status of the simulation."""
passed_on: Optional[Any] = None
"""Timestamp when status changed to passed."""
failed_on: Optional[Any] = None
"""Timestamp when status changed to failed."""
deprecated_on: Optional[Any] = None
"""Timestamp when status changed to deprecated."""
accepted_on: Optional[Any] = None
"""Timestamp when status changed to accepted."""
not_validated_on: Optional[Any] = None
"""Timestamp when status changed to not validated."""
deleted_on: Optional[Any] = None
"""Timestamp when status changed to deleted."""
replaces: Optional["SimulationTraceData"] = None
"""Simulation this one replaces."""
replaces_reason: Optional[Any] = None
"""Reason for replacement."""
[docs]
class ChunkInfo(BaseModel):
"""Information about a single chunk in a chunked file upload."""
chunk_size: int
"""Length of the chunk."""
chunk: int
"""Index of the chunk."""
num_chunks: Optional[int] = 1
"""Total amount of chunks in the file."""
[docs]
class ChunkInfoDict(RootModel):
"""Dictionary mapping file UUID hex to chunk info."""
root: Dict[str, ChunkInfo]
[docs]
class FileUploadData(BaseModel):
"""Data payload for file chunk upload (sent as JSON in 'data' field)."""
simulation: SimulationData
"""The simulation the file belongs to."""
file_type: str
"""Type of the file."""
chunk_info: Optional[Dict[str, ChunkInfo]] = None
"""Info about the chunk."""
[docs]
class FilesGetResponse(RootModel):
"""Response from the get files endpoint."""
root: List[FileData]
"""List of files."""
[docs]
class FileInfo(BaseModel):
"""Information about a single file on disk."""
path: Path
"""Path to the file."""
checksum: str
"""Checksum of the file."""
[docs]
class FileGetDataResponse(FileData):
"""Response from the get file data endpoint, extending FileData with disk info."""
files: List[FileInfo]
"""List of file info entries for the files on disk."""
[docs]
class FileUploadResponse(BaseModel):
"""Response from file upload/chunk upload endpoint."""
pass
[docs]
class FileRegistrationItem(BaseModel):
"""A single file entry in the file registration payload."""
chunks: int
"""The amount of chunks to be processed."""
file_type: str
"""The file type."""
file_uuid: HexUUID
"""The UUID of the file."""
ids_list: Optional[List[Any]] = None
"""List of IDS names associated with the file."""
[docs]
class FileRegistrationData(BaseModel):
"""Payload for final file registration after chunk uploads."""
simulation: SimulationData
"""The simulation the files belong to."""
obj_type: DataObject.Type
"""The type of the data object being registered."""
files: List[FileRegistrationItem]
"""List of file registration items."""
[docs]
class FileRegistrationResponse(BaseModel):
"""Response from file registration endpoint."""
pass
[docs]
class WatcherReference(BaseModel):
"""An watcher entry reference."""
simulation: HexUUID
"""Simulation UUID the watcher has been added to."""
watcher: str
"""Username of the added watcher."""
[docs]
class WatcherPostResponse(BaseModel):
"""Response from the add watcher endpoint."""
added: WatcherReference
"""The added watcher data."""
[docs]
class WatcherPostRequest(BaseModel):
"""Payload for adding a watcher to a simulation."""
user: Optional[str]
"""Username of the watcher, defaults to the signed in user."""
email: Optional[str]
"""Email of the watcher, defaults to the signed in user."""
notification: Literal["VALIDATION", "REVISION", "OBSOLESCENCE", "ALL"]
"""Notificaiton type of the watcher."""
[docs]
class WatcherData(BaseModel):
"""Payload describing a watcher."""
username: str
"""Username of the watcher."""
email: str
"""Email address of the watcher."""
notification: Literal["V", "R", "O", "A"]
"""Notification type of the watcher.
Types are: V(alidation), R(evision), O(bsolescence) and A(ll)
"""
[docs]
class WatcherGetResponse(RootModel):
"""Response from the get watchers endpoint."""
root: List[WatcherData]
[docs]
class WatcherDeleteRequest(BaseModel):
"""Payload for deleting a watcher from a simulation."""
user: str
"""Username to delete from the watchers."""
[docs]
class WatcherDeleteResponse(BaseModel):
"""Response from the delete watchers endpoint."""
removed: WatcherReference
"""Reference to the deleted wacher."""
[docs]
class StagingDirectoryResponse(BaseModel):
"""Response from the get staging dir endpoint."""
staging_dir: Path
"""Path to the staging dir."""
[docs]
class ErrorResponse(BaseModel):
"""Response model for server errors."""
error: str
"""Error description."""