tapestry-orm 0.0.1__py3-none-any.whl

tapestry/edge.py

@@ -0,0 +1,271 @@
+from abc import ABC
+from copy import deepcopy
+from typing import Self, Iterable, get_type_hints, Any, ClassVar, Union, TYPE_CHECKING
+from pydantic.fields import FieldInfo, ComputedFieldInfo
+from surrealdb import AsyncWsSurrealConnection, AsyncHttpSurrealConnection, RecordID
+
+
+from .base import Base
+from .node import Node
+from .utils import replace_type
+from .table import Table, Link, flatten_type
+from .field import Traversal, Direction, Condition, LogicalCondition, Field
+
+
+class Edge(Base, ABC, discarded=True):
+    """
+    Base class for SurrealDB edge (relation) tables.
+
+    Edge represents relationship tables in SurrealDB that connect two nodes.
+    All relationship tables should inherit from this class and define
+    'in_' and 'out_' fields to specify the connected node types.
+
+    Class Attributes:
+        _directed (bool): Whether the relationship is directional (default: True)
+
+    Required Fields:
+        in_: The source node of the relationship
+        out_: The target node of the relationship
+
+    Example:
+        >>> from tapestry import Edge, Node
+        >>> from datetime import date
+        >>>
+        >>> class Person(Node):
+        ...     name: str
+        ...
+        >>> class Role(Node):
+        ...     title: str
+        ...
+        >>> class BelongsTo(Edge):
+        ...     in_: Person  # Person belongs to Role
+        ...     out_: Role
+        ...     begin_date: date
+        ...     end_date: date | None = None
+        ...
+        >>> # Create a relationship
+        >>> person = Person(name="John Doe")
+        >>> role = Role(title="Manager")
+        >>> belongs = BelongsTo(
+        ...     in_=person,
+        ...     out_=role,
+        ...     begin_date=date(2020, 1, 1)
+        ... )
+        >>> await belongs.relate(db)
+
+    Notes:
+        - Requires both 'in_' and 'out_' fields to be defined
+        - Automatically creates SurrealDB RELATION tables
+        - Supports directional and bidirectional relationships
+        - Can have additional fields beyond in_ and out_
+        - Use relate() instead of create() for edge records
+    """
+    # in_: Node
+    # out_: Node
+    #
+    _directed: bool = True
+
+    def __init_subclass__(cls, directed: bool = True, **kwargs):
+        """
+        Configure edge subclasses with relationship properties.
+
+        Args:
+            directed: Whether the relationship is directional. If False,
+                     creates a bidirectional relationship where in/out
+                     order doesn't matter.
+            **kwargs: Additional configuration passed to parent
+
+        Raises:
+            TypeError: If the subclass doesn't define both 'in_' and 'out_' fields
+        """
+        super().__init_subclass__(**kwargs)
+
+        # Get type hints for the subclass
+        hints = get_type_hints(cls)
+
+        # Check if required fields are annotated
+        required_fields = {'in_', 'out_'}
+        missing_fields = required_fields - set(hints.keys())
+        if missing_fields:
+            raise TypeError(
+                f"Class {cls.__name__} must have field annotations for: {missing_fields}"
+            )
+
+        cls._directed = directed
+
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs):
+        relation = None
+        in_, out_ = cls.model_fields.pop('in_', None), cls.model_fields.pop('out_', None)
+        if in_ or out_:
+            if in_ and out_:
+                assert in_.annotation
+                assert out_.annotation
+                assert all(issubclass(t, Node) for t in flatten_type(in_.annotation))
+                assert all(issubclass(t, Node) for t in flatten_type(out_.annotation))
+                relation = Link(_in=in_.annotation, _out=out_.annotation, symetric=not cls._directed)
+                cls.is_relation = True
+            else:
+                raise Exception("You need to specify both `_in` and `_out` to define a relation")
+        else:
+            raise Exception("You need to specify `_in` and `_out` to define an Edge")
+        super().__pydantic_init_subclass__(relation=relation, **kwargs)
+        # need to add those back for validation at instanciation
+        # otherwise, the field is skipped
+        in_.annotation = replace_type(in_.annotation, Base, Union[Base, RecordID])
+        out_.annotation = replace_type(out_.annotation, Base, Union[Base, RecordID])
+        cls.model_fields["in_"] = in_
+        cls.model_fields["out_"] = out_
+        cls.model_rebuild(force=True)
+
+
+
+    def db_dump(self) -> dict[str, Any]:
+        """
+        Serialize the edge instance for database insertion.
+
+        Converts Python field names to SurrealDB format by renaming
+        'in_' to 'in' and 'out_' to 'out'.
+
+        Returns:
+            dict[str, Any]: Dictionary ready for SurrealDB relation operations
+        """
+        dump = super().db_dump()
+        dump["in"] = dump.pop("in_")
+        dump["out"] = dump.pop("out_")
+        return dump
+
+
+    async def relate(
+        self,
+        db: AsyncWsSurrealConnection | AsyncHttpSurrealConnection
+    ) -> Self:
+        """
+        Create a relationship record in the database.
+
+        Creates an edge record connecting two nodes in SurrealDB.
+        This is the primary method for creating relationships.
+
+        Args:
+            db: Active SurrealDB connection (WebSocket or HTTP)
+
+        Returns:
+            Self: The same instance with ID assigned
+
+        Example:
+            >>> person = await Person.create(db)
+            >>> role = await Role.create(db)
+            >>> belongs = BelongsTo(
+            ...     in_=person,
+            ...     out_=role,
+            ...     begin_date=date.today()
+            ... )
+            >>> await belongs.relate(db)
+            >>> print(belongs.id)  # Has an ID like belongs_to:xyz
+
+        Raises:
+            Exception: If called on a non-relation table
+
+        Notes:
+            - Both in_ and out_ nodes must exist in the database
+            - Creates a directed or bidirectional edge based on class configuration
+            - The edge record gets a unique ID from SurrealDB
+        """
+        if not self.is_relation:
+            raise Exception("You should use .create() to create a record")
+        thing = self.__class__.__name__.lower()
+        relation = await db.insert_relation(
+            thing,
+            self.db_dump()
+        )
+        if isinstance(relation, list):
+            self.id = relation[0]["id"]
+        else:
+            self.id = relation["id"]
+        return self
+
+
+    @classmethod
+    async def insert(
+        cls,
+        db: AsyncWsSurrealConnection | AsyncHttpSurrealConnection, others: Iterable[Self]
+    ) -> list[Self]:
+        """
+        Batch insert multiple edge records into the database.
+
+        Efficiently creates multiple relationships in a single operation.
+        All edge instances are updated with their assigned IDs.
+
+        Args:
+            db: Active SurrealDB connection (WebSocket or HTTP)
+            others: Iterable of edge instances to insert
+
+        Returns:
+            list[Self]: The same instances with IDs assigned
+
+        Example:
+            >>> relationships = [
+            ...     BelongsTo(in_=person1, out_=role1, begin_date=date(2020, 1, 1)),
+            ...     BelongsTo(in_=person2, out_=role2, begin_date=date(2021, 1, 1)),
+            ... ]
+            >>> inserted = await BelongsTo.insert(db, relationships)
+            >>> # All relationships now have IDs
+
+        Notes:
+            - More efficient than multiple relate() calls
+            - All edges are inserted in a single transaction
+            - Original instances are modified with IDs
+        """
+        others = list(others)
+        inserted = await db.insert_relation(
+            cls.__name__.lower(),
+            [other.db_dump() for other in others]
+        )
+        for insert, original in zip(inserted, others):
+            original.id = insert["id"]
+        return others
+
+    @classmethod
+    def where(cls, condition):
+        """
+        Add a WHERE condition to this edge for use in graph traversals.
+
+        Creates a conditional edge that can be used in query traversals
+        to filter relationships based on their properties.
+
+        Args:
+            condition: A condition expression that filters edge records
+
+        Returns:
+            EdgeWithCondition: A wrapped edge class with the condition attached
+
+        Example:
+            >>> # Find all people who belong to roles that started after 2020
+            >>> query = (Q(Person) >>
+            ...          BelongsTo.where(BelongsTo.begin_date > date(2020, 1, 1)) >>
+            ...          Role)
+            >>> results = await query.execute(db)
+
+        Notes:
+            - Used primarily in graph traversal queries
+            - Conditions are applied during traversal, not at definition
+            - Can filter based on any edge properties
+        """
+
+        # Create a new class-like object that wraps the Edge with a condition
+        class EdgeWithCondition:
+            def __init__(self, edge_class, condition):
+                self.edge_class = edge_class
+                self.condition = condition
+
+            def __repr__(self):
+                return f"{self.edge_class.__name__.lower()}[WHERE {self.condition.to_surreal()}]"
+
+            # Allow this to be used in traversals
+            def to_traversal(self, direction: Direction) -> Traversal:
+                traversal = Traversal(direction, self.edge_class)
+                traversal.where_condition = self.condition
+                return traversal
+
+        return EdgeWithCondition(cls, condition)

tapestry/edge.pyi

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import ClassVar, Iterable, TypeVar, Type
+from surrealdb import AsyncWsSurrealConnection, AsyncHttpSurrealConnection
+
+from .base import Base
+from .node import Node
+from .field import Condition, Direction, Traversal
+
+T = TypeVar('T', bound='Edge')
+
+class EdgeWithCondition:
+    edge_class: Type[Edge]
+    condition: Condition
+
+    def __init__(self, edge_class: Type[Edge], condition: Condition) -> None: ...
+    def __repr__(self) -> str: ...
+    def to_traversal(self, direction: Direction) -> Traversal: ...
+
+class Edge(Base, ABC):
+    _directed: ClassVar[bool]
+    is_relation: ClassVar[bool]
+
+    def __init_subclass__(cls, directed: bool = True, **kwargs: object) -> None: ...
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs: object) -> None: ...
+
+    def db_dump(self) -> dict[str, object]: ...
+
+    async def relate(
+        self: T,
+        db: AsyncWsSurrealConnection | AsyncHttpSurrealConnection
+    ) -> T: ...
+
+    @classmethod
+    async def insert(
+        cls: type[T],
+        db: AsyncWsSurrealConnection | AsyncHttpSurrealConnection,
+        others: Iterable[T]
+    ) -> list[T]: ...
+
+    @classmethod
+    def where(cls: type[T], condition: Condition) -> EdgeWithCondition: ...

tapestry/engine.py

@@ -0,0 +1,301 @@
+import asyncio
+import logging
+import contextlib
+
+from collections.abc import AsyncIterator, AsyncGenerator
+from surrealdb import AsyncSurreal, AsyncWsSurrealConnection, AsyncHttpSurrealConnection
+
+
+logger = logging.getLogger(__name__)
+
+class SurrealClientPool:
+    """
+    Connection pool manager for SurrealDB clients.
+
+    Manages a pool of authenticated SurrealDB connections that can be
+    reused across async operations. This improves performance by avoiding
+    the overhead of creating new connections for each operation.
+
+    Attributes:
+        url: SurrealDB connection URL
+        auth_payload: Authentication credentials dictionary
+        pool_size: Number of connections to maintain in the pool
+        connect_timeout: Timeout for connection attempts in seconds
+
+    Example:
+        >>> pool = SurrealClientPool(
+        ...     url="ws://localhost:8000/rpc",
+        ...     auth_payload={"username": "root", "password": "root"},
+        ...     pool_size=5
+        ... )
+        >>> await pool.start()
+        >>> async with pool.acquire() as client:
+        ...     await client.use("mydb", "myns")
+        ...     results = await client.select("person")
+    """
+
+    def __init__(
+        self,
+        url: str,
+        auth_payload: dict,
+        pool_size: int = 4,
+        connect_timeout: float = 5.
+    ):
+        """
+        Initialize the connection pool.
+
+        Args:
+            url: SurrealDB connection URL (ws:// or http://)
+            auth_payload: Dictionary with authentication credentials
+            pool_size: Number of connections to maintain (default: 4)
+            connect_timeout: Connection timeout in seconds (default: 5.0)
+        """
+        self.url = url
+        self.auth_payload = auth_payload
+        self.pool_size = pool_size
+        self.connect_timeout = connect_timeout
+        self._queue: asyncio.Queue = asyncio.Queue()
+        self._clients = []   # keep references for shutdown
+        self._closed = False
+
+    async def _make_one(self):
+        """
+        Create and authenticate a single SurrealDB client.
+
+        Returns:
+            An authenticated AsyncSurreal client ready for use
+
+        Raises:
+            Exception: If connection or authentication fails
+        """
+        c = AsyncSurreal(self.url)
+        await c.signin(self.auth_payload)
+        return c
+
+
+    async def start(self):
+        """
+        Initialize the connection pool with authenticated clients.
+
+        Creates the specified number of clients, authenticates them,
+        and adds them to the pool queue for reuse.
+
+        Raises:
+            Exception: If unable to create the required number of clients
+        """
+        for i in range(self.pool_size):
+            try:
+                client = await self._make_one()
+            except Exception as e:
+                logger.exception("Failed to create surreal client #%d: %s", i, e)
+                # retry/backoff before giving up; here we re-raise after a small delay
+                await asyncio.sleep(0.5)
+                client = await self._make_one()
+            self._clients.append(client)
+            self._queue.put_nowait(client)
+        logger.info("SurrealDB pool started (%d clients)", len(self._clients))
+
+    async def close(self):
+        """
+        Close all connections and shut down the pool.
+
+        Drains the queue and closes all client connections gracefully.
+        This should be called when the application shuts down.
+        """
+        if self._closed:
+            return
+        self._closed = True
+        # drain queue so other tasks don't try to reuse
+        while not self._queue.empty():
+            try:
+                _ = self._queue.get_nowait()
+                self._queue.task_done()
+            except asyncio.QueueEmpty:
+                break
+        # close actual clients
+        for c in self._clients:
+            try:
+                await c.close()
+            except Exception:
+                logger.exception("Error closing surreal client")
+        self._clients.clear()
+        logger.info("SurrealDB pool closed")
+
+    @contextlib.asynccontextmanager
+    async def acquire(self) -> AsyncIterator[AsyncWsSurrealConnection | AsyncHttpSurrealConnection]:
+        """
+        Acquire a connection from the pool as a context manager.
+
+        Gets an authenticated client from the pool, yields it for use,
+        and automatically returns it to the pool when done.
+
+        Yields:
+            AsyncWsSurrealConnection | AsyncHttpSurrealConnection: An authenticated client
+
+        Example:
+            >>> async with pool.acquire() as client:
+            ...     results = await client.select("person")
+
+        Raises:
+            Exception: If the acquired client encounters a fatal error
+        """
+        client = await self._queue.get()
+        try:
+            yield client
+        except Exception as exc:
+            # If the client had a fatal error, attempt to reconnect/replace it
+            # Simple strategy: if exception looks like connection/auth error -> recreate
+            if self._is_fatal_client_exc(exc):
+                logger.warning("Client had fatal error, replacing: %s", exc)
+                await self._replace_client(client)
+                # re-raise so caller can handle retry logic
+                raise
+            else:
+                raise
+        finally:
+            # If pool is closing, don't put back
+            if not self._closed:
+                self._queue.put_nowait(client)
+
+    async def run_query(self, *args, retries: int = 1, backoff: float = 0.2, **kwargs):
+        """
+        Execute a query using a pooled connection with automatic retry.
+
+        Acquires a client from the pool, executes the query, and handles
+        retries for transient failures.
+
+        Args:
+            *args: Arguments to pass to the query method
+            retries: Number of retry attempts for transient errors (default: 1)
+            backoff: Base backoff time in seconds between retries (default: 0.2)
+            **kwargs: Keyword arguments to pass to the query method
+
+        Returns:
+            Query results from SurrealDB
+
+        Raises:
+            Exception: If the query fails after all retry attempts
+        """
+        for attempt in range(retries + 1):
+            async with self.acquire() as client:
+                try:
+                    # Replace with client's actual query/call method
+                    return await client.query(*args, **kwargs)
+                except Exception as exc:
+                    logger.exception("Query failed on attempt %d: %s", attempt, exc)
+                    if attempt < retries and self._is_retryable(exc):
+                        await asyncio.sleep(backoff * (2 ** attempt))
+                        continue
+                    # If fatal to client, replace it so pool remains healthy
+                    if self._is_fatal_client_exc(exc):
+                        await self._replace_client(client)
+                    raise
+
+    async def _replace_client(self, dead_client):
+        """
+        Replace a failed client with a new authenticated connection.
+
+        Closes the dead client, creates a new one, and adds it back to the pool
+        to maintain the pool size.
+
+        Args:
+            dead_client: The client connection that failed
+        """
+        try:
+            await dead_client.close()
+        except Exception:
+            pass
+        try:
+            new_client = await self._make_one()
+        except Exception as e:
+            logger.exception("Failed to create replacement client: %s", e)
+            # If we can't create a replacement immediately, put back the old one to avoid starvation.
+            # In production, you might want to implement more sophisticated retry/backoff + alerts.
+            self._queue.put_nowait(dead_client)
+            return
+        # replace in internal list
+        try:
+            idx = self._clients.index(dead_client)
+            self._clients[idx] = new_client
+        except ValueError:
+            self._clients.append(new_client)
+        self._queue.put_nowait(new_client)
+
+    def _is_retryable(self, exc: Exception) -> bool:
+        """
+        Determine if an exception is retryable.
+
+        Args:
+            exc: The exception to check
+
+        Returns:
+            bool: True if the error is transient and worth retrying
+        """
+        return isinstance(exc, (asyncio.TimeoutError, ConnectionError))
+
+    def _is_fatal_client_exc(self, exc: Exception) -> bool:
+        """
+        Determine if an exception indicates a broken client connection.
+
+        Args:
+            exc: The exception to check
+
+        Returns:
+            bool: True if the client connection is likely broken and needs replacement
+        """
+        return isinstance(exc, (ConnectionError, RuntimeError))
+
+
+
+def create_engine(
+    url: str,
+    auth_payload: dict,
+    pool_size: int = 4,
+    connect_timeout: float = 5.
+) -> AsyncGenerator[AsyncWsSurrealConnection | AsyncHttpSurrealConnection, None]:
+    """
+    Create a connection pool for SurrealDB with automatic connection management.
+
+    This function creates a pool of authenticated SurrealDB connections that can
+    be reused across your application. It returns a context manager that handles
+    acquiring and releasing connections automatically.
+
+    Args:
+        url: SurrealDB connection URL (e.g., "ws://localhost:8000/rpc")
+        auth_payload: Dictionary with authentication credentials
+            Example: {"username": "root", "password": "root"}
+        pool_size: Number of connections to maintain in the pool (default: 4)
+        connect_timeout: Timeout for connection attempts in seconds (default: 5.0)
+
+    Returns:
+        AsyncGenerator: A context manager for acquiring connections from the pool
+
+    Example:
+        >>> # Create an engine
+        >>> engine = create_engine(
+        ...     "ws://localhost:8000/rpc",
+        ...     {"username": "root", "password": "root"},
+        ...     pool_size=5
+        ... )
+        >>>
+        >>> # Use the engine to get connections
+        >>> async with engine() as db:
+        ...     await db.use("mydatabase", "mynamespace")
+        ...     people = await Person.insert(db, [
+        ...         Person(name="Alice"),
+        ...         Person(name="Bob")
+        ...     ])
+
+    Notes:
+        - The pool is not started automatically; you need to call pool.start()
+        - Connections are reused for better performance
+        - Failed connections are automatically replaced
+        - The pool should be closed when your application shuts down
+
+    See Also:
+        - SurrealClientPool: The underlying pool implementation
+        - Node.insert(): For batch operations using pooled connections
+        - Edge.relate(): For creating relationships using pooled connections
+    """
+    pool = SurrealClientPool(url, auth_payload, pool_size, connect_timeout)
+    return pool.acquire