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