surrealdb-orm 0.1.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

surreal_sdk/connection/websocket.py

@@ -0,0 +1,519 @@
+"""
+WebSocket Connection Implementation for SurrealDB SDK.
+
+Provides stateful WebSocket-based connection for real-time features.
+"""
+
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
+import asyncio
+import json
+
+import aiohttp
+from aiohttp import ClientWSTimeout
+
+from .base import BaseSurrealConnection
+
+if TYPE_CHECKING:
+    from ..transaction import WebSocketTransaction
+    from ..streaming.live_select import LiveSelectStream, LiveSubscriptionParams
+from ..protocol.rpc import RPCRequest, RPCResponse
+from ..exceptions import ConnectionError, LiveQueryError, TimeoutError
+
+
+# Type alias for live query callbacks
+LiveCallback = Callable[[dict[str, Any]], Coroutine[Any, Any, None]]
+
+
+class WebSocketConnection(BaseSurrealConnection):
+    """
+    WebSocket-based connection to SurrealDB.
+
+    This connection is stateful - session is maintained across requests.
+    Required for Live Queries and session variables.
+    Ideal for real-time applications, dashboards, and collaborative features.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        namespace: str,
+        database: str,
+        timeout: float = 30.0,
+        auto_reconnect: bool = True,
+        reconnect_interval: float = 1.0,
+        max_reconnect_attempts: int = 5,
+    ):
+        """
+        Initialize WebSocket connection.
+
+        Args:
+            url: SurrealDB WebSocket URL (e.g., "ws://localhost:8000")
+            namespace: Target namespace
+            database: Target database
+            timeout: Request timeout in seconds
+            auto_reconnect: Whether to automatically reconnect on disconnect
+            reconnect_interval: Seconds between reconnection attempts
+            max_reconnect_attempts: Maximum reconnection attempts
+        """
+        # Normalize URL to WebSocket
+        if url.startswith("http://"):
+            url = url.replace("http://", "ws://", 1)
+        elif url.startswith("https://"):
+            url = url.replace("https://", "wss://", 1)
+
+        # Ensure /rpc suffix
+        if not url.endswith("/rpc"):
+            url = url.rstrip("/") + "/rpc"
+
+        super().__init__(url, namespace, database, timeout)
+
+        self.auto_reconnect = auto_reconnect
+        self.reconnect_interval = reconnect_interval
+        self.max_reconnect_attempts = max_reconnect_attempts
+
+        self._session: aiohttp.ClientSession | None = None
+        self._ws: aiohttp.ClientWebSocketResponse | None = None
+        self._request_id = 0
+        self._pending: dict[int, asyncio.Future[RPCResponse]] = {}
+        self._live_callbacks: dict[str, LiveCallback] = {}
+        self._live_subscriptions: dict[str, "LiveSubscriptionParams"] = {}
+        self._reader_task: asyncio.Task[None] | None = None
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._closing = False
+
+    def _next_request_id(self) -> int:
+        """Generate next request ID."""
+        self._request_id += 1
+        return self._request_id
+
+    async def connect(self) -> None:
+        """Establish WebSocket connection."""
+        if self._connected:
+            return
+
+        self._closing = False
+        self._session = aiohttp.ClientSession()
+
+        try:
+            # Specify protocol format explicitly to avoid deprecation warning
+            # SurrealDB 2.0+ requires explicit format specification (json or cbor)
+            self._ws = await self._session.ws_connect(
+                self.url,
+                timeout=ClientWSTimeout(ws_close=self.timeout),
+                protocols=["json"],  # Explicit JSON format for RPC
+            )
+            self._connected = True
+
+            # Start message reader loop
+            self._reader_task = asyncio.create_task(self._read_loop())
+
+            # Yield to event loop to allow read loop to start
+            await asyncio.sleep(0)
+
+            # Set namespace and database
+            await self.use(self.namespace, self.database)
+
+        except aiohttp.ClientError as e:
+            await self._cleanup()
+            raise ConnectionError(f"WebSocket connection failed: {e}")
+
+    async def close(self) -> None:
+        """Close WebSocket connection."""
+        self._closing = True
+        await self._cleanup()
+
+    async def _cleanup(self) -> None:
+        """Clean up connection resources."""
+        self._connected = False
+        self._authenticated = False
+
+        # Cancel reader task
+        if self._reader_task:
+            self._reader_task.cancel()
+            try:
+                await self._reader_task
+            except asyncio.CancelledError:
+                pass
+            self._reader_task = None
+
+        # Cancel reconnect task
+        if self._reconnect_task:
+            self._reconnect_task.cancel()
+            try:
+                await self._reconnect_task
+            except asyncio.CancelledError:
+                pass
+            self._reconnect_task = None
+
+        # Fail all pending requests
+        for future in self._pending.values():
+            if not future.done():
+                future.set_exception(ConnectionError("Connection closed"))
+        self._pending.clear()
+
+        # Close WebSocket
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+
+        # Close session
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+    async def _read_loop(self) -> None:
+        """Background task to read WebSocket messages."""
+        if not self._ws:
+            return
+
+        try:
+            async for msg in self._ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    await self._handle_message(msg.data)
+                elif msg.type == aiohttp.WSMsgType.BINARY:
+                    # CBOR support could be added here
+                    pass
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    break
+                elif msg.type == aiohttp.WSMsgType.CLOSED:
+                    break
+                elif msg.type == aiohttp.WSMsgType.CLOSE:
+                    break
+
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            pass
+        finally:
+            if not self._closing and self.auto_reconnect:
+                self._reconnect_task = asyncio.create_task(self._reconnect())
+
+    async def _handle_message(self, data: str) -> None:
+        """Handle incoming WebSocket message."""
+        try:
+            message = json.loads(data)
+        except json.JSONDecodeError:
+            return
+
+        msg_id = message.get("id")
+
+        # Check if this is a response to a pending request
+        # Convert msg_id to int for comparison since SurrealDB may return string IDs
+        if msg_id is not None:
+            try:
+                msg_id_int = int(msg_id)
+            except (ValueError, TypeError):
+                msg_id_int = None
+
+            if msg_id_int is not None and msg_id_int in self._pending:
+                response = RPCResponse.from_dict(message)
+                future = self._pending.pop(msg_id_int)
+                if not future.done():
+                    future.set_result(response)
+                return
+
+        # Check if this is a live query notification
+        if "action" in message:
+            live_id = message.get("id")
+            if live_id and live_id in self._live_callbacks:
+                callback = self._live_callbacks[live_id]
+                asyncio.create_task(callback(message))
+
+    async def _reconnect(self) -> None:
+        """Attempt to reconnect after disconnection."""
+        attempts = 0
+        while attempts < self.max_reconnect_attempts and not self._closing:
+            attempts += 1
+            await asyncio.sleep(self.reconnect_interval)
+
+            try:
+                self._session = aiohttp.ClientSession()
+                self._ws = await self._session.ws_connect(
+                    self.url,
+                    timeout=ClientWSTimeout(ws_close=self.timeout),
+                    protocols=["json"],
+                )
+                self._connected = True
+                self._reader_task = asyncio.create_task(self._read_loop())
+
+                # Re-authenticate if we had a token
+                if self._token:
+                    await self.rpc("authenticate", [self._token])
+
+                # Set namespace and database
+                await self.use(self.namespace, self.database)
+
+                # Re-establish live queries with auto-resubscribe
+                await self._resubscribe_all()
+
+                return
+
+            except Exception:
+                await self._cleanup()
+                continue
+
+    async def _resubscribe_all(self) -> None:
+        """Re-establish all live subscriptions after reconnect."""
+        old_subscriptions = dict(self._live_subscriptions)
+        self._live_subscriptions.clear()
+        self._live_callbacks.clear()
+
+        for old_id, params in old_subscriptions.items():
+            try:
+                new_id = await self._resubscribe_one(params)
+
+                # Call reconnect callback if provided
+                if params.on_reconnect:
+                    asyncio.create_task(params.on_reconnect(old_id, new_id))
+
+            except Exception:
+                # Failed to resubscribe, skip this one
+                pass
+
+    async def _resubscribe_one(self, params: "LiveSubscriptionParams") -> str:
+        """Resubscribe a single live query."""
+        # Set session variables for parameters
+        for key, value in params.params.items():
+            await self.let(key, value)
+
+        # Build query
+        sql = f"LIVE SELECT * FROM {params.table}"
+        if params.where:
+            sql += f" WHERE {params.where}"
+        if params.diff:
+            sql += " DIFF"
+
+        response = await self.query(sql)
+
+        if response.results:
+            first_result = response.results[0]
+            if first_result.is_ok:
+                result_data = first_result.result
+                if isinstance(result_data, str):
+                    new_id = result_data
+                elif isinstance(result_data, dict) and "result" in result_data:
+                    new_id = str(result_data["result"])
+                else:
+                    raise LiveQueryError("Invalid live query response")
+
+                # Re-register callback if provided
+                if params.callback:
+                    self._live_callbacks[new_id] = params.callback  # type: ignore[assignment]
+
+                # Store subscription params for future reconnects
+                self._live_subscriptions[new_id] = params
+
+                return new_id
+
+        raise LiveQueryError("No live query ID returned")
+
+    def _register_live_subscription(self, live_id: str, params: "LiveSubscriptionParams") -> None:
+        """Register a live subscription for auto-resubscribe."""
+        self._live_subscriptions[live_id] = params
+
+    def _unregister_live_subscription(self, live_id: str) -> None:
+        """Unregister a live subscription."""
+        self._live_subscriptions.pop(live_id, None)
+
+    async def _send_rpc(self, request: RPCRequest) -> RPCResponse:
+        """
+        Send RPC request via WebSocket.
+
+        Args:
+            request: The RPC request to send
+
+        Returns:
+            The RPC response
+
+        Raises:
+            ConnectionError: If not connected
+            TimeoutError: If request times out
+        """
+        if not self._ws or not self._connected:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        request.id = self._next_request_id()
+
+        # Create future for response
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future[RPCResponse] = loop.create_future()
+        self._pending[request.id] = future
+
+        try:
+            # Send request
+            await self._ws.send_str(request.to_json())
+
+            # Wait for response with timeout
+            response = await asyncio.wait_for(future, timeout=self.timeout)
+            return response
+
+        except asyncio.TimeoutError:
+            self._pending.pop(request.id, None)
+            raise TimeoutError(f"Request timed out after {self.timeout}s")
+        except Exception as e:
+            self._pending.pop(request.id, None)
+            raise ConnectionError(f"Request failed: {e}")
+
+    # WebSocket-specific methods
+
+    async def live(
+        self,
+        table: str,
+        callback: LiveCallback,
+        diff: bool = False,
+    ) -> str:
+        """
+        Start a live query subscription.
+
+        Args:
+            table: Table to watch
+            callback: Async callback for change notifications
+            diff: If True, receive only changed fields
+
+        Returns:
+            Live query UUID
+
+        Raises:
+            LiveQueryError: If live query fails to start
+        """
+        sql = f"LIVE SELECT * FROM {table}"
+        if diff:
+            sql += " DIFF"
+
+        try:
+            response = await self.query(sql)
+
+            # Extract live query UUID from result
+            if response.results:
+                first_result = response.results[0]
+                if first_result.is_ok:
+                    # Result can be string UUID directly or dict with "result" key
+                    if isinstance(first_result.result, str):
+                        live_id = first_result.result
+                    elif isinstance(first_result.result, dict) and "result" in first_result.result:
+                        live_id = str(first_result.result["result"])
+                    else:
+                        raise LiveQueryError("Invalid live query response")
+
+                    self._live_callbacks[live_id] = callback
+                    return live_id
+
+            raise LiveQueryError("No live query ID returned")
+
+        except Exception as e:
+            raise LiveQueryError(f"Failed to start live query: {e}")
+
+    async def kill(self, live_id: str) -> None:
+        """
+        Stop a live query subscription.
+
+        Args:
+            live_id: Live query UUID to stop
+        """
+        await self.rpc("kill", [live_id])
+        self._live_callbacks.pop(live_id, None)
+
+    async def let(self, name: str, value: Any) -> None:
+        """
+        Set a session variable.
+
+        Args:
+            name: Variable name
+            value: Variable value
+        """
+        await self.rpc("let", [name, value])
+
+    async def unset(self, name: str) -> None:
+        """
+        Remove a session variable.
+
+        Args:
+            name: Variable name to remove
+        """
+        await self.rpc("unset", [name])
+
+    @property
+    def live_queries(self) -> list[str]:
+        """Get list of active live query IDs."""
+        return list(self._live_callbacks.keys())
+
+    async def kill_all_live_queries(self) -> None:
+        """Stop all active live queries."""
+        for live_id in list(self._live_callbacks.keys()):
+            try:
+                await self.kill(live_id)
+            except Exception:
+                pass
+
+    # Transaction support
+
+    def transaction(self) -> "WebSocketTransaction":
+        """
+        Create a new WebSocket transaction.
+
+        WebSocket transactions use server-side state with BEGIN/COMMIT/ROLLBACK.
+        Operations are executed immediately within the transaction context.
+
+        Usage:
+            async with conn.transaction() as tx:
+                await tx.create("users", {"name": "Alice"})
+                await tx.create("orders", {"user": "users:alice"})
+                # Committed on successful exit, rolled back on exception
+
+        Returns:
+            WebSocketTransaction context manager
+        """
+        from ..transaction import WebSocketTransaction
+
+        return WebSocketTransaction(self)
+
+    # Live Select Stream API
+
+    def live_select(
+        self,
+        table: str,
+        where: str | None = None,
+        params: dict[str, Any] | None = None,
+        diff: bool = False,
+        auto_resubscribe: bool = True,
+        on_reconnect: Callable[[str, str], Coroutine[Any, Any, None]] | None = None,
+    ) -> "LiveSelectStream":
+        """
+        Create a live select stream for real-time change notifications.
+
+        This method returns an async iterator that yields LiveChange objects
+        whenever records matching the query are created, updated, or deleted.
+
+        Args:
+            table: Table to watch (e.g., "players", "game_tables")
+            where: Optional WHERE clause filter (e.g., "table_id = $id")
+            params: Parameters for the WHERE clause (e.g., {"id": "game_tables:xyz"})
+            diff: If True, receive only changed fields
+            auto_resubscribe: If True, automatically resubscribe on reconnect
+            on_reconnect: Optional callback when resubscribed (old_id, new_id)
+
+        Returns:
+            LiveSelectStream async context manager and iterator
+
+        Usage:
+            async with conn.live_select("players", where="table_id = $id", params={"id": table_id}) as stream:
+                async for change in stream:
+                    match change.action:
+                        case LiveAction.CREATE:
+                            print(f"New player: {change.result}")
+                        case LiveAction.UPDATE:
+                            print(f"Player updated: {change.record_id}")
+                        case LiveAction.DELETE:
+                            print(f"Player left: {change.record_id}")
+        """
+        from ..streaming.live_select import LiveSelectStream
+
+        return LiveSelectStream(
+            connection=self,
+            table=table,
+            where=where,
+            params=params,
+            diff=diff,
+            auto_resubscribe=auto_resubscribe,
+            on_reconnect=on_reconnect,
+        )

surreal_sdk/exceptions.py

@@ -0,0 +1,71 @@
+"""
+SurrealDB SDK Exceptions.
+
+Custom exception hierarchy for the SDK.
+"""
+
+
+class SurrealDBError(Exception):
+    """Base exception for all SurrealDB SDK errors."""
+
+    def __init__(self, message: str, code: int | None = None):
+        self.message = message
+        self.code = code
+        super().__init__(message)
+
+
+class ConnectionError(SurrealDBError):
+    """Raised when connection to SurrealDB fails."""
+
+    pass
+
+
+class AuthenticationError(SurrealDBError):
+    """Raised when authentication fails."""
+
+    pass
+
+
+class QueryError(SurrealDBError):
+    """Raised when a query execution fails."""
+
+    def __init__(self, message: str, query: str | None = None, code: int | None = None):
+        self.query = query
+        super().__init__(message, code)
+
+
+class TimeoutError(SurrealDBError):
+    """Raised when an operation times out."""
+
+    pass
+
+
+class ValidationError(SurrealDBError):
+    """Raised when data validation fails."""
+
+    pass
+
+
+class LiveQueryError(SurrealDBError):
+    """Raised when a live query operation fails."""
+
+    pass
+
+
+class ChangeFeedError(SurrealDBError):
+    """Raised when a change feed operation fails."""
+
+    pass
+
+
+class TransactionError(SurrealDBError):
+    """Raised when a transaction operation fails."""
+
+    def __init__(
+        self,
+        message: str,
+        code: int | None = None,
+        rollback_succeeded: bool | None = None,
+    ):
+        self.rollback_succeeded = rollback_succeeded
+        super().__init__(message, code)