229 lines
7.2 KiB
Python
229 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional, Sequence
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from sync_state_machine.common.registry import DomainRegistry
|
|
from sync_state_machine.common.sync_node import SyncNode
|
|
from sync_state_machine.datasource.datasource import BaseDataSource
|
|
from sync_state_machine.datasource.handler import ValidationResult
|
|
from sync_state_machine.datasource.task_result import TaskResult
|
|
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
|
|
|
|
|
|
class TestProjectSchema(BaseModel):
|
|
__test__ = False
|
|
id: str
|
|
name: str
|
|
code: str
|
|
|
|
|
|
class TestContractSchema(BaseModel):
|
|
__test__ = False
|
|
id: str
|
|
name: str
|
|
code: str
|
|
project_ref: str
|
|
|
|
|
|
class TestProjectNode(SyncNode[TestProjectSchema]):
|
|
__test__ = False
|
|
node_type = "test_project"
|
|
schema = TestProjectSchema
|
|
|
|
|
|
class TestContractNode(SyncNode[TestContractSchema]):
|
|
__test__ = False
|
|
node_type = "test_contract"
|
|
schema = TestContractSchema
|
|
|
|
|
|
class TestProjectStrategy(DefaultSyncStrategy[TestProjectSchema]):
|
|
__test__ = False
|
|
schema = TestProjectSchema
|
|
|
|
|
|
class TestContractStrategy(DefaultSyncStrategy[TestContractSchema]):
|
|
__test__ = False
|
|
schema = TestContractSchema
|
|
|
|
|
|
def register_test_domain() -> None:
|
|
if not DomainRegistry.is_registered("test_project"):
|
|
DomainRegistry.register(
|
|
node_type="test_project",
|
|
schema=TestProjectSchema,
|
|
node_class=TestProjectNode,
|
|
strategy_class=TestProjectStrategy,
|
|
)
|
|
if not DomainRegistry.is_registered("test_contract"):
|
|
DomainRegistry.register(
|
|
node_type="test_contract",
|
|
schema=TestContractSchema,
|
|
node_class=TestContractNode,
|
|
strategy_class=TestContractStrategy,
|
|
)
|
|
|
|
|
|
def build_project_node(
|
|
node_id: str,
|
|
data_id: str,
|
|
*,
|
|
name: str,
|
|
code: str,
|
|
data_present: bool = True,
|
|
):
|
|
data = {"id": data_id, "name": name, "code": code} if data_present else None
|
|
return TestProjectNode(node_id=node_id, data_id=data_id, data=data)
|
|
|
|
|
|
def build_contract_node(
|
|
node_id: str,
|
|
data_id: str,
|
|
*,
|
|
name: str,
|
|
code: str,
|
|
project_ref: str,
|
|
data_present: bool = True,
|
|
):
|
|
data = (
|
|
{
|
|
"id": data_id,
|
|
"name": name,
|
|
"code": code,
|
|
"project_ref": project_ref,
|
|
}
|
|
if data_present
|
|
else None
|
|
)
|
|
return TestContractNode(node_id=node_id, data_id=data_id, data=data)
|
|
|
|
|
|
class InMemoryNodeHandler:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
node_type: str,
|
|
node_class,
|
|
schema,
|
|
load_nodes: Sequence[SyncNode],
|
|
create_mode_by_code: Optional[Dict[str, str]] = None,
|
|
update_mode_by_code: Optional[Dict[str, str]] = None,
|
|
):
|
|
self._node_type = node_type
|
|
self._node_class = node_class
|
|
self._schema = schema
|
|
self._load_nodes = load_nodes
|
|
self._collection = None
|
|
self._create_mode_by_code = create_mode_by_code or {}
|
|
self._update_mode_by_code = update_mode_by_code or {}
|
|
self._poll_sequences: Dict[str, List[str]] = defaultdict(list)
|
|
|
|
@property
|
|
def node_type(self) -> str:
|
|
return self._node_type
|
|
|
|
@property
|
|
def node_class(self):
|
|
return self._node_class
|
|
|
|
@property
|
|
def schema(self):
|
|
return self._schema
|
|
|
|
def set_collection(self, collection):
|
|
self._collection = collection
|
|
|
|
def extract_id(self, data: Dict[str, str]) -> str:
|
|
return str(data.get("id", ""))
|
|
|
|
async def validate(self, data: Dict[str, str]) -> ValidationResult:
|
|
try:
|
|
self._schema.model_validate(data)
|
|
return ValidationResult(is_valid=True)
|
|
except Exception as exc:
|
|
return ValidationResult(is_valid=False, errors=[str(exc)])
|
|
|
|
def _create_node(self, data: Dict[str, str]):
|
|
data_id = self.extract_id(data)
|
|
return self._node_class(
|
|
node_id=f"{self._node_type}-{data_id}",
|
|
data_id=data_id,
|
|
data=data,
|
|
)
|
|
|
|
async def load(self):
|
|
return [
|
|
self._node_class(
|
|
node_id=node.node_id,
|
|
data_id=node.data_id,
|
|
data=node.get_data(),
|
|
origin_data=node.get_data(),
|
|
action=node.action,
|
|
status=node.status,
|
|
binding_status=node.binding_status,
|
|
depend_ids=list(node.depend_ids),
|
|
context=dict(node.context),
|
|
)
|
|
for node in self._load_nodes
|
|
]
|
|
|
|
async def create_all(self, nodes):
|
|
results = []
|
|
for node in nodes:
|
|
data = node.get_data() or {}
|
|
mode = self._create_mode_by_code.get(data.get("code", ""), "success")
|
|
results.append(self._result_for_mode(mode, node.node_id, node_code=data.get("code", ""), phase="create"))
|
|
return results
|
|
|
|
async def update_all(self, nodes):
|
|
results = []
|
|
for node in nodes:
|
|
data = node.get_data() or {}
|
|
mode = self._update_mode_by_code.get(data.get("code", ""), "success")
|
|
results.append(self._result_for_mode(mode, node.node_id, node_code=data.get("code", ""), phase="update"))
|
|
return results
|
|
|
|
async def delete_all(self, nodes):
|
|
return [TaskResult.success(node_id=node.node_id, data_id=node.data_id) for node in nodes]
|
|
|
|
async def poll_tasks(self, task_ids: List[str]) -> Dict[str, TaskResult]:
|
|
out: Dict[str, TaskResult] = {}
|
|
for task_id in task_ids:
|
|
sequence = self._poll_sequences.get(task_id) or ["success"]
|
|
mode = sequence.pop(0)
|
|
if not sequence:
|
|
sequence = [mode]
|
|
self._poll_sequences[task_id] = sequence
|
|
node_id = task_id.split("::", 1)[1]
|
|
if mode == "success":
|
|
out[task_id] = TaskResult.success(node_id=node_id, data_id=f"polled-{node_id}")
|
|
elif mode == "failed":
|
|
out[task_id] = TaskResult.failed(node_id=node_id, error="poll failed")
|
|
else:
|
|
out[task_id] = TaskResult.in_progress(node_id=node_id, task_id=task_id)
|
|
return out
|
|
|
|
def _result_for_mode(self, mode: str, node_id: str, *, node_code: str, phase: str) -> TaskResult:
|
|
if mode == "success":
|
|
return TaskResult.success(node_id=node_id, data_id=f"{phase}-ok-{node_id}")
|
|
if mode == "failed":
|
|
return TaskResult.failed(node_id=node_id, error=f"{phase} failed for {node_code}")
|
|
if mode == "skipped":
|
|
return TaskResult.skipped(node_id=node_id, reason=f"{phase} skipped for {node_code}")
|
|
if mode == "in_progress_success":
|
|
task_id = f"task::{node_id}"
|
|
self._poll_sequences[task_id] = ["success"]
|
|
return TaskResult.in_progress(node_id=node_id, task_id=task_id)
|
|
if mode == "in_progress_failed":
|
|
task_id = f"task::{node_id}"
|
|
self._poll_sequences[task_id] = ["failed"]
|
|
return TaskResult.in_progress(node_id=node_id, task_id=task_id)
|
|
raise ValueError(f"Unsupported mode: {mode}")
|
|
|
|
|
|
class InMemoryDataSource(BaseDataSource):
|
|
pass
|