整理配置和代码入口，增加部分测试。

2026-03-10 12:22:54 +08:00
parent 515c3d9388
commit a638e4203b
35 changed files with 1314 additions and 199 deletions
@@ -7,6 +7,7 @@ from .full_sync_pipeline import FullSyncPipeline
 from ..config import (
    PipelineRunConfig,
    ApiDataSourceConfig,
+    DbDataSourceConfig,
    JsonlDataSourceConfig,
    DataSourceConfig,
    PersistConfig,
@@ -35,6 +36,7 @@ __all__ = [
    "run_pipeline_from_config",
    "PipelineRunConfig",
    "ApiDataSourceConfig",
+    "DbDataSourceConfig",
    "JsonlDataSourceConfig",
    "DataSourceConfig",
    "PersistConfig",
@@ -23,11 +23,13 @@ from ..common.collection import DataCollection
 from ..common.persistence import PersistenceBackend
 from ..common.registry import DomainRegistry
 from ..datasource.api.datasource import ApiDataSource
+from ..datasource.db.datasource import DbDataSource
 from ..datasource.jsonl.datasource import JsonlDataSource
 from ..config import (
    PipelineRunConfig,
    DataSourceConfig,
    ApiDataSourceConfig,
+    DbDataSourceConfig,
    JsonlDataSourceConfig,
    PersistConfig,
    LoggingConfig,
@@ -97,16 +99,41 @@ def _ensure_pipeline_logger(
    )


-def _api_handler_kwargs(
+def _build_handler_config(
    node_type: str,
    ds_config: DataSourceConfig,
    target_project_ids: List[str],
 ) -> Dict[str, Any]:
-    kwargs = dict(ds_config.handler_configs.get(node_type, {}))
+    config = dict(ds_config.handler_configs.get(node_type, {}))

    if node_type == "project" and target_project_ids:
-        kwargs.setdefault("filter_project_id", target_project_ids)
-    return kwargs
+        config.setdefault("filter_project_id", target_project_ids)
+    return config
+
+
+def _create_handler(
+    *,
+    node_type: str,
+    ds_config: DataSourceConfig,
+    datasource,
+    target_project_ids: List[str],
+):
+    datasource_type = ds_config.type
+    handler_class = DomainRegistry.get_handler(node_type, datasource_type)
+    if handler_class is None:
+        raise RuntimeError(
+            f"Missing {datasource_type.upper()} handler for node_type: {node_type}"
+        )
+
+    handler_config = _build_handler_config(node_type, ds_config, target_project_ids)
+    if datasource_type == "api":
+        handler = handler_class(**handler_config)
+    else:
+        handler = handler_class(datasource)
+
+    if hasattr(handler, "set_handler_config"):
+        handler.set_handler_config(handler_config)
+    return handler


 def _resolve_datasource_target_project_ids(config: PipelineRunConfig) -> tuple[List[str], List[str]]:
@@ -120,48 +147,58 @@ def _resolve_datasource_target_project_ids(config: PipelineRunConfig) -> tuple[L
    return local_target_ids, remote_target_ids


-async def _create_datasources(config: PipelineRunConfig):
-    local_cfg = config.local_datasource
-    remote_cfg = config.remote_datasource
+async def _create_datasource(
+    ds_config: DataSourceConfig,
+    *,
+    datasource_override=None,
+    endpoint_name: str,
+):
+    if datasource_override is not None:
+        return datasource_override

-    if isinstance(local_cfg, JsonlDataSourceConfig):
-        local_dir = Path(local_cfg.jsonl_dir)
-        if not local_dir.exists() or not local_dir.is_dir():
+    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"local_datasource.jsonl_dir does not exist or is not a directory: {local_dir}"
+                f"{endpoint_name}_datasource.jsonl_dir does not exist or is not a directory: {dir_path}"
            )
-        local_ds = JsonlDataSource(local_dir, read_only=local_cfg.read_only)
-    else:
-        local_ds = ApiDataSource(
-            base_url=local_cfg.api_base_url,
-            uid=local_cfg.api_uid,
-            secret=local_cfg.api_secret,
-            debug=local_cfg.api_debug,
-            poll_max_retries=local_cfg.poll_max_retries,
-            poll_interval=local_cfg.poll_interval,
-            rate_limit_max_requests=local_cfg.api_rate_limit_max_requests,
-            rate_limit_window_seconds=local_cfg.api_rate_limit_window_seconds,
-        )
-        await local_ds.initialize()
+        return JsonlDataSource(dir_path, read_only=ds_config.read_only)

-    if isinstance(remote_cfg, JsonlDataSourceConfig):
-        remote_dir = Path(remote_cfg.jsonl_dir)
-        remote_dir = _prepare_remote_jsonl_dir(remote_dir)
-        remote_ds = JsonlDataSource(remote_dir, read_only=remote_cfg.read_only)
-    else:
-        remote_ds = ApiDataSource(
-            base_url=remote_cfg.api_base_url,
-            uid=remote_cfg.api_uid,
-            secret=remote_cfg.api_secret,
-            debug=remote_cfg.api_debug,
-            poll_max_retries=remote_cfg.poll_max_retries,
-            poll_interval=remote_cfg.poll_interval,
-            rate_limit_max_requests=remote_cfg.api_rate_limit_max_requests,
-            rate_limit_window_seconds=remote_cfg.api_rate_limit_window_seconds,
-        )
-        await remote_ds.initialize()
+    if isinstance(ds_config, DbDataSourceConfig):
+        return DbDataSource()

-    return local_ds, remote_ds
+    datasource = ApiDataSource(
+        base_url=ds_config.api_base_url,
+        uid=ds_config.api_uid,
+        secret=ds_config.api_secret,
+        debug=ds_config.api_debug,
+        poll_max_retries=ds_config.poll_max_retries,
+        poll_interval=ds_config.poll_interval,
+        rate_limit_max_requests=ds_config.api_rate_limit_max_requests,
+        rate_limit_window_seconds=ds_config.api_rate_limit_window_seconds,
+    )
+    await datasource.initialize()
+    return datasource
+
+
+def _register_handlers_for_endpoint(
+    *,
+    datasource,
+    ds_config: DataSourceConfig,
+    node_types: List[str],
+    target_project_ids: List[str],
+) -> None:
+    for node_type in node_types:
+        datasource.register_handler(
+            _create_handler(
+                node_type=node_type,
+                ds_config=ds_config,
+                datasource=datasource,
+                target_project_ids=target_project_ids,
+            )
+        )


 def _register_handlers(
@@ -170,44 +207,18 @@ def _register_handlers(
    remote_ds,
    all_types: List[str],
 ) -> None:
-    for node_type in all_types:
-        if isinstance(config.local_datasource, JsonlDataSourceConfig):
-            jsonl_handler_class = DomainRegistry.get_jsonl_handler(node_type)
-            if jsonl_handler_class is None:
-                raise RuntimeError(f"Missing JSONL handler for node_type: {node_type}")
-            local_ds.register_handler(jsonl_handler_class(local_ds))
-        else:
-            api_handler_class = DomainRegistry.get_api_handler(node_type)
-            if api_handler_class is None:
-                raise RuntimeError(f"Missing API handler for node_type: {node_type}")
-            local_ds.register_handler(
-                api_handler_class(
-                    **_api_handler_kwargs(
-                        node_type,
-                        config.local_datasource,
-                        config.target_project_ids,
-                    )
-                )
-            )
-
-        if isinstance(config.remote_datasource, JsonlDataSourceConfig):
-            jsonl_handler_class = DomainRegistry.get_jsonl_handler(node_type)
-            if jsonl_handler_class is None:
-                raise RuntimeError(f"Missing JSONL handler for node_type: {node_type}")
-            remote_ds.register_handler(jsonl_handler_class(remote_ds))
-        else:
-            api_handler_class = DomainRegistry.get_api_handler(node_type)
-            if api_handler_class is None:
-                raise RuntimeError(f"Missing API handler for node_type: {node_type}")
-            remote_ds.register_handler(
-                api_handler_class(
-                    **_api_handler_kwargs(
-                        node_type,
-                        config.remote_datasource,
-                        config.target_project_ids,
-                    )
-                )
-            )
+    _register_handlers_for_endpoint(
+        datasource=local_ds,
+        ds_config=config.local_datasource,
+        node_types=all_types,
+        target_project_ids=[str(pid) for pid in config.local_datasource.target_project_ids if str(pid)],
+    )
+    _register_handlers_for_endpoint(
+        datasource=remote_ds,
+        ds_config=config.remote_datasource,
+        node_types=all_types,
+        target_project_ids=[str(pid) for pid in config.remote_datasource.target_project_ids if str(pid)],
+    )


 def _resolve_strategy_map(config: PipelineRunConfig, node_types: List[str]) -> Dict[str, StrategyRuntimeConfig]:
@@ -260,7 +271,12 @@ def _build_strategies(
    return strategies


-async def create_pipeline_from_config(config: PipelineRunConfig) -> FullSyncPipeline:
+async def create_pipeline_from_config(
+    config: PipelineRunConfig,
+    *,
+    local_datasource_override=None,
+    remote_datasource_override=None,
+) -> FullSyncPipeline:
    if config.logging.initialize and not config.logging.file_path:
        config.logging.file_path = resolve_log_file_path(config.logging)

@@ -289,7 +305,16 @@ async def create_pipeline_from_config(config: PipelineRunConfig) -> FullSyncPipe
        persistence = PersistenceBackend(str(config.persist.db_path), backend=config.persist.backend)
        await persistence.initialize()

-        local_ds, remote_ds = await _create_datasources(config)
+        local_ds = await _create_datasource(
+            config.local_datasource,
+            datasource_override=local_datasource_override,
+            endpoint_name="local",
+        )
+        remote_ds = await _create_datasource(
+            config.remote_datasource,
+            datasource_override=remote_datasource_override,
+            endpoint_name="remote",
+        )
        local_cfg_with_target = config.local_datasource.model_copy(
            update={"target_project_ids": list(local_target_project_ids)},
            deep=True,
@@ -340,7 +365,7 @@ async def create_pipeline_from_config(config: PipelineRunConfig) -> FullSyncPipe

        return pipeline
    except Exception:
-        async def _safe_close(name: str, resource: ApiDataSource | JsonlDataSource | PersistenceBackend | None) -> None:
+        async def _safe_close(name: str, resource) -> None:
            if resource is None:
                return
            try:
@@ -359,6 +384,8 @@ async def run_pipeline_from_config(
    *,
    title: str = "sync_state_machine pipeline",
    print_summary: bool = True,
+    local_datasource_override=None,
+    remote_datasource_override=None,
 ) -> Dict[str, Any]:
    if config.logging.initialize and not config.logging.file_path:
        config.logging.file_path = resolve_log_file_path(config.logging)
@@ -401,7 +428,11 @@ async def run_pipeline_from_config(
            print(f"🔧 Resolved Config:\n{resolved_cfg_text}")
            print("=" * 80)

-        pipeline = await create_pipeline_from_config(config)
+        pipeline = await create_pipeline_from_config(
+            config,
+            local_datasource_override=local_datasource_override,
+            remote_datasource_override=remote_datasource_override,
+        )
        stats = await pipeline.run()

        if print_summary: