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

This commit is contained in:
strepsiades
2026-03-10 12:22:54 +08:00
parent 515c3d9388
commit a638e4203b
35 changed files with 1314 additions and 199 deletions
+114 -83
View File
@@ -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: