整理配置和代码入口，增加部分测试。
This commit is contained in:
strepsiades
@@ -14,6 +14,7 @@ from . import domain as _domain
|
||||
from .common.registry import DomainRegistry
|
||||
from .config import (
|
||||
ApiDataSourceConfig,
|
||||
DbDataSourceConfig,
|
||||
JsonlDataSourceConfig,
|
||||
PipelineRunConfig,
|
||||
build_config,
|
||||
@@ -21,7 +22,7 @@ from .config import (
|
||||
load_named_preset_overrides,
|
||||
load_overrides_from_file,
|
||||
)
|
||||
from .datasource import ApiClient, ApiDataSource, BaseApiHandler, BaseJsonlHandler, JsonlDataSource
|
||||
from .datasource import ApiClient, ApiDataSource, BaseApiHandler, BaseDbHandler, BaseJsonlHandler, DbDataSource, JsonlDataSource
|
||||
from .pipeline import FullSyncPipeline, create_pipeline_from_config, run_pipeline_from_config
|
||||
|
||||
try:
|
||||
@@ -41,6 +42,7 @@ __all__ = [
|
||||
"DomainRegistry",
|
||||
"PipelineRunConfig",
|
||||
"ApiDataSourceConfig",
|
||||
"DbDataSourceConfig",
|
||||
"JsonlDataSourceConfig",
|
||||
"build_config",
|
||||
"build_default_config",
|
||||
@@ -48,8 +50,10 @@ __all__ = [
|
||||
"load_named_preset_overrides",
|
||||
"ApiClient",
|
||||
"ApiDataSource",
|
||||
"DbDataSource",
|
||||
"JsonlDataSource",
|
||||
"BaseApiHandler",
|
||||
"BaseDbHandler",
|
||||
"BaseJsonlHandler",
|
||||
"FullSyncPipeline",
|
||||
"create_pipeline_from_config",
|
||||
|
||||
@@ -26,7 +26,7 @@ Domain 注册中心
|
||||
# DataSource 查询 Handler
|
||||
handler_class = DomainRegistry.get_jsonl_handler("contract")
|
||||
"""
|
||||
from typing import Dict, Type, Optional, Any
|
||||
from typing import Dict, Type, Optional, Any, Literal
|
||||
from dataclasses import dataclass, field
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -64,6 +64,12 @@ class DomainRegistry:
|
||||
"""
|
||||
|
||||
_registry: Dict[str, DomainRegistration] = {}
|
||||
|
||||
_HANDLER_ATTRS: Dict[str, str] = {
|
||||
"jsonl": "jsonl_handler_class",
|
||||
"api": "api_handler_class",
|
||||
"db": "db_handler_class",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def register(
|
||||
@@ -143,36 +149,30 @@ class DomainRegistry:
|
||||
return reg.strategy_class
|
||||
|
||||
@classmethod
|
||||
def get_jsonl_handler(cls, node_type: str) -> Optional[Type[Any]]:
|
||||
"""
|
||||
获取 JSONL Handler 类（用于 DataSource）
|
||||
|
||||
Args:
|
||||
node_type: 节点类型
|
||||
|
||||
Returns:
|
||||
Handler 类，如果未注册则返回 None
|
||||
"""
|
||||
def get_handler(
|
||||
cls,
|
||||
node_type: str,
|
||||
datasource_type: Literal["jsonl", "api", "db"],
|
||||
) -> Optional[Type[Any]]:
|
||||
"""按 datasource 类型获取 Handler 类。"""
|
||||
reg = cls._registry.get(node_type)
|
||||
if not reg:
|
||||
raise ValueError(f"Unknown node_type: '{node_type}'")
|
||||
return reg.jsonl_handler_class
|
||||
|
||||
attr_name = cls._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 getattr(reg, attr_name)
|
||||
|
||||
@classmethod
|
||||
def get_api_handler(cls, node_type: str) -> Optional[Type[Any]]:
|
||||
"""获取 API Handler 类"""
|
||||
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}'")
|
||||
return reg.api_handler_class
|
||||
|
||||
@classmethod
|
||||
def get_db_handler(cls, node_type: str) -> Optional[Type[Any]]:
|
||||
"""获取 DB Handler 类"""
|
||||
reg = cls._registry.get(node_type)
|
||||
if not reg:
|
||||
raise ValueError(f"Unknown node_type: '{node_type}'")
|
||||
return reg.db_handler_class
|
||||
reg.db_handler_class = db_handler_class
|
||||
|
||||
@classmethod
|
||||
def get_schema(cls, node_type: str) -> Type[BaseModel]:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from .datasource_config import ApiDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .datasource_config import ApiDataSourceConfig, DbDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .persist_config import PersistConfig
|
||||
from .logging_config import LoggingConfig, resolve_log_file_path
|
||||
from .run_presets import (
|
||||
@@ -26,6 +26,7 @@ from .builder import (
|
||||
|
||||
__all__ = [
|
||||
"ApiDataSourceConfig",
|
||||
"DbDataSourceConfig",
|
||||
"JsonlDataSourceConfig",
|
||||
"PipelineRunConfig",
|
||||
"DataSourceConfig",
|
||||
|
||||
@@ -15,6 +15,14 @@ class JsonlDataSourceConfig(BaseModel):
|
||||
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class DbDataSourceConfig(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", validate_assignment=True)
|
||||
|
||||
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)
|
||||
|
||||
@@ -31,4 +39,7 @@ class ApiDataSourceConfig(BaseModel):
|
||||
handler_configs: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
DataSourceConfig = Annotated[Union[JsonlDataSourceConfig, ApiDataSourceConfig], Field(discriminator="type")]
|
||||
DataSourceConfig = Annotated[
|
||||
Union[JsonlDataSourceConfig, DbDataSourceConfig, ApiDataSourceConfig],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Prefer importing from `sync_state_machine.config` directly.
|
||||
"""
|
||||
|
||||
from .datasource_config import ApiDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .datasource_config import ApiDataSourceConfig, DbDataSourceConfig, JsonlDataSourceConfig, DataSourceConfig
|
||||
from .persist_config import PersistConfig
|
||||
from .logging_config import LoggingConfig
|
||||
from .strategy_config import (
|
||||
@@ -24,6 +24,7 @@ from .builder import (
|
||||
|
||||
__all__ = [
|
||||
"ApiDataSourceConfig",
|
||||
"DbDataSourceConfig",
|
||||
"JsonlDataSourceConfig",
|
||||
"DataSourceConfig",
|
||||
"PersistConfig",
|
||||
|
||||
@@ -17,6 +17,9 @@ from .task_result import TaskResult, TaskStatus
|
||||
# API DataSource
|
||||
from .api import ApiClient, ApiDataSource, BaseApiHandler
|
||||
|
||||
# DB DataSource
|
||||
from .db import DbDataSource, BaseDbHandler
|
||||
|
||||
# JSONL DataSource
|
||||
from .jsonl import JsonlDataSource, BaseJsonlHandler
|
||||
|
||||
@@ -34,6 +37,10 @@ __all__ = [
|
||||
"ApiClient",
|
||||
"ApiDataSource",
|
||||
"BaseApiHandler",
|
||||
|
||||
# DB DataSource
|
||||
"DbDataSource",
|
||||
"BaseDbHandler",
|
||||
|
||||
# JSONL DataSource
|
||||
"JsonlDataSource",
|
||||
|
||||
@@ -44,6 +44,7 @@ class BaseApiHandler(NodeHandler[T]):
|
||||
self.api_client: 'ApiClient' = None # type: ignore # 由 DataSource 注入
|
||||
self.poll_mode: str = "async" # async | sync
|
||||
self._collection: 'DataCollection' = None # type: ignore # 由 DataSource 注入
|
||||
self.handler_config: Dict[str, Any] = {}
|
||||
# 子类在 __init__ 中设置此列表以声明 create/update 中涉及的 Pydantic schema 类；
|
||||
# post-check 将只比较这些 schema 覆盖的字段，过滤后端只读的派生字段。
|
||||
self.update_schemas: List[type] = []
|
||||
@@ -63,6 +64,13 @@ class BaseApiHandler(NodeHandler[T]):
|
||||
def set_collection(self, collection: 'DataCollection') -> None:
|
||||
"""设置 Collection（由 DataSource 调用）"""
|
||||
self._collection = collection
|
||||
|
||||
def set_handler_config(self, config: Dict[str, Any]) -> None:
|
||||
self.handler_config = dict(config)
|
||||
|
||||
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
|
||||
"""默认忽略项目过滤列表，具体子类按需覆盖。"""
|
||||
return
|
||||
|
||||
@property
|
||||
def node_type(self) -> str:
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
"""DB DataSource 模块。"""
|
||||
|
||||
from .datasource import DbDataSource
|
||||
from .handler import BaseDbHandler
|
||||
|
||||
__all__ = ["DbDataSource", "BaseDbHandler"]
|
||||
@@ -0,0 +1,22 @@
|
||||
"""DbDataSource - 数据库数据源协调层。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ..datasource import BaseDataSource
|
||||
|
||||
|
||||
class DbDataSource(BaseDataSource):
|
||||
"""基于数据库 Handler 的数据源实现。"""
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
def _on_node_loaded(self, handler, node) -> None:
|
||||
node.append_log(
|
||||
f"DB链路[load] node_type={node.node_type} data_id={node.data_id or '<empty>'}"
|
||||
)
|
||||
|
||||
def _on_in_progress_result(self, handler, node, result) -> None:
|
||||
node.append_log(
|
||||
f"DB链路[submit] action={node.action.value} data_id={node.data_id or '<empty>'}"
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
"""BaseDbHandler - DB Handler 基类。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional, Set, TypeVar, TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ...common.type_safety import get_pydantic_model_fields
|
||||
from ..handler import BaseNodeHandler
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...common.sync_node import SyncNode
|
||||
from .datasource import DbDataSource
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
class BaseDbHandler(BaseNodeHandler[T]):
|
||||
"""数据库 Handler 基类。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
datasource: "DbDataSource",
|
||||
node_type: str,
|
||||
node_class: type["SyncNode[T]"],
|
||||
schema: type[T],
|
||||
):
|
||||
super().__init__(node_type, node_class, schema)
|
||||
self.datasource = datasource
|
||||
self.update_schemas: List[type] = []
|
||||
self._target_project_ids: Optional[Set[str]] = None
|
||||
|
||||
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
|
||||
if target_project_ids:
|
||||
self._target_project_ids = {str(item) for item in target_project_ids if str(item)}
|
||||
else:
|
||||
self._target_project_ids = None
|
||||
|
||||
def get_update_fields(self, *, exclude: frozenset = frozenset({"id"})) -> List[str]:
|
||||
fields: set[str] = set()
|
||||
for schema in self.update_schemas:
|
||||
for field_name in get_pydantic_model_fields(schema):
|
||||
if field_name not in exclude:
|
||||
fields.add(field_name)
|
||||
return sorted(fields)
|
||||
|
||||
def _create_basic_node(self, data: Dict[str, Any], depend_ids: Optional[List[str]] = None) -> "SyncNode[T]":
|
||||
data_id = self.extract_id(data)
|
||||
if self._collection is None:
|
||||
raise RuntimeError("Collection not set. Call set_collection() before creating nodes.")
|
||||
|
||||
if data_id:
|
||||
existing_node = self._collection.get_by_data_id(self.node_type, data_id)
|
||||
if existing_node is not None:
|
||||
existing_node.depend_ids = list(depend_ids or [])
|
||||
existing_node.set_data(data.copy())
|
||||
existing_node.set_origin_data(data.copy())
|
||||
return existing_node
|
||||
|
||||
node_id = self._collection.generate_node_id()
|
||||
return self.node_class(
|
||||
node_id=node_id,
|
||||
data_id=data_id,
|
||||
data=data.copy(),
|
||||
depend_ids=depend_ids or [],
|
||||
origin_data=data.copy(),
|
||||
)
|
||||
@@ -70,6 +70,10 @@ class NodeHandler(Protocol[T]):
|
||||
"""
|
||||
...
|
||||
|
||||
def set_handler_config(self, config: Dict[str, Any]) -> None:
|
||||
"""注入 handler 级别配置。默认可忽略。"""
|
||||
...
|
||||
|
||||
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
|
||||
"""注入项目过滤列表。默认可忽略。"""
|
||||
...
|
||||
@@ -345,6 +349,7 @@ class BaseNodeHandler(ABC, Generic[T]):
|
||||
self._node_class = node_class
|
||||
self._schema = schema
|
||||
self._collection: Optional["DataCollection"] = None
|
||||
self.handler_config: Dict[str, Any] = {}
|
||||
|
||||
def set_collection(self, collection: "DataCollection") -> None:
|
||||
"""
|
||||
@@ -359,6 +364,10 @@ class BaseNodeHandler(ABC, Generic[T]):
|
||||
"""默认忽略项目过滤能力。"""
|
||||
return
|
||||
|
||||
def set_handler_config(self, config: Dict[str, Any]) -> None:
|
||||
"""默认保存 handler 级别配置，供子类按需读取。"""
|
||||
self.handler_config = dict(config)
|
||||
|
||||
def set_api_client(self, client: "ApiClient") -> None:
|
||||
"""默认忽略 API 客户端注入。"""
|
||||
return
|
||||
@@ -491,6 +500,11 @@ class CompatibleNodeHandler(Generic[T]):
|
||||
def set_collection(self, collection: "DataCollection") -> None:
|
||||
self._handler.set_collection(collection)
|
||||
|
||||
def set_handler_config(self, config: Dict[str, Any]) -> None:
|
||||
method = getattr(self._handler, "set_handler_config", None)
|
||||
if callable(method):
|
||||
method(config)
|
||||
|
||||
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
|
||||
method = getattr(self._handler, "set_target_project_ids", None)
|
||||
if callable(method):
|
||||
|
||||
@@ -188,7 +188,7 @@ class MaterialDetailApiHandler(BaseApiHandler):
|
||||
)
|
||||
# 整组成功
|
||||
for node, _ in group_items:
|
||||
results.append(TaskResult.success(node_id=node.node_id))
|
||||
results.append(TaskResult.in_progress(node_id=node.node_id, task_id=push_id))
|
||||
except Exception as e:
|
||||
# 整组失败
|
||||
error_msg = str(e)
|
||||
|
||||
@@ -18,6 +18,7 @@ PUT (6个更新接口):
|
||||
不支持: CREATE, DELETE
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Dict, Any
|
||||
from ...datasource.api.handler import BaseApiHandler
|
||||
from ...datasource.task_result import TaskResult
|
||||
@@ -35,6 +36,35 @@ from schemas.project.project_base import (
|
||||
from schemas.project.project import ProjectInfoResponse
|
||||
|
||||
|
||||
_PROJECT_EXTENSION_FIELD_NAMES = {
|
||||
"implementation_estimate",
|
||||
"static_total_investment",
|
||||
"completed_investment",
|
||||
"total_contract_quantity",
|
||||
"completed_settlement_quantity",
|
||||
}
|
||||
|
||||
|
||||
def _merge_project_all_info_fields(
|
||||
project_data: Dict[str, Any],
|
||||
all_info: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""把 all_info 中项目扩展字段拍平到 project 主数据上。"""
|
||||
merged = dict(project_data)
|
||||
|
||||
for key, value in all_info.items():
|
||||
if not key.endswith("_extension") or not isinstance(value, dict):
|
||||
continue
|
||||
for field_name in _PROJECT_EXTENSION_FIELD_NAMES:
|
||||
if field_name not in value:
|
||||
continue
|
||||
current_value = merged.get(field_name)
|
||||
if current_value in (None, ""):
|
||||
merged[field_name] = value.get(field_name)
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
class ProjectApiHandler(BaseApiHandler):
|
||||
"""项目 API Handler"""
|
||||
|
||||
@@ -136,12 +166,15 @@ class ProjectApiHandler(BaseApiHandler):
|
||||
|
||||
# 处理所有项目
|
||||
for project_data in projects:
|
||||
node = self._create_node(project_data.model_dump())
|
||||
# 设置 context
|
||||
node.context["project_id"] = project_data.id
|
||||
|
||||
# 预加载并缓存 all_info，同时存储到 node.context
|
||||
all_info = await self.get_all_info(self.api_client, project_data.id)
|
||||
normalized_project_data = _merge_project_all_info_fields(
|
||||
project_data.model_dump(),
|
||||
all_info,
|
||||
)
|
||||
node = self._create_node(normalized_project_data)
|
||||
# 设置 context
|
||||
node.context["project_id"] = project_data.id
|
||||
node.context["all_info"] = all_info
|
||||
|
||||
nodes.append(node)
|
||||
@@ -217,10 +250,12 @@ class ProjectApiHandler(BaseApiHandler):
|
||||
|
||||
# 执行更新（差异已由 _get_schema_diff 打印）
|
||||
try:
|
||||
push_id = self._generate_push_id()
|
||||
# project 需要使用 _filter_update_data 返回的完整 schema 数据
|
||||
full_update_data = self._filter_update_data(data, original_data, schema)
|
||||
await api_func(self.api_client, full_update_data, push_id, biz_id)
|
||||
await self._submit_update_and_wait(
|
||||
api_func=api_func,
|
||||
data=full_update_data,
|
||||
biz_id=biz_id,
|
||||
)
|
||||
node_updated = True
|
||||
except Exception as e:
|
||||
error = str(e)
|
||||
@@ -249,13 +284,39 @@ class ProjectApiHandler(BaseApiHandler):
|
||||
async def poll_tasks(self, task_ids: List[str]) -> Dict[str, TaskResult]:
|
||||
"""轮询异步任务状态"""
|
||||
return await super().poll_tasks(task_ids)
|
||||
|
||||
async def _submit_update_and_wait(self, *, api_func, data: Dict[str, Any], biz_id: str) -> None:
|
||||
push_id = self._generate_push_id()
|
||||
await api_func(self.api_client, data, push_id, biz_id)
|
||||
|
||||
max_attempts = 20
|
||||
poll_interval = 0.2
|
||||
last_result: TaskResult | None = None
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
polled = await super().poll_tasks([push_id])
|
||||
result = polled.get(push_id)
|
||||
last_result = result
|
||||
|
||||
if result is None:
|
||||
raise RuntimeError(f"Project update poll missing result: push_id={push_id}")
|
||||
if result.status == result.status.SUCCESS:
|
||||
return
|
||||
if result.status == result.status.FAILED:
|
||||
raise RuntimeError(result.error or f"Project update failed: push_id={push_id}")
|
||||
if attempt < max_attempts - 1:
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
raise RuntimeError(
|
||||
f"Project update poll timeout: push_id={push_id}, last_status={getattr(last_result, 'status', None)}"
|
||||
)
|
||||
|
||||
def _create_node(self, data: Dict[str, Any]) -> SyncNode:
|
||||
node = self._create_basic_node(data, depend_ids=[])
|
||||
|
||||
|
||||
# 设置 context
|
||||
node.context["project_id"] = node.data_id
|
||||
|
||||
|
||||
return node
|
||||
|
||||
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
"""Project domain registration"""
|
||||
from ...common.registry import DomainRegistry
|
||||
from schemas.project.project_base import ProjectResponseBase
|
||||
from .sync_node import ProjectSyncNode
|
||||
from .sync_node import ProjectSyncNode, ProjectSyncSchema
|
||||
from .sync_strategy import ProjectSyncStrategy
|
||||
from .jsonl_handler import ProjectJsonlHandler
|
||||
from .api_handler import ProjectApiHandler
|
||||
@@ -9,7 +8,7 @@ from .api_handler import ProjectApiHandler
|
||||
# 注册 project domain
|
||||
DomainRegistry.register(
|
||||
node_type="project",
|
||||
schema=ProjectResponseBase,
|
||||
schema=ProjectSyncSchema,
|
||||
node_class=ProjectSyncNode,
|
||||
strategy_class=ProjectSyncStrategy,
|
||||
jsonl_handler_class=ProjectJsonlHandler,
|
||||
|
||||
@@ -1,11 +1,23 @@
|
||||
from pydantic import Field
|
||||
|
||||
from ...common.sync_node import SyncNode
|
||||
from schemas.project.project_base import ProjectResponseBase
|
||||
|
||||
|
||||
class ProjectSyncNode(SyncNode[ProjectResponseBase]):
|
||||
class ProjectSyncSchema(ProjectResponseBase):
|
||||
"""项目同步内部 schema，补充 all_info 扩展字段。"""
|
||||
|
||||
implementation_estimate: float | None = Field(None, description="执行概算")
|
||||
static_total_investment: float | None = Field(None, description="静态总投资")
|
||||
completed_investment: float | None = Field(None, description="已完成投资")
|
||||
total_contract_quantity: float | None = Field(None, description="合同总量")
|
||||
completed_settlement_quantity: float | None = Field(None, description="已结算量")
|
||||
|
||||
|
||||
class ProjectSyncNode(SyncNode[ProjectSyncSchema]):
|
||||
"""项目同步节点"""
|
||||
node_type = "project"
|
||||
schema = ProjectResponseBase
|
||||
schema = ProjectSyncSchema
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@@ -17,6 +17,7 @@ PUT:
|
||||
不支持: CREATE, DELETE
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Dict, Any, Optional, Tuple, Set
|
||||
from pydantic import ValidationError
|
||||
|
||||
@@ -270,6 +271,8 @@ class ProjectDetailApiHandler(BaseApiHandler):
|
||||
|
||||
try:
|
||||
resolved_data_id = self._resolve_data_id_from_all_info(all_info, project_id=project_id, key=key)
|
||||
if not resolved_data_id:
|
||||
resolved_data_id = await self._resolve_data_id_with_retry(project_id=project_id, key=key)
|
||||
except Exception as exc:
|
||||
self._pending_push_context.pop(task_id, None)
|
||||
results[task_id] = TaskResult.failed(error=f"Poll resolve data_id failed: {exc}", push_id=task_id)
|
||||
@@ -342,6 +345,9 @@ class ProjectDetailApiHandler(BaseApiHandler):
|
||||
if isinstance(all_info, dict):
|
||||
return all_info
|
||||
|
||||
if force_refresh:
|
||||
ProjectApiHandler._all_info_cache.pop(project_id, None)
|
||||
|
||||
all_info = await ProjectApiHandler.get_all_info(self.api_client, project_id)
|
||||
if project_node is not None:
|
||||
project_node.context["all_info"] = all_info
|
||||
@@ -370,6 +376,20 @@ class ProjectDetailApiHandler(BaseApiHandler):
|
||||
return matched[0]
|
||||
return None
|
||||
|
||||
async def _resolve_data_id_with_retry(self, *, project_id: str, key: str) -> Optional[str]:
|
||||
attempts = 6
|
||||
delay = 0.2
|
||||
|
||||
for attempt in range(attempts):
|
||||
all_info = await self._get_project_all_info(project_id, force_refresh=True)
|
||||
resolved = self._resolve_data_id_from_all_info(all_info, project_id=project_id, key=key)
|
||||
if resolved:
|
||||
return resolved
|
||||
if attempt < attempts - 1:
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
return None
|
||||
|
||||
async def _resolve_data_id_from_project_all_info(self, project_id: str, key: str) -> Optional[str]:
|
||||
all_info = await self._get_project_all_info(project_id)
|
||||
return self._resolve_data_id_from_all_info(all_info, project_id=project_id, key=key)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user