strepsiades
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
"""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(),
|
|
) |