增强配置系统
This commit is contained in:
@@ -14,13 +14,20 @@ from . import domain as _domain
|
||||
from .common.registry import DomainRegistry
|
||||
from .config import (
|
||||
ApiDataSourceConfig,
|
||||
BaseDataSourceConfig,
|
||||
DbDataSourceConfig,
|
||||
JsonlDataSourceConfig,
|
||||
PipelineRunConfig,
|
||||
build_config,
|
||||
build_default_config,
|
||||
get_datasource_config_model,
|
||||
get_datasource_factory,
|
||||
list_registered_datasource_types,
|
||||
load_named_preset_overrides,
|
||||
load_overrides_from_file,
|
||||
register_datasource_config_model,
|
||||
register_datasource_factory,
|
||||
register_datasource_type,
|
||||
)
|
||||
from .datasource import ApiClient, ApiDataSource, BaseApiHandler, BaseDbHandler, BaseJsonlHandler, DbDataSource, JsonlDataSource
|
||||
from .pipeline import FullSyncPipeline, create_pipeline_from_config, run_pipeline_from_config
|
||||
@@ -42,12 +49,19 @@ __all__ = [
|
||||
"DomainRegistry",
|
||||
"PipelineRunConfig",
|
||||
"ApiDataSourceConfig",
|
||||
"BaseDataSourceConfig",
|
||||
"DbDataSourceConfig",
|
||||
"JsonlDataSourceConfig",
|
||||
"build_config",
|
||||
"build_default_config",
|
||||
"get_datasource_config_model",
|
||||
"get_datasource_factory",
|
||||
"list_registered_datasource_types",
|
||||
"load_overrides_from_file",
|
||||
"load_named_preset_overrides",
|
||||
"register_datasource_config_model",
|
||||
"register_datasource_factory",
|
||||
"register_datasource_type",
|
||||
"ApiClient",
|
||||
"ApiDataSource",
|
||||
"DbDataSource",
|
||||
|
||||
@@ -26,7 +26,7 @@ Domain 注册中心
|
||||
# DataSource 查询 Handler
|
||||
handler_class = DomainRegistry.get_jsonl_handler("contract")
|
||||
"""
|
||||
from typing import Dict, Type, Optional, Any, Literal
|
||||
from typing import Dict, Type, Optional, Any
|
||||
from dataclasses import dataclass, field
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -48,6 +48,7 @@ class DomainRegistration:
|
||||
jsonl_handler_class: Optional[Type[Any]] = None
|
||||
api_handler_class: Optional[Type[Any]] = None
|
||||
db_handler_class: Optional[Type[Any]] = None
|
||||
handler_classes: Dict[str, Type[Any]] = field(default_factory=dict)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict) # 额外元数据
|
||||
|
||||
|
||||
@@ -65,7 +66,7 @@ class DomainRegistry:
|
||||
|
||||
_registry: Dict[str, DomainRegistration] = {}
|
||||
|
||||
_HANDLER_ATTRS: Dict[str, str] = {
|
||||
_BUILTIN_HANDLER_ATTRS: Dict[str, str] = {
|
||||
"jsonl": "jsonl_handler_class",
|
||||
"api": "api_handler_class",
|
||||
"db": "db_handler_class",
|
||||
@@ -102,7 +103,7 @@ class DomainRegistry:
|
||||
if node_type in cls._registry:
|
||||
raise ValueError(f"Domain '{node_type}' already registered")
|
||||
|
||||
cls._registry[node_type] = DomainRegistration(
|
||||
reg = DomainRegistration(
|
||||
node_type=node_type,
|
||||
schema=schema,
|
||||
node_class=node_class,
|
||||
@@ -112,6 +113,12 @@ class DomainRegistry:
|
||||
db_handler_class=db_handler_class,
|
||||
metadata=metadata
|
||||
)
|
||||
for datasource_type, attr_name in cls._BUILTIN_HANDLER_ATTRS.items():
|
||||
handler_class = getattr(reg, attr_name)
|
||||
if handler_class is not None:
|
||||
reg.handler_classes[datasource_type] = handler_class
|
||||
|
||||
cls._registry[node_type] = reg
|
||||
|
||||
@classmethod
|
||||
def get_node_class(cls, node_type: str) -> Type[Any]:
|
||||
@@ -152,27 +159,37 @@ class DomainRegistry:
|
||||
def get_handler(
|
||||
cls,
|
||||
node_type: str,
|
||||
datasource_type: Literal["jsonl", "api", "db"],
|
||||
datasource_type: str,
|
||||
) -> Optional[Type[Any]]:
|
||||
"""按 datasource 类型获取 Handler 类。"""
|
||||
reg = cls._registry.get(node_type)
|
||||
if not reg:
|
||||
raise ValueError(f"Unknown node_type: '{node_type}'")
|
||||
attr_name = cls._HANDLER_ATTRS.get(datasource_type)
|
||||
handler_class = reg.handler_classes.get(datasource_type)
|
||||
if handler_class is not None:
|
||||
return handler_class
|
||||
|
||||
attr_name = cls._BUILTIN_HANDLER_ATTRS.get(datasource_type)
|
||||
if attr_name is None:
|
||||
raise ValueError(
|
||||
f"Unsupported datasource_type: '{datasource_type}'. "
|
||||
f"Available: {sorted(cls._HANDLER_ATTRS.keys())}"
|
||||
)
|
||||
return None
|
||||
return getattr(reg, attr_name)
|
||||
|
||||
@classmethod
|
||||
def register_handler(cls, node_type: str, datasource_type: str, handler_class: Type[Any]) -> None:
|
||||
"""为已注册 domain 补充任意 datasource 类型的 Handler。"""
|
||||
reg = cls._registry.get(node_type)
|
||||
if not reg:
|
||||
raise ValueError(f"Unknown node_type: '{node_type}'")
|
||||
|
||||
reg.handler_classes[datasource_type] = handler_class
|
||||
attr_name = cls._BUILTIN_HANDLER_ATTRS.get(datasource_type)
|
||||
if attr_name is not None:
|
||||
setattr(reg, attr_name, handler_class)
|
||||
|
||||
@classmethod
|
||||
def register_db_handler(cls, node_type: str, db_handler_class: Type[Any]) -> None:
|
||||
"""为已注册 domain 补充 DB Handler。"""
|
||||
reg = cls._registry.get(node_type)
|
||||
if not reg:
|
||||
raise ValueError(f"Unknown node_type: '{node_type}'")
|
||||
reg.db_handler_class = db_handler_class
|
||||
cls.register_handler(node_type, "db", db_handler_class)
|
||||
|
||||
@classmethod
|
||||
def get_schema(cls, node_type: str) -> Type[BaseModel]:
|
||||
@@ -229,6 +246,11 @@ class DomainRegistry:
|
||||
print(f" JSONL Handler: {reg.jsonl_handler_class.__name__ if reg.jsonl_handler_class else '❌ None'}")
|
||||
print(f" API Handler: {reg.api_handler_class.__name__ if reg.api_handler_class else '❌ None'}")
|
||||
print(f" DB Handler: {reg.db_handler_class.__name__ if reg.db_handler_class else '❌ None'}")
|
||||
extra_handler_types = sorted(
|
||||
handler_type for handler_type in reg.handler_classes if handler_type not in cls._BUILTIN_HANDLER_ATTRS
|
||||
)
|
||||
if extra_handler_types:
|
||||
print(f" Extra Handlers: {extra_handler_types}")
|
||||
if reg.metadata:
|
||||
print(f" Metadata: {reg.metadata}")
|
||||
print()
|
||||
@@ -236,7 +258,7 @@ class DomainRegistry:
|
||||
print(f"{'='*80}\n")
|
||||
|
||||
@classmethod
|
||||
def get_summary(cls) -> Dict[str, Dict[str, bool]]:
|
||||
def get_summary(cls) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
获取注册摘要(用于自动化检查)
|
||||
|
||||
@@ -263,5 +285,6 @@ class DomainRegistry:
|
||||
"has_jsonl_handler": reg.jsonl_handler_class is not None,
|
||||
"has_api_handler": reg.api_handler_class is not None,
|
||||
"has_db_handler": reg.db_handler_class is not None,
|
||||
"handler_types": sorted(reg.handler_classes.keys()),
|
||||
}
|
||||
return summary
|
||||
|
||||
@@ -1,6 +1,14 @@
|
||||
from .datasource_config import ApiDataSourceConfig, DbDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .datasource_config import ApiDataSourceConfig, BaseDataSourceConfig, DbDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .persist_config import PersistConfig
|
||||
from .logging_config import LoggingConfig, resolve_log_file_path
|
||||
from ..datasource.type_registry import (
|
||||
get_datasource_config_model,
|
||||
get_datasource_factory,
|
||||
list_registered_datasource_types,
|
||||
register_datasource_config_model,
|
||||
register_datasource_factory,
|
||||
register_datasource_type,
|
||||
)
|
||||
from .run_presets import (
|
||||
list_preset_names,
|
||||
get_preset_file_candidates,
|
||||
@@ -26,12 +34,19 @@ from .builder import (
|
||||
|
||||
__all__ = [
|
||||
"ApiDataSourceConfig",
|
||||
"BaseDataSourceConfig",
|
||||
"DbDataSourceConfig",
|
||||
"JsonlDataSourceConfig",
|
||||
"PipelineRunConfig",
|
||||
"DataSourceConfig",
|
||||
"get_datasource_config_model",
|
||||
"get_datasource_factory",
|
||||
"PersistConfig",
|
||||
"LoggingConfig",
|
||||
"list_registered_datasource_types",
|
||||
"register_datasource_config_model",
|
||||
"register_datasource_factory",
|
||||
"register_datasource_type",
|
||||
"resolve_log_file_path",
|
||||
"list_preset_names",
|
||||
"get_preset_file_candidates",
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Mapping, Optional, Union
|
||||
|
||||
@@ -9,6 +8,7 @@ from pydantic import BaseModel
|
||||
from ..common.registry import DomainRegistry
|
||||
from .datasource_config import DataSourceConfig
|
||||
from .pipeline_config import PipelineRunConfig
|
||||
from .run_presets import load_overrides_from_file
|
||||
from .strategy_config import ConfigPresets, StrategyRuntimeConfig
|
||||
|
||||
|
||||
@@ -148,7 +148,7 @@ def apply_env_overrides(
|
||||
dotenv_file: Optional[Union[str, Path]] = None,
|
||||
) -> PipelineRunConfig:
|
||||
raise NotImplementedError(
|
||||
"Environment-based overrides are removed. Use JSON config files via overrides_file instead."
|
||||
"Environment-based overrides are removed. Use config files via overrides_file instead."
|
||||
)
|
||||
|
||||
|
||||
@@ -162,7 +162,7 @@ def build_config(
|
||||
dotenv_file: Optional[Union[str, Path]] = None,
|
||||
) -> PipelineRunConfig:
|
||||
if dotenv_file is not None:
|
||||
raise ValueError("dotenv_file is no longer supported; use overrides_file JSON instead")
|
||||
raise ValueError("dotenv_file is no longer supported; use overrides_file instead")
|
||||
|
||||
config = build_default_config(project_root)
|
||||
|
||||
@@ -180,7 +180,7 @@ def build_config(
|
||||
|
||||
if overrides_file:
|
||||
override_path = Path(overrides_file)
|
||||
file_data = json.loads(override_path.read_text(encoding="utf-8"))
|
||||
file_data = load_overrides_from_file(override_path, project_root)
|
||||
config = apply_config_overrides(config, file_data)
|
||||
|
||||
if overrides:
|
||||
|
||||
@@ -1,30 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Any, Dict, List, Literal, Union
|
||||
from typing import Annotated, Any, Dict, List, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator
|
||||
|
||||
from ..datasource.type_registry import register_datasource_config_model, validate_registered_datasource_config
|
||||
|
||||
|
||||
class JsonlDataSourceConfig(BaseModel):
|
||||
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
|
||||
target_project_ids: List[str] = Field(default_factory=list)
|
||||
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class DbDataSourceConfig(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
||||
class DbDataSourceConfig(BaseDataSourceConfig):
|
||||
|
||||
type: Literal["db"] = "db"
|
||||
target_project_ids: List[str] = Field(default_factory=list)
|
||||
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ApiDataSourceConfig(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
||||
class ApiDataSourceConfig(BaseDataSourceConfig):
|
||||
|
||||
type: Literal["api"] = "api"
|
||||
api_base_url: str = ""
|
||||
@@ -36,10 +39,18 @@ class ApiDataSourceConfig(BaseModel):
|
||||
api_rate_limit_max_requests: int = 10
|
||||
api_rate_limit_window_seconds: float = 10.0
|
||||
target_project_ids: List[str] = Field(min_length=1)
|
||||
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
DataSourceConfig = Annotated[
|
||||
Union[JsonlDataSourceConfig, DbDataSourceConfig, ApiDataSourceConfig],
|
||||
Field(discriminator="type"),
|
||||
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)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
@@ -28,6 +29,43 @@ def _replace_placeholders(value: Any, project_root: Path) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
def _import_from_path(import_path: str) -> Any:
|
||||
module_path, sep, attr_path = import_path.partition(":")
|
||||
if not module_path:
|
||||
raise ValueError(f"Invalid import path: {import_path}")
|
||||
|
||||
module = importlib.import_module(module_path)
|
||||
if not sep:
|
||||
return module
|
||||
|
||||
obj: Any = module
|
||||
for attr_name in attr_path.split("."):
|
||||
obj = getattr(obj, attr_name)
|
||||
return obj
|
||||
|
||||
|
||||
def _apply_datasource_bootstraps(payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
extensions = payload.pop("extensions", None)
|
||||
if extensions is None:
|
||||
return payload
|
||||
if not isinstance(extensions, dict):
|
||||
raise ValueError("extensions must be an object mapping")
|
||||
|
||||
bootstrap_paths = extensions.get("datasource_bootstraps") or []
|
||||
if not isinstance(bootstrap_paths, list):
|
||||
raise ValueError("extensions.datasource_bootstraps must be a list")
|
||||
|
||||
for item in bootstrap_paths:
|
||||
if not isinstance(item, str) or not item.strip():
|
||||
raise ValueError("extensions.datasource_bootstraps entries must be non-empty strings")
|
||||
loaded = _import_from_path(item)
|
||||
if not callable(loaded):
|
||||
raise TypeError(f"datasource bootstrap must resolve to a callable: {item}")
|
||||
loaded()
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def list_preset_names() -> List[str]:
|
||||
return sorted(PRESET_FILES.keys())
|
||||
|
||||
@@ -77,7 +115,7 @@ def load_overrides_from_file(file_path: Path, project_root: Path) -> Dict[str, A
|
||||
resolved = _replace_placeholders(payload, project_root)
|
||||
if not isinstance(resolved, dict):
|
||||
raise ValueError(f"Preset file must contain object mapping: {file_path}")
|
||||
return resolved
|
||||
return _apply_datasource_bootstraps(resolved)
|
||||
|
||||
|
||||
def load_named_preset_overrides(project_root: Path, preset_name: str) -> Tuple[Dict[str, Any], Path, bool]:
|
||||
|
||||
@@ -6,6 +6,14 @@ Handler 负责类型转换和业务逻辑。
|
||||
"""
|
||||
|
||||
from .datasource import BaseDataSource
|
||||
from .type_registry import (
|
||||
get_datasource_config_model,
|
||||
get_datasource_factory,
|
||||
list_registered_datasource_types,
|
||||
register_datasource_config_model,
|
||||
register_datasource_factory,
|
||||
register_datasource_type,
|
||||
)
|
||||
from .handler import (
|
||||
NodeHandler,
|
||||
BaseNodeHandler,
|
||||
@@ -26,9 +34,15 @@ from .jsonl import JsonlDataSource, BaseJsonlHandler
|
||||
__all__ = [
|
||||
# 基础类
|
||||
"BaseDataSource",
|
||||
"get_datasource_config_model",
|
||||
"get_datasource_factory",
|
||||
"list_registered_datasource_types",
|
||||
"NodeHandler",
|
||||
"BaseNodeHandler",
|
||||
"NodeHandlerRegistry",
|
||||
"register_datasource_config_model",
|
||||
"register_datasource_factory",
|
||||
"register_datasource_type",
|
||||
"TaskResult",
|
||||
"TaskStatus",
|
||||
"ValidationResult",
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
DataSourceFactory = Callable[..., Any | Awaitable[Any]]
|
||||
|
||||
|
||||
_CONFIG_MODELS: dict[str, type[BaseModel]] = {}
|
||||
_FACTORIES: dict[str, DataSourceFactory] = {}
|
||||
|
||||
|
||||
def _normalize_type_name(type_name: str) -> str:
|
||||
normalized = str(type_name).strip()
|
||||
if not normalized:
|
||||
raise ValueError("datasource type name must not be empty")
|
||||
return normalized
|
||||
|
||||
|
||||
def register_datasource_config_model(
|
||||
type_name: str,
|
||||
config_model: type[BaseModel],
|
||||
*,
|
||||
replace: bool = False,
|
||||
) -> None:
|
||||
normalized = _normalize_type_name(type_name)
|
||||
existing = _CONFIG_MODELS.get(normalized)
|
||||
if existing is not None and existing is not config_model and not replace:
|
||||
raise ValueError(f"datasource config model already registered for type: {normalized}")
|
||||
_CONFIG_MODELS[normalized] = config_model
|
||||
|
||||
|
||||
def register_datasource_factory(
|
||||
type_name: str,
|
||||
factory: DataSourceFactory,
|
||||
*,
|
||||
replace: bool = False,
|
||||
) -> None:
|
||||
normalized = _normalize_type_name(type_name)
|
||||
existing = _FACTORIES.get(normalized)
|
||||
if existing is not None and existing is not factory and not replace:
|
||||
raise ValueError(f"datasource factory already registered for type: {normalized}")
|
||||
_FACTORIES[normalized] = factory
|
||||
|
||||
|
||||
def register_datasource_type(
|
||||
type_name: str,
|
||||
*,
|
||||
config_model: type[BaseModel] | None = None,
|
||||
factory: DataSourceFactory | None = None,
|
||||
replace: bool = False,
|
||||
) -> None:
|
||||
if config_model is None and factory is None:
|
||||
raise ValueError("register_datasource_type() requires config_model, factory, or both")
|
||||
|
||||
if config_model is not None:
|
||||
register_datasource_config_model(type_name, config_model, replace=replace)
|
||||
if factory is not None:
|
||||
register_datasource_factory(type_name, factory, replace=replace)
|
||||
|
||||
|
||||
def get_datasource_config_model(type_name: str) -> type[BaseModel] | None:
|
||||
return _CONFIG_MODELS.get(_normalize_type_name(type_name))
|
||||
|
||||
|
||||
def get_datasource_factory(type_name: str) -> DataSourceFactory | None:
|
||||
return _FACTORIES.get(_normalize_type_name(type_name))
|
||||
|
||||
|
||||
def list_registered_datasource_types() -> list[str]:
|
||||
return sorted(set(_CONFIG_MODELS) | set(_FACTORIES))
|
||||
|
||||
|
||||
def validate_registered_datasource_config(value: Any) -> BaseModel:
|
||||
if isinstance(value, BaseModel):
|
||||
payload = value.model_dump(mode="python")
|
||||
elif isinstance(value, dict):
|
||||
payload = dict(value)
|
||||
else:
|
||||
raise TypeError(f"datasource config must be a mapping or BaseModel, got {type(value).__name__}")
|
||||
|
||||
type_name = payload.get("type")
|
||||
if not isinstance(type_name, str) or not type_name.strip():
|
||||
raise ValueError("datasource config must define a non-empty string field 'type'")
|
||||
|
||||
config_model = get_datasource_config_model(type_name)
|
||||
if config_model is None:
|
||||
available = ", ".join(list_registered_datasource_types()) or "<none>"
|
||||
raise ValueError(f"Unsupported datasource type: '{type_name}'. Registered types: {available}")
|
||||
|
||||
return config_model.model_validate(payload)
|
||||
@@ -13,6 +13,7 @@ Pipeline Factory - 工厂函数,负责初始化所有组件并创建 Pipeline
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
@@ -25,6 +26,7 @@ from ..common.registry import DomainRegistry
|
||||
from ..datasource.api.datasource import ApiDataSource
|
||||
from ..datasource.db.datasource import DbDataSource
|
||||
from ..datasource.jsonl.datasource import JsonlDataSource
|
||||
from ..datasource.type_registry import get_datasource_factory, register_datasource_factory
|
||||
from ..config import (
|
||||
PipelineRunConfig,
|
||||
DataSourceConfig,
|
||||
@@ -156,19 +158,34 @@ async def _create_datasource(
|
||||
if datasource_override is not None:
|
||||
return datasource_override
|
||||
|
||||
if isinstance(ds_config, JsonlDataSourceConfig):
|
||||
dir_path = Path(ds_config.jsonl_dir)
|
||||
if endpoint_name == "remote":
|
||||
dir_path = _prepare_remote_jsonl_dir(dir_path)
|
||||
elif not dir_path.exists() or not dir_path.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"{endpoint_name}_datasource.jsonl_dir does not exist or is not a directory: {dir_path}"
|
||||
)
|
||||
return JsonlDataSource(dir_path, read_only=ds_config.read_only)
|
||||
datasource_factory = get_datasource_factory(ds_config.type)
|
||||
if datasource_factory is None:
|
||||
raise RuntimeError(f"No datasource factory registered for type: {ds_config.type}")
|
||||
|
||||
if isinstance(ds_config, DbDataSourceConfig):
|
||||
return DbDataSource()
|
||||
datasource = datasource_factory(ds_config, endpoint_name)
|
||||
if inspect.isawaitable(datasource):
|
||||
datasource = await datasource
|
||||
return datasource
|
||||
|
||||
|
||||
def _build_jsonl_datasource(ds_config: JsonlDataSourceConfig, endpoint_name: str):
|
||||
dir_path = Path(ds_config.jsonl_dir)
|
||||
if endpoint_name == "remote":
|
||||
dir_path = _prepare_remote_jsonl_dir(dir_path)
|
||||
elif not dir_path.exists() or not dir_path.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"{endpoint_name}_datasource.jsonl_dir does not exist or is not a directory: {dir_path}"
|
||||
)
|
||||
return JsonlDataSource(dir_path, read_only=ds_config.read_only)
|
||||
|
||||
|
||||
def _build_db_datasource(ds_config: DbDataSourceConfig, endpoint_name: str):
|
||||
del ds_config, endpoint_name
|
||||
return DbDataSource()
|
||||
|
||||
|
||||
async def _build_api_datasource(ds_config: ApiDataSourceConfig, endpoint_name: str):
|
||||
del endpoint_name
|
||||
datasource = ApiDataSource(
|
||||
base_url=ds_config.api_base_url,
|
||||
uid=ds_config.api_uid,
|
||||
@@ -183,6 +200,18 @@ async def _create_datasource(
|
||||
return datasource
|
||||
|
||||
|
||||
def _ensure_builtin_datasource_factories_registered() -> None:
|
||||
if get_datasource_factory("jsonl") is None:
|
||||
register_datasource_factory("jsonl", _build_jsonl_datasource)
|
||||
if get_datasource_factory("db") is None:
|
||||
register_datasource_factory("db", _build_db_datasource)
|
||||
if get_datasource_factory("api") is None:
|
||||
register_datasource_factory("api", _build_api_datasource)
|
||||
|
||||
|
||||
_ensure_builtin_datasource_factories_registered()
|
||||
|
||||
|
||||
def _register_handlers_for_endpoint(
|
||||
*,
|
||||
datasource,
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from sync_state_machine.common.registry import DomainRegistry
|
||||
from sync_state_machine.common.sync_node import SyncNode
|
||||
from sync_state_machine.config import (
|
||||
BaseDataSourceConfig,
|
||||
DbDataSourceConfig,
|
||||
LoggingConfig,
|
||||
PersistConfig,
|
||||
PipelineRunConfig,
|
||||
build_config,
|
||||
register_datasource_type,
|
||||
)
|
||||
from sync_state_machine.datasource import BaseDataSource
|
||||
from sync_state_machine.datasource.db import BaseDbHandler, DbDataSource
|
||||
from sync_state_machine.pipeline.factory import create_pipeline_from_config
|
||||
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
|
||||
|
||||
|
||||
class MemoryConfig(BaseDataSourceConfig):
|
||||
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
||||
|
||||
type: str = "memory"
|
||||
namespace: str = "default"
|
||||
|
||||
|
||||
class MemorySchema(BaseModel):
|
||||
__test__ = False
|
||||
|
||||
id: str
|
||||
name: str | None = None
|
||||
|
||||
|
||||
class MemoryNode(SyncNode[MemorySchema]):
|
||||
__test__ = False
|
||||
|
||||
node_type = "memory_demo"
|
||||
schema = MemorySchema
|
||||
|
||||
|
||||
class MemoryStrategy(DefaultSyncStrategy[MemorySchema]):
|
||||
__test__ = False
|
||||
|
||||
schema = MemorySchema
|
||||
|
||||
|
||||
class MemoryDataSource(BaseDataSource):
|
||||
def __init__(self, namespace: str):
|
||||
super().__init__()
|
||||
self.namespace = namespace
|
||||
|
||||
|
||||
class MemoryHandler(BaseDbHandler[MemorySchema]):
|
||||
__test__ = False
|
||||
|
||||
def __init__(self, datasource: MemoryDataSource):
|
||||
super().__init__(datasource, "memory_demo", MemoryNode, MemorySchema)
|
||||
|
||||
def extract_id(self, data: dict[str, str]) -> str:
|
||||
return str(data.get("id", ""))
|
||||
|
||||
async def load(self):
|
||||
return []
|
||||
|
||||
async def create_all(self, nodes):
|
||||
return []
|
||||
|
||||
async def update_all(self, nodes):
|
||||
return []
|
||||
|
||||
|
||||
class MemoryRemoteDbHandler(BaseDbHandler[MemorySchema]):
|
||||
__test__ = False
|
||||
|
||||
def __init__(self, datasource: DbDataSource):
|
||||
super().__init__(datasource, "memory_demo", MemoryNode, MemorySchema)
|
||||
|
||||
def extract_id(self, data: dict[str, str]) -> str:
|
||||
return str(data.get("id", ""))
|
||||
|
||||
async def load(self):
|
||||
return []
|
||||
|
||||
async def create_all(self, nodes):
|
||||
return []
|
||||
|
||||
async def update_all(self, nodes):
|
||||
return []
|
||||
|
||||
|
||||
def _ensure_memory_domain_registered() -> None:
|
||||
if DomainRegistry.is_registered("memory_demo"):
|
||||
reg = DomainRegistry.get_registration("memory_demo")
|
||||
reg.schema = MemorySchema
|
||||
reg.node_class = MemoryNode
|
||||
reg.strategy_class = MemoryStrategy
|
||||
reg.db_handler_class = MemoryRemoteDbHandler
|
||||
reg.handler_classes["db"] = MemoryRemoteDbHandler
|
||||
reg.handler_classes["memory"] = MemoryHandler
|
||||
return
|
||||
|
||||
DomainRegistry.register(
|
||||
node_type="memory_demo",
|
||||
schema=MemorySchema,
|
||||
node_class=MemoryNode,
|
||||
strategy_class=MemoryStrategy,
|
||||
db_handler_class=MemoryRemoteDbHandler,
|
||||
)
|
||||
DomainRegistry.register_handler("memory_demo", "memory", MemoryHandler)
|
||||
|
||||
|
||||
def _memory_factory(config: MemoryConfig, endpoint_name: str) -> MemoryDataSource:
|
||||
return MemoryDataSource(namespace=f"{endpoint_name}:{config.namespace}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_pipeline_from_config_supports_registered_custom_datasource_type(tmp_path: Path) -> None:
|
||||
register_datasource_type("memory", config_model=MemoryConfig, factory=_memory_factory, replace=True)
|
||||
_ensure_memory_domain_registered()
|
||||
|
||||
config = PipelineRunConfig(
|
||||
node_types=["memory_demo"],
|
||||
local_datasource=MemoryConfig(type="memory", namespace="alpha"),
|
||||
remote_datasource=DbDataSourceConfig(type="db"),
|
||||
strategies=[],
|
||||
persist=PersistConfig(db_path=str(tmp_path / "memory-registry.db")),
|
||||
logging=LoggingConfig(initialize=False),
|
||||
)
|
||||
|
||||
pipeline = await create_pipeline_from_config(config)
|
||||
try:
|
||||
assert isinstance(pipeline.local_datasource, MemoryDataSource)
|
||||
assert pipeline.local_datasource.namespace == "local:alpha"
|
||||
assert isinstance(pipeline.local_datasource.get_handler("memory_demo"), MemoryHandler)
|
||||
assert isinstance(pipeline.remote_datasource.get_handler("memory_demo"), MemoryRemoteDbHandler)
|
||||
finally:
|
||||
await pipeline.close()
|
||||
|
||||
|
||||
def test_build_config_supports_same_file_datasource_bootstrap(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
bootstrap_module = tmp_path / "memory_bootstrap_demo.py"
|
||||
bootstrap_module.write_text(
|
||||
"from pydantic import ConfigDict\n"
|
||||
"from sync_state_machine.config import BaseDataSourceConfig, register_datasource_type\n"
|
||||
"from sync_state_machine.datasource import BaseDataSource\n"
|
||||
"\n"
|
||||
"class BootstrapMemoryConfig(BaseDataSourceConfig):\n"
|
||||
" model_config = ConfigDict(extra=\"forbid\", validate_assignment=True)\n"
|
||||
" type: str = \"memory_bootstrap\"\n"
|
||||
" namespace: str = \"demo\"\n"
|
||||
"\n"
|
||||
"class BootstrapMemoryDataSource(BaseDataSource):\n"
|
||||
" def __init__(self, namespace: str):\n"
|
||||
" super().__init__()\n"
|
||||
" self.namespace = namespace\n"
|
||||
"\n"
|
||||
"def build_memory_datasource(config, endpoint_name):\n"
|
||||
" return BootstrapMemoryDataSource(namespace=f\"{endpoint_name}:{config.namespace}\")\n"
|
||||
"\n"
|
||||
"def register():\n"
|
||||
" register_datasource_type(\n"
|
||||
" \"memory_bootstrap\",\n"
|
||||
" config_model=BootstrapMemoryConfig,\n"
|
||||
" factory=build_memory_datasource,\n"
|
||||
" replace=True,\n"
|
||||
" )\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
|
||||
profile_path = tmp_path / "custom-memory.yaml"
|
||||
profile_path.write_text(
|
||||
"extensions:\n"
|
||||
" datasource_bootstraps:\n"
|
||||
" - memory_bootstrap_demo:register\n"
|
||||
"node_types: []\n"
|
||||
"local_datasource:\n"
|
||||
" type: memory_bootstrap\n"
|
||||
" namespace: local-demo\n"
|
||||
"remote_datasource:\n"
|
||||
" type: db\n"
|
||||
"persist:\n"
|
||||
" db_path: ${PROJECT_ROOT}/runtime/test-memory-bootstrap.db\n"
|
||||
"logging:\n"
|
||||
" initialize: false\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
config = build_config(project_root=tmp_path, overrides_file=profile_path)
|
||||
|
||||
assert config.local_datasource.type == "memory_bootstrap"
|
||||
assert getattr(config.local_datasource, "namespace") == "local-demo"
|
||||
assert config.remote_datasource.type == "db"
|
||||
Reference in New Issue
Block a user