strepsiades
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Annotated, Any, Dict, List, Literal
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator
|
|
|
|
from ..datasource.type_registry import register_datasource_config_model, validate_registered_datasource_config
|
|
|
|
|
|
class BaseDataSourceConfig(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
|
|
|
type: str
|
|
target_project_ids: List[str] = Field(default_factory=list)
|
|
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
|
|
|
|
|
class JsonlDataSourceConfig(BaseDataSourceConfig):
|
|
|
|
type: Literal["jsonl"] = "jsonl"
|
|
jsonl_dir: str
|
|
read_only: bool = False
|
|
|
|
|
|
class DbDataSourceConfig(BaseDataSourceConfig):
|
|
|
|
type: Literal["db"] = "db"
|
|
|
|
|
|
class ApiDataSourceConfig(BaseDataSourceConfig):
|
|
|
|
type: Literal["api"] = "api"
|
|
api_base_url: str = ""
|
|
api_uid: str = ""
|
|
api_secret: str = ""
|
|
api_debug: bool = False
|
|
poll_max_retries: int = 10
|
|
poll_interval: float = 0.5
|
|
api_rate_limit_max_requests: int = 10
|
|
api_rate_limit_window_seconds: float = 10.0
|
|
target_project_ids: List[str] = Field(default_factory=list)
|
|
|
|
|
|
DataSourceConfig = Annotated[
|
|
BaseDataSourceConfig,
|
|
PlainValidator(validate_registered_datasource_config),
|
|
PlainSerializer(
|
|
lambda value: value.model_dump(mode="python") if isinstance(value, BaseModel) else value,
|
|
return_type=dict[str, Any],
|
|
),
|
|
]
|
|
|
|
|
|
register_datasource_config_model("jsonl", JsonlDataSourceConfig)
|
|
register_datasource_config_model("db", DbDataSourceConfig)
|
|
register_datasource_config_model("api", ApiDataSourceConfig)
|