整理了初始化过程,现在更容易被继承了。

This commit is contained in:
strepsiades
2026-03-18 16:48:53 +08:00
parent 84b3b7652a
commit 841986740c
22 changed files with 432 additions and 277 deletions
+17 -4
View File
@@ -23,13 +23,25 @@ from .handler import (
from .task_result import TaskResult, TaskStatus
# API DataSource
from .api import ApiClient, ApiDataSource, BaseApiHandler
from .api import ApiClient, ApiDataSource, BaseApiHandler, register as register_api_datasource
# DB DataSource
# DB base helpers
from .db import DbDataSource, BaseDbHandler
# JSONL DataSource
from .jsonl import JsonlDataSource, BaseJsonlHandler
from .jsonl import JsonlDataSource, BaseJsonlHandler, register as register_jsonl_datasource
_BUILTIN_DATASOURCE_TYPES_REGISTERED = False
def ensure_builtin_datasource_types_registered(*, replace: bool = False) -> None:
global _BUILTIN_DATASOURCE_TYPES_REGISTERED
if _BUILTIN_DATASOURCE_TYPES_REGISTERED and not replace:
return
register_jsonl_datasource(replace=replace)
register_api_datasource(replace=replace)
_BUILTIN_DATASOURCE_TYPES_REGISTERED = True
__all__ = [
# 基础类
@@ -46,13 +58,14 @@ __all__ = [
"TaskResult",
"TaskStatus",
"ValidationResult",
"ensure_builtin_datasource_types_registered",
# API DataSource
"ApiClient",
"ApiDataSource",
"BaseApiHandler",
# DB DataSource
# DB base helpers
"DbDataSource",
"BaseDbHandler",
@@ -1,15 +1,50 @@
from __future__ import annotations
"""
API DataSource 模块
提供基于 HTTP API 的数据源实现。
"""
from typing import TYPE_CHECKING
from ..type_registry import register_datasource_type
from .client import ApiClient
from .datasource import ApiDataSource
from .handler import BaseApiHandler
if TYPE_CHECKING:
from ...config.datasource_config import ApiDataSourceConfig
def build_api_datasource(config: ApiDataSourceConfig, endpoint_name: str, **_: object):
del endpoint_name
return ApiDataSource(
base_url=config.api_base_url,
uid=config.api_uid,
secret=config.api_secret,
debug=config.api_debug,
poll_max_retries=config.poll_max_retries,
poll_interval=config.poll_interval,
rate_limit_max_requests=config.api_rate_limit_max_requests,
rate_limit_window_seconds=config.api_rate_limit_window_seconds,
)
def register(*, replace: bool = False) -> None:
from ...config.datasource_config import ApiDataSourceConfig
register_datasource_type(
"api",
config_model=ApiDataSourceConfig,
factory=build_api_datasource,
replace=replace,
)
__all__ = [
"ApiClient",
"ApiDataSource",
"BaseApiHandler",
"build_api_datasource",
"register",
]
+12 -1
View File
@@ -88,6 +88,13 @@ class BaseDataSource(ABC):
self._poll_interval = max(0, poll_interval)
self._sm_runtime: Optional[StateMachineRuntime] = None
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
self._target_project_ids = [str(item) for item in target_project_ids if str(item)]
@property
def target_project_ids(self) -> List[str]:
return list(self._target_project_ids)
def _ensure_sm_runtime(self) -> StateMachineRuntime:
if self._sm_runtime is None:
cfg_path = Path(__file__).resolve().parents[1] / "config" / "node_state_machine.yaml"
@@ -957,6 +964,10 @@ class BaseDataSource(ABC):
return summary
async def initialize(self) -> None:
"""初始化数据源资源(子类可覆盖)"""
return None
async def close(self) -> None:
"""关闭数据源资源(子类可覆盖以实现特定的清理逻辑)"""
pass # 基类默认空实现
return None
+3 -1
View File
@@ -1,4 +1,6 @@
"""DB DataSource 模块。"""
from __future__ import annotations
"""DB DataSource 基类模块。"""
from .datasource import DbDataSource
from .handler import BaseDbHandler
@@ -1,13 +1,53 @@
from __future__ import annotations
"""
JSONL DataSource 模块
提供基于 JSONL 文件的数据源实现,用于测试和本地开发。
"""
from pathlib import Path
from typing import TYPE_CHECKING
from ..type_registry import register_datasource_type
from .datasource import JsonlDataSource
from .handler import BaseJsonlHandler
from .path_utils import prepare_remote_jsonl_dir
if TYPE_CHECKING:
from ...config.datasource_config import JsonlDataSourceConfig
def build_jsonl_datasource(
config: JsonlDataSourceConfig,
endpoint_name: str,
*,
project_root: Path | None = None,
):
dir_path = Path(config.jsonl_dir)
if endpoint_name == "remote":
dir_path = prepare_remote_jsonl_dir(dir_path, project_root=project_root or Path.cwd())
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=config.read_only)
def register(*, replace: bool = False) -> None:
from ...config.datasource_config import JsonlDataSourceConfig
register_datasource_type(
"jsonl",
config_model=JsonlDataSourceConfig,
factory=build_jsonl_datasource,
replace=replace,
)
__all__ = [
"JsonlDataSource",
"BaseJsonlHandler",
"build_jsonl_datasource",
"prepare_remote_jsonl_dir",
"register",
]
@@ -50,24 +50,22 @@ class BaseJsonlHandler(BaseNodeHandler):
):
super().__init__(node_type, node_class, schema)
self.datasource = datasource
self._target_project_ids: Optional[Set[str]] = None
self._target_project_ids: List[str] = []
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
"""注入项目过滤列表,供 load 阶段预过滤使用"""
if target_project_ids:
self._target_project_ids = {str(pid) for pid in target_project_ids if str(pid)}
else:
self._target_project_ids = None
"""注入项目过滤列表的本地副本。真正的过滤来源仍然是 datasource"""
self._target_project_ids = [str(pid) for pid in target_project_ids if str(pid)]
def _should_skip_item_by_project_filter(self, item: Dict[str, Any], item_id: str) -> bool:
if not self._target_project_ids:
target_project_ids = set(self.datasource.target_project_ids)
if not target_project_ids:
return False
if self.node_type == "project":
return item_id not in self._target_project_ids
return item_id not in target_project_ids
project_id = item.get("project_id")
if project_id:
return str(project_id) not in self._target_project_ids
return str(project_id) not in target_project_ids
return False
async def load(self) -> List['SyncNode']:
@@ -0,0 +1,31 @@
from __future__ import annotations
import logging
from pathlib import Path
logger = logging.getLogger(__name__)
def prepare_remote_jsonl_dir(remote_dir: Path, *, project_root: Path) -> Path:
project_root_path = project_root.resolve()
resolved_remote_dir = remote_dir.resolve(strict=False)
if resolved_remote_dir.exists():
if not resolved_remote_dir.is_dir():
raise FileNotFoundError(
f"remote_datasource.jsonl_dir exists but is not a directory: {resolved_remote_dir}"
)
return resolved_remote_dir
try:
resolved_remote_dir.relative_to(project_root_path)
except ValueError as exc:
raise FileNotFoundError(
"remote_datasource.jsonl_dir does not exist and auto-create is only allowed "
f"under project root: {resolved_remote_dir} (project_root={project_root_path})"
) from exc
resolved_remote_dir.mkdir(parents=True, exist_ok=True)
logger.info("Created missing remote JSONL directory: %s", resolved_remote_dir)
return resolved_remote_dir