整理代码

sync_state_machine/datasource/datasource.py

@@ -163,6 +163,11 @@ class BaseDataSource(ABC):
         """Datasource-specific hook after a FAILED result is applied to node."""
         return
 
+    def _prepare_handler_for_registration(self, handler: "NodeHandler") -> None:
+        """Allow subclasses to inject datasource-specific dependencies before registration."""
+        if self._collection is not None:
+            handler.set_collection(self._collection)
+
     async def _handle_success_result(self, node: "SyncNode", result: "TaskResult") -> None:
         previous_data_id = node.data_id
         ok = self._apply_sync_execute_state(
@@ -371,6 +376,7 @@ class BaseDataSource(ABC):
         示例：
             datasource.add_handler(ContractNodeHandler(api_client))
         """
+        self._prepare_handler_for_registration(handler)
         self._handlers[handler.node_type] = handler
     
     def get_handler(self, node_type: str) -> "NodeHandler":