159class UniNode(BaseModel, metaclass=UniModelMeta):
160    """
161    Base class for graph node models.
162
163    Subclass this to define your node types. Each UniNode subclass
164    represents a vertex label in the graph database.
165
166    Attributes:
167        __label__: The vertex label name. Defaults to the class name.
168        __relationships__: Dictionary of relationship configurations.
169
170    Private Attributes:
171        _vid: The vertex ID assigned by the database.
172        _uid: The unique identifier (content-addressed hash).
173        _session: Reference to the owning session.
174        _dirty: Set of modified field names.
175
176    Example:
177        >>> class Person(UniNode):
178        ...     __label__ = "Person"
179        ...
180        ...     name: str
181        ...     age: int | None = None
182        ...     email: str = Field(unique=True)
183        ...
184        ...     friends: list["Person"] = Relationship("FRIEND_OF", direction="both")
185    """
186
187    model_config = ConfigDict(
188        # Allow extra fields for future extensibility
189        extra="forbid",
190        # Validate on assignment for dirty tracking
191        validate_assignment=True,
192        # Allow arbitrary types (for Vector, etc.)
193        arbitrary_types_allowed=True,
194        # Use enum values
195        use_enum_values=True,
196    )
197
198    # Class-level configuration
199    __label__: ClassVar[str] = ""
200    __relationships__: ClassVar[dict[str, RelationshipConfig]] = {}
201
202    # Private attributes for session tracking
203    _vid: int | None = PrivateAttr(default=None)
204    _uid: str | None = PrivateAttr(default=None)
205    _session: UniSession | None = PrivateAttr(default=None)
206    _dirty: set[str] = PrivateAttr(default_factory=set)
207    _is_new: bool = PrivateAttr(default=True)
208
209    def __init_subclass__(cls, **kwargs: Any) -> None:
210        super().__init_subclass__(**kwargs)
211        # Set default label to class name if not specified
212        if not cls.__label__:
213            cls.__label__ = cls.__name__
214
215    def model_post_init(self, __context: Any) -> None:
216        """Clear dirty tracking after construction."""
217        super().model_post_init(__context)
218        self._dirty = set()
219
220    @property
221    def vid(self) -> int | None:
222        """The vertex ID assigned by the database."""
223        return self._vid
224
225    @property
226    def uid(self) -> str | None:
227        """The unique identifier (content-addressed hash)."""
228        return self._uid
229
230    @property
231    def is_persisted(self) -> bool:
232        """Whether this node has been saved to the database."""
233        return self._vid is not None
234
235    @property
236    def is_dirty(self) -> bool:
237        """Whether this node has unsaved changes."""
238        return bool(self._dirty)
239
240    def __setattr__(self, name: str, value: Any) -> None:
241        # Track dirty fields (but not private attributes)
242        if not name.startswith("_") and hasattr(self, "_dirty"):
243            self._dirty.add(name)
244        super().__setattr__(name, value)
245
246    def _mark_clean(self) -> None:
247        """Mark all fields as clean (called after commit)."""
248        self._dirty.clear()
249        self._is_new = False
250
251    def _attach_session(
252        self, session: UniSession, vid: int, uid: str | None = None
253    ) -> None:
254        """Attach this node to a session with its database IDs."""
255        self._session = session
256        self._vid = vid
257        self._uid = uid
258        self._is_new = False
259
260    @classmethod
261    def get_property_fields(cls) -> dict[str, FieldInfo]:
262        """Get all property fields (excluding relationships)."""
263        return {
264            name: info
265            for name, info in cls.model_fields.items()
266            if name not in cls.__relationships__
267        }
268
269    @classmethod
270    def get_relationship_fields(cls) -> dict[str, RelationshipConfig]:
271        """Get all relationship field configurations."""
272        return cls.__relationships__
273
274    def to_properties(self) -> dict[str, Any]:
275        """Convert to a property dictionary for database storage.
276
277        Uses python_to_db_value() for type conversion. Includes None
278        explicitly so null-outs work.
279        """
280        return _model_to_properties(self, self.get_property_fields())
281
282    @classmethod
283    def from_properties(
284        cls,
285        props: dict[str, Any],
286        *,
287        vid: int | None = None,
288        uid: str | None = None,
289        session: UniSession | None = None,
290    ) -> UniNode:
291        """Create an instance from a property dictionary.
292
293        Accepts _id (string->int vid) and _label from uni-db node dicts.
294        Does not mutate the input dict.
295        """
296        data = dict(props)
297
298        raw_id = data.pop("_id", None)
299        if raw_id is not None and vid is None:
300            vid = int(raw_id) if not isinstance(raw_id, int) else raw_id
301        data.pop("_label", None)
302
303        converted = _convert_db_values(data, cls)
304
305        instance = cls.model_validate(converted)
306        if vid is not None:
307            instance._vid = vid
308        if uid is not None:
309            instance._uid = uid
310        if session is not None:
311            instance._session = session
312        instance._is_new = vid is None
313        instance._dirty = set()
314        return instance
315
316    def __repr__(self) -> str:
317        vid_str = f"vid={self._vid}" if self._vid else "unsaved"
318        return f"{self.__class__.__name__}({vid_str}, {super().__repr__()})"

321class UniEdge(BaseModel, metaclass=UniModelMeta):
322    """
323    Base class for graph edge models with properties.
324
325    Subclass this to define edge types with typed properties.
326    Edges represent relationships between nodes.
327
328    Attributes:
329        __edge_type__: The edge type name.
330        __from__: The source node type(s).
331        __to__: The target node type(s).
332
333    Private Attributes:
334        _eid: The edge ID assigned by the database.
335        _src_vid: The source vertex ID.
336        _dst_vid: The destination vertex ID.
337        _session: Reference to the owning session.
338
339    Example:
340        >>> class FriendshipEdge(UniEdge):
341        ...     __edge_type__ = "FRIEND_OF"
342        ...     __from__ = Person
343        ...     __to__ = Person
344        ...
345        ...     since: date
346        ...     strength: float = 1.0
347    """
348
349    model_config = ConfigDict(
350        extra="forbid",
351        validate_assignment=True,
352        arbitrary_types_allowed=True,
353        use_enum_values=True,
354    )
355
356    # Class-level configuration
357    __edge_type__: ClassVar[str] = ""
358    __from__: ClassVar[type[UniNode] | tuple[type[UniNode], ...] | None] = None
359    __to__: ClassVar[type[UniNode] | tuple[type[UniNode], ...] | None] = None
360    __relationships__: ClassVar[dict[str, RelationshipConfig]] = {}
361
362    # Private attributes
363    _eid: int | None = PrivateAttr(default=None)
364    _src_vid: int | None = PrivateAttr(default=None)
365    _dst_vid: int | None = PrivateAttr(default=None)
366    _session: UniSession | None = PrivateAttr(default=None)
367    _is_new: bool = PrivateAttr(default=True)
368
369    def __init_subclass__(cls, **kwargs: Any) -> None:
370        super().__init_subclass__(**kwargs)
371        # Set default edge type to class name if not specified
372        if not cls.__edge_type__:
373            cls.__edge_type__ = cls.__name__
374
375    @property
376    def eid(self) -> int | None:
377        """The edge ID assigned by the database."""
378        return self._eid
379
380    @property
381    def src_vid(self) -> int | None:
382        """The source vertex ID."""
383        return self._src_vid
384
385    @property
386    def dst_vid(self) -> int | None:
387        """The destination vertex ID."""
388        return self._dst_vid
389
390    @property
391    def is_persisted(self) -> bool:
392        """Whether this edge has been saved to the database."""
393        return self._eid is not None
394
395    def _attach(
396        self,
397        session: UniSession,
398        eid: int,
399        src_vid: int,
400        dst_vid: int,
401    ) -> None:
402        """Attach this edge to a session with its database IDs."""
403        self._session = session
404        self._eid = eid
405        self._src_vid = src_vid
406        self._dst_vid = dst_vid
407        self._is_new = False
408
409    @classmethod
410    def get_from_labels(cls) -> list[str]:
411        """Get the source label names."""
412        if cls.__from__ is None:
413            return []
414        if isinstance(cls.__from__, tuple):
415            return [n.__label__ for n in cls.__from__]
416        return [cls.__from__.__label__]
417
418    @classmethod
419    def get_to_labels(cls) -> list[str]:
420        """Get the target label names."""
421        if cls.__to__ is None:
422            return []
423        if isinstance(cls.__to__, tuple):
424            return [n.__label__ for n in cls.__to__]
425        return [cls.__to__.__label__]
426
427    @classmethod
428    def get_property_fields(cls) -> dict[str, FieldInfo]:
429        """Get all property fields."""
430        return dict(cls.model_fields)
431
432    def to_properties(self) -> dict[str, Any]:
433        """Convert to a property dictionary for database storage."""
434        return _model_to_properties(self, self.get_property_fields())
435
436    @classmethod
437    def from_properties(
438        cls,
439        props: dict[str, Any],
440        *,
441        eid: int | None = None,
442        src_vid: int | None = None,
443        dst_vid: int | None = None,
444        session: UniSession | None = None,
445    ) -> UniEdge:
446        """Create an instance from a property dictionary.
447
448        Accepts _id, _type, _src, _dst from uni-db edge dicts.
449        Does not mutate the input dict.
450        """
451        data = dict(props)
452
453        raw_id = data.pop("_id", None)
454        if raw_id is not None and eid is None:
455            eid = int(raw_id) if not isinstance(raw_id, int) else raw_id
456        data.pop("_type", None)
457        raw_src = data.pop("_src", None)
458        if raw_src is not None and src_vid is None:
459            src_vid = int(raw_src) if not isinstance(raw_src, int) else raw_src
460        raw_dst = data.pop("_dst", None)
461        if raw_dst is not None and dst_vid is None:
462            dst_vid = int(raw_dst) if not isinstance(raw_dst, int) else raw_dst
463
464        converted = _convert_db_values(data, cls)
465
466        instance = cls.model_validate(converted)
467        if eid is not None:
468            instance._eid = eid
469        if src_vid is not None:
470            instance._src_vid = src_vid
471        if dst_vid is not None:
472            instance._dst_vid = dst_vid
473        if session is not None:
474            instance._session = session
475        instance._is_new = eid is None
476        return instance
477
478    @classmethod
479    def from_edge_result(
480        cls,
481        data: dict[str, Any],
482        *,
483        session: UniSession | None = None,
484    ) -> UniEdge:
485        """Create an instance from a uni-db edge result dict.
486
487        Convenience method that handles _id, _type, _src, _dst keys.
488        """
489        return cls.from_properties(data, session=session)
490
491    def __repr__(self) -> str:
492        eid_str = f"eid={self._eid}" if self._eid else "unsaved"
493        return f"{self.__class__.__name__}({eid_str}, {super().__repr__()})"

160class UniSession:
161    """
162    Session for interacting with the graph database using Pydantic models.
163
164    The session manages model registration, schema synchronization,
165    and provides CRUD operations and query building.
166
167    Example:
168        >>> from uni_db import Uni
169        >>> from uni_pydantic import UniSession
170        >>>
171        >>> db = Uni("./my_graph")
172        >>> session = UniSession(db)
173        >>> session.register(Person, Company)
174        >>> session.sync_schema()
175        >>>
176        >>> alice = Person(name="Alice", age=30)
177        >>> session.add(alice)
178        >>> session.commit()
179    """
180
181    def __init__(self, db: uni_db.Uni) -> None:
182        self._db = db
183        self._db_session = db.session()
184        self._schema_gen = SchemaGenerator()
185        self._identity_map: WeakValueDictionary[tuple[str, int], UniNode] = (
186            WeakValueDictionary()
187        )
188        self._pending_new: list[UniNode] = []
189        self._pending_delete: list[UniNode] = []
190
191    def __enter__(self) -> UniSession:
192        return self
193
194    def __exit__(
195        self,
196        exc_type: type[BaseException] | None,
197        exc_val: BaseException | None,
198        exc_tb: TracebackType | None,
199    ) -> None:
200        self.close()
201
202    def close(self) -> None:
203        """Close the session and clear all pending state."""
204        self._pending_new.clear()
205        self._pending_delete.clear()
206
207    @property
208    def db(self) -> uni_db.Uni:
209        """Access the underlying uni_db.Uni for low-level operations."""
210        return self._db
211
212    def locy(
213        self, program: str, params: dict[str, Any] | None = None
214    ) -> uni_db.LocyResult:
215        """
216        Evaluate a Locy program and return derived facts, stats, and warnings.
217
218        Delegates to the underlying ``uni_db.Session.locy()``.
219        """
220        return self._db_session.locy(program, params)
221
222    def register(self, *models: type[UniNode] | type[UniEdge]) -> None:
223        """
224        Register model classes with the session.
225
226        Registered models can be used for schema generation and queries.
227
228        Args:
229            *models: UniNode or UniEdge subclasses to register.
230        """
231        self._schema_gen.register(*models)
232
233    def sync_schema(self) -> None:
234        """
235        Synchronize database schema with registered models.
236
237        Creates labels, edge types, properties, and indexes as needed.
238        This is additive-only; it won't remove existing schema elements.
239        """
240        self._schema_gen.apply_to_database(self._db)
241
242    def query(self, model: type[NodeT]) -> QueryBuilder[NodeT]:
243        """
244        Create a query builder for the given model.
245
246        Args:
247            model: The UniNode subclass to query.
248
249        Returns:
250            A QueryBuilder for constructing queries.
251        """
252        return QueryBuilder(self, model)
253
254    def add(self, entity: UniNode) -> None:
255        """
256        Add a new entity to be persisted.
257
258        The entity will be inserted on the next commit().
259        """
260        if entity.is_persisted:
261            raise SessionError(f"Entity {entity!r} is already persisted")
262        entity._session = self
263        self._pending_new.append(entity)
264
265    def add_all(self, entities: Sequence[UniNode]) -> None:
266        """Add multiple entities to be persisted."""
267        for entity in entities:
268            self.add(entity)
269
270    def delete(self, entity: UniNode) -> None:
271        """Mark an entity for deletion."""
272        if not entity.is_persisted:
273            raise NotPersisted(entity)
274        self._pending_delete.append(entity)
275
276    def get(
277        self,
278        model: type[NodeT],
279        vid: int | None = None,
280        uid: str | None = None,
281        **kwargs: Any,
282    ) -> NodeT | None:
283        """
284        Get an entity by ID or unique properties.
285
286        Args:
287            model: The model type to retrieve.
288            vid: Vertex ID to look up.
289            uid: Unique ID to look up.
290            **kwargs: Property equality filters.
291
292        Returns:
293            The model instance or None if not found.
294        """
295        # Check identity map first
296        if vid is not None:
297            cached = self._identity_map.get((model.__label__, vid))
298            if cached is not None:
299                return cached  # type: ignore[return-value]
300
301        # Build query
302        label = model.__label__
303        params: dict[str, Any] = {}
304
305        if vid is not None:
306            cypher = f"MATCH (n:{label}) WHERE id(n) = $vid RETURN {_NODE_RETURN}"
307            params["vid"] = vid
308        elif uid is not None:
309            cypher = f"MATCH (n:{label}) WHERE n._uid = $uid RETURN {_NODE_RETURN}"
310            params["uid"] = uid
311        elif kwargs:
312            # Validate property names
313            for k in kwargs:
314                _validate_property(k, model)
315            conditions = [f"n.{k} = ${k}" for k in kwargs]
316            cypher = f"MATCH (n:{label}) WHERE {' AND '.join(conditions)} RETURN {_NODE_RETURN} LIMIT 1"
317            params.update(kwargs)
318        else:
319            raise ValueError("Must provide vid, uid, or property filters")
320
321        results = self._db_session.query(cypher, params)
322        if not results:
323            return None
324
325        node_data = _row_to_node_dict(results[0].to_dict())
326        if node_data is None:
327            return None
328        return self._result_to_model(node_data, model)
329
330    def refresh(self, entity: UniNode) -> None:
331        """Refresh an entity's properties from the database."""
332        if not entity.is_persisted:
333            raise NotPersisted(entity)
334
335        label = entity.__class__.__label__
336        cypher = f"MATCH (n:{label}) WHERE id(n) = $vid RETURN {_NODE_RETURN}"
337        results = self._db_session.query(cypher, {"vid": entity._vid})
338
339        if not results:
340            raise SessionError(f"Entity with vid={entity._vid} no longer exists")
341
342        # Update properties
343        props = _row_to_node_dict(results[0].to_dict())
344        if props is None:
345            raise SessionError(f"Entity with vid={entity._vid} no longer exists")
346        try:
347            hints = get_type_hints(type(entity))
348        except Exception:
349            hints = {}
350
351        for field_name in entity.get_property_fields():
352            if field_name in props:
353                value = props[field_name]
354                if field_name in hints:
355                    value = db_to_python_value(value, hints[field_name])
356                setattr(entity, field_name, value)
357
358        entity._mark_clean()
359
360    def commit(self) -> None:
361        """
362        Commit all pending changes to the database.
363
364        This persists new entities, updates dirty entities,
365        and deletes marked entities.
366        """
367        # Insert new entities
368        for entity in self._pending_new:
369            self._create_node(entity)
370
371        # Update dirty entities in identity map
372        for (label, vid), entity in list(self._identity_map.items()):
373            if entity.is_dirty and entity.is_persisted:
374                self._update_node(entity)
375
376        # Delete marked entities
377        for entity in self._pending_delete:
378            self._delete_node(entity)
379
380        # Flush to storage
381        self._db.flush()
382
383        # Clear pending lists
384        self._pending_new.clear()
385        self._pending_delete.clear()
386
387    def rollback(self) -> None:
388        """Discard all pending changes."""
389        # Clear pending new — detach entities
390        for entity in self._pending_new:
391            entity._session = None
392        self._pending_new.clear()
393
394        # Clear pending deletes
395        self._pending_delete.clear()
396
397        # Invalidate dirty identity map entries
398        for entity in list(self._identity_map.values()):
399            if entity.is_dirty:
400                self.refresh(entity)
401
402    @contextmanager
403    def transaction(self) -> Iterator[UniTransaction]:
404        """Create a transaction context."""
405        tx = UniTransaction(self)
406        with tx:
407            yield tx
408
409    def begin(self) -> UniTransaction:
410        """Begin a new transaction."""
411        tx = UniTransaction(self)
412        tx._tx = self._db_session.tx()
413        return tx
414
415    def cypher(
416        self,
417        query: str,
418        params: dict[str, Any] | None = None,
419        result_type: type[NodeT] | None = None,
420    ) -> list[NodeT] | list[dict[str, Any]]:
421        """
422        Execute a raw Cypher query.
423
424        Args:
425            query: Cypher query string.
426            params: Query parameters.
427            result_type: Optional model type for result mapping.
428
429        Returns:
430            List of results (model instances if result_type provided).
431        """
432        results = self._db_session.query(query, params)
433
434        if result_type is None:
435            return [r.to_dict() for r in results]
436
437        # Map results to model instances
438        mapped = []
439        for raw_row in results:
440            row = raw_row.to_dict()
441            # Try to find node data in the row
442            for key, value in row.items():
443                if isinstance(value, dict):
444                    # Check for _id/_label keys (uni-db node dict)
445                    if "_id" in value and "_label" in value:
446                        instance = self._result_to_model(value, result_type)
447                        if instance:
448                            mapped.append(instance)
449                            break
450                    # Also check if _label matches registered model
451                    elif "_label" in value:
452                        label = value["_label"]
453                        if label in self._schema_gen._node_models:
454                            model = self._schema_gen._node_models[label]
455                            instance = self._result_to_model(value, model)
456                            if instance:
457                                mapped.append(instance)
458                                break
459            else:
460                # Try the first column
461                first_value = next(iter(row.values()), None)
462                if isinstance(first_value, dict):
463                    instance = self._result_to_model(first_value, result_type)
464                    if instance:
465                        mapped.append(instance)
466
467        return mapped
468
469    @staticmethod
470    def _validate_edge_endpoints(
471        source: UniNode, target: UniNode
472    ) -> tuple[int, int, str, str]:
473        """Validate that both endpoints are persisted and return (src_vid, dst_vid, src_label, dst_label)."""
474        if not source.is_persisted:
475            raise NotPersisted(source)
476        if not target.is_persisted:
477            raise NotPersisted(target)
478        return (
479            source._vid,
480            target._vid,
481            source.__class__.__label__,
482            target.__class__.__label__,
483        )
484
485    @staticmethod
486    def _normalize_edge_properties(
487        properties: dict[str, Any] | UniEdge | None,
488    ) -> dict[str, Any]:
489        """Normalize edge properties from dict, UniEdge, or None."""
490        if isinstance(properties, UniEdge):
491            return properties.to_properties()
492        if properties:
493            return properties
494        return {}
495
496    def create_edge(
497        self,
498        source: UniNode,
499        edge_type: str,
500        target: UniNode,
501        properties: dict[str, Any] | UniEdge | None = None,
502    ) -> None:
503        """Create an edge between two nodes."""
504        src_vid, dst_vid, src_label, dst_label = self._validate_edge_endpoints(
505            source, target
506        )
507        props = self._normalize_edge_properties(properties)
508
509        # Build CREATE edge query with labels (required by Cypher implementation)
510        props_str = ", ".join(f"{k}: ${k}" for k in props)
511        if props_str:
512            cypher = f"MATCH (a:{src_label}), (b:{dst_label}) WHERE a._vid = $src AND b._vid = $dst CREATE (a)-[r:{edge_type} {{{props_str}}}]->(b)"
513        else:
514            cypher = f"MATCH (a:{src_label}), (b:{dst_label}) WHERE a._vid = $src AND b._vid = $dst CREATE (a)-[r:{edge_type}]->(b)"
515
516        params = {"src": src_vid, "dst": dst_vid, **props}
517        with self._db_session.tx() as tx:
518            tx.execute(cypher, params)
519            tx.commit()
520
521    def delete_edge(
522        self,
523        source: UniNode,
524        edge_type: str,
525        target: UniNode,
526    ) -> int:
527        """Delete edges between two nodes. Returns the number of deleted edges."""
528        src_vid, dst_vid, src_label, dst_label = self._validate_edge_endpoints(
529            source, target
530        )
531        cypher = (
532            f"MATCH (a:{src_label})-[r:{edge_type}]->(b:{dst_label}) "
533            f"WHERE a._vid = $src AND b._vid = $dst "
534            f"DELETE r RETURN count(r) as count"
535        )
536        with self._db_session.tx() as tx:
537            results = tx.query(cypher, {"src": src_vid, "dst": dst_vid})
538            tx.commit()
539        return cast(int, results[0]["count"]) if results else 0
540
541    def update_edge(
542        self,
543        source: UniNode,
544        edge_type: str,
545        target: UniNode,
546        properties: dict[str, Any],
547    ) -> int:
548        """Update properties on edges between two nodes. Returns the number of updated edges."""
549        src_vid, dst_vid, src_label, dst_label = self._validate_edge_endpoints(
550            source, target
551        )
552        set_parts = [f"r.{k} = ${k}" for k in properties]
553        params: dict[str, Any] = {"src": src_vid, "dst": dst_vid, **properties}
554        cypher = (
555            f"MATCH (a:{src_label})-[r:{edge_type}]->(b:{dst_label}) "
556            f"WHERE a._vid = $src AND b._vid = $dst "
557            f"SET {', '.join(set_parts)} "
558            f"RETURN count(r) as count"
559        )
560        with self._db_session.tx() as tx:
561            results = tx.query(cypher, params)
562            tx.commit()
563        return cast(int, results[0]["count"]) if results else 0
564
565    def get_edge(
566        self,
567        source: UniNode,
568        edge_type: str,
569        target: UniNode,
570        edge_model: type[EdgeT] | None = None,
571    ) -> list[dict[str, Any]] | list[EdgeT]:
572        """Get edges between two nodes. Returns dicts or edge model instances."""
573        src_vid, dst_vid, src_label, dst_label = self._validate_edge_endpoints(
574            source, target
575        )
576        cypher = (
577            f"MATCH (a:{src_label})-[r:{edge_type}]->(b:{dst_label}) "
578            f"WHERE a._vid = $src AND b._vid = $dst "
579            f"RETURN properties(r) AS _props, id(r) AS _eid"
580        )
581        results = self._db_session.query(cypher, {"src": src_vid, "dst": dst_vid})
582        rows = [r.to_dict() for r in results]
583
584        if edge_model is None:
585            edge_dicts: list[dict[str, Any]] = []
586            for row in rows:
587                props = row.get("_props", {})
588                if isinstance(props, dict):
589                    edge_dict = dict(props)
590                    edge_dict["_eid"] = row.get("_eid")
591                    edge_dicts.append(edge_dict)
592            return edge_dicts
593
594        edges = []
595        for row in rows:
596            r_data = row.get("_props", {})
597            if isinstance(r_data, dict):
598                edge = edge_model.from_properties(
599                    r_data,
600                    src_vid=src_vid,
601                    dst_vid=dst_vid,
602                    session=self,
603                )
604                edges.append(edge)
605        return edges
606
607    def bulk_add(self, entities: Sequence[UniNode]) -> list[int]:
608        """
609        Bulk-add entities using bulk_writer for performance.
610
611        Groups entities by label and uses db.bulk_writer().
612        Returns VIDs and attaches sessions.
613
614        Args:
615            entities: Sequence of UniNode instances to bulk-insert.
616
617        Returns:
618            List of assigned vertex IDs.
619
620        Raises:
621            BulkLoadError: If bulk insertion fails.
622        """
623        if not entities:
624            return []
625
626        # Group by label
627        by_label: dict[str, list[UniNode]] = {}
628        for entity in entities:
629            label = entity.__class__.__label__
630            if label not in by_label:
631                by_label[label] = []
632            by_label[label].append(entity)
633
634        all_vids: list[int] = []
635        try:
636            for label, group in by_label.items():
637                # Run before_create hooks
638                for entity in group:
639                    run_hooks(entity, _BEFORE_CREATE)
640
641                # Convert to property dicts
642                prop_dicts = [e.to_properties() for e in group]
643
644                # Bulk insert via transaction
645                tx = self._db_session.tx()
646                with tx.bulk_writer().build() as bw:
647                    vids = bw.insert_vertices(label, prop_dicts)
648                    bw.commit()
649                tx.commit()
650
651                # Attach sessions and record VIDs
652                for entity, vid in zip(group, vids):
653                    entity._attach_session(self, vid)
654                    self._identity_map[(label, vid)] = entity
655                    run_hooks(entity, _AFTER_CREATE)
656                    entity._mark_clean()
657
658                all_vids.extend(vids)
659        except Exception as e:
660            raise BulkLoadError(f"Bulk insert failed: {e}") from e
661
662        return all_vids
663
664    def explain(self, cypher: str) -> uni_db.ExplainOutput:
665        """Get the query execution plan without running it."""
666        return self._db_session.explain(cypher)
667
668    def profile(self, cypher: str) -> tuple[uni_db.QueryResult, uni_db.ProfileOutput]:
669        """Run the query with profiling and return results + stats."""
670        return self._db_session.profile(cypher)
671
672    def save_schema(self, path: str) -> None:
673        """Save the database schema to a file."""
674        self._db.save_schema(path)
675
676    def load_schema(self, path: str) -> None:
677        """Load a database schema from a file."""
678        self._db.load_schema(path)
679
680    # -------------------------------------------------------------------------
681    # Internal Methods
682    # -------------------------------------------------------------------------
683
684    def _create_node(self, entity: UniNode) -> None:
685        """Create a node in the database."""
686        # Run before_create hooks
687        run_hooks(entity, _BEFORE_CREATE)
688
689        label = entity.__class__.__label__
690        props = entity.to_properties()
691
692        # Build CREATE query
693        props_str = ", ".join(f"{k}: ${k}" for k in props)
694        cypher = f"CREATE (n:{label} {{{props_str}}}) RETURN id(n) as vid"
695
696        with self._db_session.tx() as tx:
697            results = tx.query(cypher, props)
698            tx.commit()
699        if results:
700            vid = results[0]["vid"]
701            entity._attach_session(self, vid)
702
703            # Add to identity map
704            self._identity_map[(label, vid)] = entity
705
706        # Run after_create hooks
707        run_hooks(entity, _AFTER_CREATE)
708        entity._mark_clean()
709
710    def _create_node_in_tx(self, entity: UniNode, tx: uni_db.Transaction) -> None:
711        """Create a node within a transaction."""
712        run_hooks(entity, _BEFORE_CREATE)
713
714        label = entity.__class__.__label__
715        props = entity.to_properties()
716
717        props_str = ", ".join(f"{k}: ${k}" for k in props)
718        cypher = f"CREATE (n:{label} {{{props_str}}}) RETURN id(n) as vid"
719
720        results = tx.query(cypher, props)
721        if results:
722            vid = results[0]["vid"]
723            entity._attach_session(self, vid)
724            self._identity_map[(label, vid)] = entity
725
726        run_hooks(entity, _AFTER_CREATE)
727
728    def _create_edge_in_tx(
729        self,
730        source: UniNode,
731        edge_type: str,
732        target: UniNode,
733        properties: UniEdge | None,
734        tx: uni_db.Transaction,
735    ) -> None:
736        """Create an edge within a transaction."""
737        props = properties.to_properties() if properties else {}
738        src_label = source.__class__.__label__
739        dst_label = target.__class__.__label__
740
741        props_str = ", ".join(f"{k}: ${k}" for k in props)
742        if props_str:
743            cypher = f"MATCH (a:{src_label}), (b:{dst_label}) WHERE a._vid = $src AND b._vid = $dst CREATE (a)-[:{edge_type} {{{props_str}}}]->(b)"
744        else:
745            cypher = f"MATCH (a:{src_label}), (b:{dst_label}) WHERE a._vid = $src AND b._vid = $dst CREATE (a)-[:{edge_type}]->(b)"
746
747        params = {"src": source._vid, "dst": target._vid, **props}
748        tx.query(cypher, params)
749
750    def _update_node(self, entity: UniNode) -> None:
751        """Update a node in the database."""
752        run_hooks(entity, _BEFORE_UPDATE)
753
754        label = entity.__class__.__label__
755
756        # Convert dirty prop values via python_to_db_value
757        try:
758            hints = get_type_hints(type(entity))
759        except Exception:
760            hints = {}
761
762        dirty_props = {}
763        for name in entity._dirty:
764            value = getattr(entity, name)
765            if name in hints:
766                value = python_to_db_value(value, hints[name])
767            dirty_props[name] = value
768
769        if not dirty_props:
770            return
771
772        set_clause = ", ".join(f"n.{k} = ${k}" for k in dirty_props)
773        cypher = f"MATCH (n:{label}) WHERE id(n) = $vid SET {set_clause}"
774        params = {"vid": entity._vid, **dirty_props}
775
776        with self._db_session.tx() as tx:
777            tx.execute(cypher, params)
778            tx.commit()
779
780        run_hooks(entity, _AFTER_UPDATE)
781        entity._mark_clean()
782
783    def _delete_node(self, entity: UniNode) -> None:
784        """Delete a node from the database."""
785        run_hooks(entity, _BEFORE_DELETE)
786
787        label = entity.__class__.__label__
788        vid = entity._vid
789
790        # DETACH DELETE to also remove connected edges
791        cypher = f"MATCH (n:{label}) WHERE id(n) = $vid DETACH DELETE n"
792        with self._db_session.tx() as tx:
793            tx.execute(cypher, {"vid": vid})
794            tx.commit()
795
796        # Remove from identity map
797        if vid is not None and (label, vid) in self._identity_map:
798            del self._identity_map[(label, vid)]
799
800        # Clear entity IDs
801        entity._vid = None
802        entity._uid = None
803        entity._session = None
804
805        run_hooks(entity, _AFTER_DELETE)
806
807    def _result_to_model(
808        self,
809        data: dict[str, Any],
810        model: type[NodeT],
811    ) -> NodeT | None:
812        """Convert a query result row to a model instance.
813
814        Does not mutate the input dict.
815        """
816        if not data:
817            return None
818
819        # Work on a copy
820        data = dict(data)
821
822        # Run before_load hooks
823        data = run_class_hooks(model, _BEFORE_LOAD, data) or data
824
825        # Extract _id → vid (uni-db returns _id as string or int)
826        vid = data.pop("_id", None)
827        if vid is None:
828            vid = data.pop("_vid", None)
829        if vid is None:
830            vid = data.pop("vid", None)
831        if vid is not None and not isinstance(vid, int):
832            vid = int(vid)
833
834        # Remove _label (informational)
835        data.pop("_label", None)
836
837        try:
838            instance = cast(
839                NodeT,
840                model.from_properties(
841                    data,
842                    vid=vid,
843                    session=self,
844                ),
845            )
846        except Exception:
847            # If validation fails, return None
848            return None
849
850        # Add to identity map if we have a vid
851        if vid is not None:
852            existing = self._identity_map.get((model.__label__, vid))
853            if existing is not None:
854                return cast(NodeT, existing)
855            self._identity_map[(model.__label__, vid)] = instance
856
857        # Run after_load hooks
858        run_hooks(instance, _AFTER_LOAD)
859
860        return instance
861
862    def _load_relationship(
863        self,
864        entity: UniNode,
865        descriptor: RelationshipDescriptor[Any],
866    ) -> list[UniNode] | UniNode | None:
867        """Load a relationship for an entity."""
868        if not entity.is_persisted:
869            raise NotPersisted(entity)
870
871        config = descriptor.config
872        label = entity.__class__.__label__
873        pattern = _edge_pattern(config.edge_type, config.direction)
874
875        cypher = (
876            f"MATCH (a:{label}){pattern}(b) WHERE id(a) = $vid "
877            f"RETURN properties(b) AS _props, id(b) AS _vid, labels(b) AS _labels"
878        )
879        results = self._db_session.query(cypher, {"vid": entity._vid})
880
881        nodes = []
882        for raw_row in results:
883            row = raw_row.to_dict()
884            node_data = _row_to_node_dict(row)
885            if node_data is None:
886                continue
887            # Try to find the model for this node
888            node_label = node_data.get("_label")
889            if node_label and node_label in self._schema_gen._node_models:
890                model = self._schema_gen._node_models[node_label]
891                instance = self._result_to_model(node_data, model)
892                if instance:
893                    nodes.append(instance)
894
895        if not descriptor.is_list:
896            return nodes[0] if nodes else None
897        return nodes
898
899    def _eager_load_relationships(
900        self,
901        entities: list[NodeT],
902        relationships: list[str],
903    ) -> None:
904        """Eager load relationships for a list of entities."""
905        if not entities:
906            return
907
908        model = type(entities[0])
909        rel_configs = model.get_relationship_fields()
910
911        for rel_name in relationships:
912            if rel_name not in rel_configs:
913                continue
914
915            config = rel_configs[rel_name]
916            label = model.__label__
917            vids = [e._vid for e in entities if e._vid is not None]
918
919            if not vids:
920                continue
921
922            pattern = _edge_pattern(config.edge_type, config.direction)
923            cypher = (
924                f"MATCH (a:{label}){pattern}(b) WHERE id(a) IN $vids "
925                f"RETURN id(a) as src_vid, properties(b) AS _props, id(b) AS _vid, labels(b) AS _labels"
926            )
927            results = self._db_session.query(cypher, {"vids": vids})
928
929            # Group results by source vid
930            by_source: dict[int, list[Any]] = {}
931            for raw_row in results:
932                row = raw_row.to_dict()
933                src_vid = row["src_vid"]
934                node_data = _row_to_node_dict(row)
935                if node_data is None:
936                    continue
937                if src_vid not in by_source:
938                    by_source[src_vid] = []
939                by_source[src_vid].append(node_data)
940
941            # Set cached values on entities
942            for entity in entities:
943                if entity._vid in by_source:
944                    related = by_source[entity._vid]
945                    cache_attr = f"_rel_cache_{rel_name}"
946                    setattr(entity, cache_attr, related)