govoplan-datasources/src/govoplan_datasources/backend/schemas.py

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator


DatasourceModeValue = Literal["live", "cached", "static"]
DatasourceKindValue = Literal[
    "upload",
    "database",
    "http",
    "rest",
    "directory",
    "file",
    "feed",
    "custom",
]
DatasourceShapeValue = Literal["tabular", "document", "binary", "directory", "stream"]


class DatasourceFieldResponse(BaseModel):
    name: str
    data_type: str
    nullable: bool


class DatasourceResponse(BaseModel):
    ref: str
    source_name: str
    name: str
    description: str | None
    kind: DatasourceKindValue
    mode: DatasourceModeValue
    shape: DatasourceShapeValue
    status: str
    provider: str | None
    provider_ref: str | None
    fields: list[DatasourceFieldResponse] = Field(alias="schema")
    schema_version: str
    fingerprint: str
    current_materialization_ref: str | None
    row_count: int | None
    byte_count: int | None
    updated_at: str | None
    capabilities: list[str]
    provenance: dict[str, Any]
    metadata: dict[str, Any]


class DatasourceListResponse(BaseModel):
    datasources: list[DatasourceResponse]


class DatasourceMaterializationResponse(BaseModel):
    ref: str
    datasource_ref: str
    revision: int
    state: str
    fingerprint: str
    fields: list[DatasourceFieldResponse] = Field(alias="schema")
    row_count: int | None
    byte_count: int | None
    frozen_at: str | None
    frozen_label: str | None
    source_timestamp: str | None
    created_at: str | None
    provenance: dict[str, Any]
    metadata: dict[str, Any]


class DatasourceMaterializationListResponse(BaseModel):
    materializations: list[DatasourceMaterializationResponse]


class DatasourceStageResponse(BaseModel):
    ref: str
    name: str
    source_name: str
    kind: DatasourceKindValue
    mode: DatasourceModeValue
    shape: DatasourceShapeValue
    state: str
    target_datasource_ref: str | None
    fingerprint: str
    fields: list[DatasourceFieldResponse] = Field(alias="schema")
    row_count: int | None
    byte_count: int | None
    validation: dict[str, Any]
    created_at: str | None
    promoted_at: str | None
    promoted_materialization_ref: str | None
    provenance: dict[str, Any]
    metadata: dict[str, Any]


class DatasourceStageListResponse(BaseModel):
    stages: list[DatasourceStageResponse]


class DatasourceStageCreateRequest(BaseModel):
    name: str = Field(min_length=1, max_length=300)
    source_name: str = Field(min_length=1, max_length=120)
    description: str | None = Field(default=None, max_length=4_000)
    kind: DatasourceKindValue = "upload"
    mode: Literal["static", "cached"] = "static"
    shape: Literal["tabular"] = "tabular"
    format: Literal["json", "csv"] = "json"
    rows: list[dict[str, Any]] | None = Field(default=None, max_length=10_000)
    csv_text: str | None = Field(default=None, max_length=5_000_000)
    delimiter: Literal[",", ";", "\t", "|"] = ","
    target_datasource_ref: str | None = None
    provenance: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def matching_payload(self) -> "DatasourceStageCreateRequest":
        if self.format == "json" and self.rows is None:
            raise ValueError("JSON stages require rows.")
        if self.format == "csv" and not self.csv_text:
            raise ValueError("CSV stages require csv_text.")
        return self


class DatasourceStagePromoteRequest(BaseModel):
    freeze: bool = False
    frozen_label: str | None = Field(default=None, max_length=300)


class DatasourceStagePromoteResponse(BaseModel):
    datasource: DatasourceResponse
    materialization: DatasourceMaterializationResponse


class DatasourceOriginResponse(BaseModel):
    ref: str
    source_name: str
    name: str
    description: str | None
    kind: DatasourceKindValue
    shape: DatasourceShapeValue
    supported_modes: list[DatasourceModeValue]
    provider: str
    fields: list[DatasourceFieldResponse] = Field(alias="schema")
    schema_version: str
    fingerprint: str
    row_count: int | None
    byte_count: int | None
    updated_at: str | None
    capabilities: list[str]
    metadata: dict[str, Any]


class DatasourceOriginListResponse(BaseModel):
    available: bool
    origins: list[DatasourceOriginResponse]


class DatasourceOriginRegisterRequest(BaseModel):
    origin_ref: str = Field(min_length=1, max_length=500)
    name: str = Field(min_length=1, max_length=300)
    source_name: str = Field(min_length=1, max_length=120)
    mode: Literal["live", "cached"]
    description: str | None = Field(default=None, max_length=4_000)


class DatasourcePreviewResponse(BaseModel):
    datasource: DatasourceResponse
    rows: list[dict[str, Any]]
    total_rows: int
    truncated: bool
    materialization: DatasourceMaterializationResponse | None


class DatasourceFreezeRequest(BaseModel):
    label: str | None = Field(default=None, max_length=300)


class DatasourceRetireResponse(BaseModel):
    retired: bool
    datasource_ref: str