进一步清理代码和注册过程

This commit is contained in:
strepsiades
2026-03-18 17:23:31 +08:00
parent 32fa374e86
commit f9f175ee79
7 changed files with 42 additions and 49 deletions
@@ -20,9 +20,9 @@ from sync_state_machine.pipeline.factory import create_pipeline_from_config
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
class EndpointInitDbConfig(BaseDataSourceConfig):
class EndpointInitConfig(BaseDataSourceConfig):
__test__ = False
type: str = "endpoint_init_db"
type: str = "endpoint_init_memory"
class EndpointInitSchema(BaseModel):
@@ -42,7 +42,7 @@ class EndpointInitStrategy(DefaultSyncStrategy[EndpointInitSchema]):
schema = EndpointInitSchema
class EndpointInitDbHandler(BaseNodeHandler[EndpointInitSchema]):
class EndpointInitHandler(BaseNodeHandler[EndpointInitSchema]):
__test__ = False
def __init__(self, datasource: BaseDataSource):
@@ -69,8 +69,7 @@ def _ensure_registered() -> None:
reg.schema = EndpointInitSchema
reg.node_class = EndpointInitNode
reg.strategy_class = EndpointInitStrategy
reg.db_handler_class = EndpointInitDbHandler
reg.handler_classes["endpoint_init_db"] = EndpointInitDbHandler
reg.handler_classes["endpoint_init_memory"] = EndpointInitHandler
return
DomainRegistry.register(
@@ -78,15 +77,14 @@ def _ensure_registered() -> None:
schema=EndpointInitSchema,
node_class=EndpointInitNode,
strategy_class=EndpointInitStrategy,
db_handler_class=EndpointInitDbHandler,
)
DomainRegistry.register_handler(node_type, "endpoint_init_db", EndpointInitDbHandler)
DomainRegistry.register_handler(node_type, "endpoint_init_memory", EndpointInitHandler)
def _register_endpoint_init_datasource_type() -> None:
register_datasource_type(
"endpoint_init_db",
config_model=EndpointInitDbConfig,
"endpoint_init_memory",
config_model=EndpointInitConfig,
factory=lambda config, endpoint_name: BaseDataSource(),
replace=True,
)
@@ -102,13 +100,13 @@ async def test_factory_initializes_same_datasource_type_with_distinct_handler_co
config = PipelineRunConfig(
node_types=["factory_endpoint_demo"],
local_datasource=EndpointInitDbConfig(
type="endpoint_init_db",
local_datasource=EndpointInitConfig(
type="endpoint_init_memory",
handler_configs={"factory_endpoint_demo": {"side": "local", "batch_size": 10}},
target_project_ids=["project_local"],
),
remote_datasource=EndpointInitDbConfig(
type="endpoint_init_db",
remote_datasource=EndpointInitConfig(
type="endpoint_init_memory",
handler_configs={"factory_endpoint_demo": {"side": "remote", "batch_size": 20}},
target_project_ids=["project_remote"],
),
@@ -122,8 +120,8 @@ async def test_factory_initializes_same_datasource_type_with_distinct_handler_co
local_handler = pipeline.local_datasource.get_handler("factory_endpoint_demo")
remote_handler = pipeline.remote_datasource.get_handler("factory_endpoint_demo")
assert isinstance(local_handler, EndpointInitDbHandler)
assert isinstance(remote_handler, EndpointInitDbHandler)
assert isinstance(local_handler, EndpointInitHandler)
assert isinstance(remote_handler, EndpointInitHandler)
assert local_handler.handler_config == {"side": "local", "batch_size": 10}
assert remote_handler.handler_config == {"side": "remote", "batch_size": 20}
finally: