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

surreal_sdk/connection/http.py

@@ -0,0 +1,421 @@
+"""
+HTTP Connection Implementation for SurrealDB SDK.
+
+Provides stateless HTTP-based connection, ideal for microservices and serverless.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from .base import BaseSurrealConnection
+
+if TYPE_CHECKING:
+    from ..transaction import HTTPTransaction
+from ..protocol.rpc import RPCRequest, RPCResponse
+from ..types import AuthResponse
+from ..exceptions import ConnectionError, QueryError
+
+
+class HTTPConnection(BaseSurrealConnection):
+    """
+    HTTP-based connection to SurrealDB.
+
+    This connection is stateless - each request is independent.
+    Ideal for microservices, serverless, and horizontally scaled applications.
+
+    Authentication is performed via headers on each request.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        namespace: str,
+        database: str,
+        timeout: float = 30.0,
+    ):
+        """
+        Initialize HTTP connection.
+
+        Args:
+            url: SurrealDB HTTP URL (e.g., "http://localhost:8000")
+            namespace: Target namespace
+            database: Target database
+            timeout: Request timeout in seconds
+        """
+        # Normalize URL to HTTP if needed
+        if url.startswith("ws://"):
+            url = url.replace("ws://", "http://", 1)
+        elif url.startswith("wss://"):
+            url = url.replace("wss://", "https://", 1)
+
+        super().__init__(url, namespace, database, timeout)
+        self._client: httpx.AsyncClient | None = None
+        self._request_id = 0
+
+    @property
+    def headers(self) -> dict[str, str]:
+        """Build request headers."""
+        h = {
+            "Surreal-NS": self.namespace,
+            "Surreal-DB": self.database,
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+        if self._token:
+            h["Authorization"] = f"Bearer {self._token}"
+        return h
+
+    def _next_request_id(self) -> int:
+        """Generate next request ID."""
+        self._request_id += 1
+        return self._request_id
+
+    async def connect(self) -> None:
+        """Establish HTTP client connection."""
+        if self._connected:
+            return
+
+        self._client = httpx.AsyncClient(
+            base_url=self.url,
+            timeout=self.timeout,
+            # Disable connection pooling to avoid event loop binding issues in tests
+            limits=httpx.Limits(max_keepalive_connections=0, max_connections=100),
+        )
+        self._connected = True
+
+    async def close(self) -> None:
+        """Close HTTP client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+        self._connected = False
+        self._authenticated = False
+
+    async def _send_rpc(self, request: RPCRequest) -> RPCResponse:
+        """
+        Send RPC request via HTTP POST to /rpc endpoint.
+
+        Args:
+            request: The RPC request to send
+
+        Returns:
+            The RPC response
+
+        Raises:
+            ConnectionError: If not connected
+            QueryError: If request fails
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        request.id = self._next_request_id()
+
+        try:
+            response = await self._client.post(
+                "/rpc",
+                json=request.to_dict(),
+                headers=self.headers,
+            )
+            response.raise_for_status()
+            return RPCResponse.from_dict(response.json())
+
+        except httpx.HTTPStatusError as e:
+            raise QueryError(
+                message=f"HTTP error: {e.response.status_code} - {e.response.text}",
+                code=e.response.status_code,
+            )
+        except httpx.RequestError as e:
+            raise ConnectionError(f"Request failed: {e}")
+
+    async def signin(
+        self,
+        user: str | None = None,
+        password: str | None = None,
+        namespace: str | None = None,
+        database: str | None = None,
+        access: str | None = None,
+        **credentials: Any,
+    ) -> AuthResponse:
+        """
+        Authenticate with SurrealDB via HTTP.
+
+        For HTTP connections, this obtains a JWT token that will be
+        included in subsequent request headers.
+
+        Args:
+            user: Username (for root/namespace/database auth)
+            password: Password (for root/namespace/database auth)
+            namespace: Optional namespace scope
+            database: Optional database scope
+            access: Optional access method (for record access auth)
+            **credentials: Additional credentials for record access (email, password, etc.)
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        from ..exceptions import AuthenticationError
+
+        payload: dict[str, Any] = {}
+        if user:
+            payload["user"] = user
+        if namespace:
+            payload["ns"] = namespace
+        if database:
+            payload["db"] = database
+        if access:
+            # Record access auth: password goes as 'password' in credentials
+            payload["ac"] = access
+            if password:
+                payload["password"] = password
+        else:
+            # Root/namespace/database auth: password goes as 'pass'
+            if password:
+                payload["pass"] = password
+        # Add any additional credentials for record access
+        payload.update(credentials)
+
+        try:
+            response = await self._client.post(
+                "/signin",
+                json=payload,
+                headers={"Accept": "application/json", "Content-Type": "application/json"},
+            )
+
+            if response.status_code != 200:
+                raise AuthenticationError(f"Authentication failed: {response.text}")
+
+            data = response.json()
+            token = data.get("token")
+            self._token = token
+            self._authenticated = True
+            return AuthResponse(token=token, success=True, raw=data)
+
+        except httpx.RequestError as e:
+            raise AuthenticationError(f"Authentication request failed: {e}")
+
+    async def sql(self, query: str, vars: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """
+        Execute raw SurrealQL via POST /sql endpoint.
+
+        This is a direct SQL execution endpoint, alternative to RPC.
+
+        Args:
+            query: SurrealQL query string
+            vars: Query variables (passed as query params)
+
+        Returns:
+            Query results
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        try:
+            response = await self._client.post(
+                "/sql",
+                content=query,
+                headers=self.headers,
+                params=vars,
+            )
+            response.raise_for_status()
+            result: list[dict[str, Any]] = response.json()
+            return result
+
+        except httpx.HTTPStatusError as e:
+            raise QueryError(
+                message=f"SQL query failed: {e.response.text}",
+                query=query,
+                code=e.response.status_code,
+            )
+
+    async def health(self) -> bool:
+        """
+        Check server health via GET /health endpoint.
+
+        Returns:
+            True if server is healthy
+        """
+        if not self._client:
+            return False
+
+        try:
+            response = await self._client.get("/health")
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    async def status(self) -> bool:
+        """
+        Check server status via GET /status endpoint.
+
+        Returns:
+            True if server is running
+        """
+        if not self._client:
+            return False
+
+        try:
+            response = await self._client.get("/status")
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    # REST-style CRUD endpoints (alternative to RPC)
+
+    async def rest_select(self, table: str, record_id: str | None = None) -> list[dict[str, Any]]:
+        """
+        Select via REST GET /key/:table or /key/:table/:id.
+
+        Args:
+            table: Table name
+            record_id: Optional record ID
+
+        Returns:
+            Records
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        path = f"/key/{table}"
+        if record_id:
+            path += f"/{record_id}"
+
+        response = await self._client.get(path, headers=self.headers)
+        response.raise_for_status()
+        result = response.json()
+        return result if isinstance(result, list) else [result]
+
+    async def rest_create(
+        self,
+        table: str,
+        data: dict[str, Any],
+        record_id: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Create via REST POST /key/:table or /key/:table/:id.
+
+        Args:
+            table: Table name
+            data: Record data
+            record_id: Optional record ID
+
+        Returns:
+            Created record
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        path = f"/key/{table}"
+        if record_id:
+            path += f"/{record_id}"
+
+        response = await self._client.post(path, json=data, headers=self.headers)
+        response.raise_for_status()
+        result: dict[str, Any] = response.json()
+        return result
+
+    async def rest_update(
+        self,
+        table: str,
+        record_id: str,
+        data: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Update via REST PUT /key/:table/:id.
+
+        Args:
+            table: Table name
+            record_id: Record ID
+            data: New record data
+
+        Returns:
+            Updated record
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        response = await self._client.put(
+            f"/key/{table}/{record_id}",
+            json=data,
+            headers=self.headers,
+        )
+        response.raise_for_status()
+        result: dict[str, Any] = response.json()
+        return result
+
+    async def rest_patch(
+        self,
+        table: str,
+        record_id: str,
+        data: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Patch via REST PATCH /key/:table/:id.
+
+        Args:
+            table: Table name
+            record_id: Record ID
+            data: Fields to update
+
+        Returns:
+            Updated record
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        response = await self._client.patch(
+            f"/key/{table}/{record_id}",
+            json=data,
+            headers=self.headers,
+        )
+        response.raise_for_status()
+        result: dict[str, Any] = response.json()
+        return result
+
+    async def rest_delete(
+        self,
+        table: str,
+        record_id: str | None = None,
+    ) -> dict[str, Any] | list[dict[str, Any]]:
+        """
+        Delete via REST DELETE /key/:table or /key/:table/:id.
+
+        Args:
+            table: Table name
+            record_id: Optional record ID
+
+        Returns:
+            Deleted record(s)
+        """
+        if not self._client:
+            raise ConnectionError("Not connected. Call connect() first.")
+
+        path = f"/key/{table}"
+        if record_id:
+            path += f"/{record_id}"
+
+        response = await self._client.delete(path, headers=self.headers)
+        response.raise_for_status()
+        result: dict[str, Any] | list[dict[str, Any]] = response.json()
+        return result
+
+    # Transaction support
+
+    def transaction(self) -> "HTTPTransaction":
+        """
+        Create a new HTTP transaction.
+
+        HTTP transactions batch all statements and execute them atomically on commit.
+
+        Usage:
+            async with conn.transaction() as tx:
+                await tx.create("users", {"name": "Alice"})
+                await tx.create("orders", {"user": "users:alice"})
+                # All statements executed atomically on exit
+
+        Returns:
+            HTTPTransaction context manager
+        """
+        from ..transaction import HTTPTransaction
+
+        return HTTPTransaction(self)

surreal_sdk/connection/pool.py

@@ -0,0 +1,244 @@
+"""
+Connection Pool Implementation for SurrealDB SDK.
+
+Provides connection pooling for both HTTP and WebSocket connections.
+"""
+
+from typing import Any, AsyncGenerator, Self
+import asyncio
+from contextlib import asynccontextmanager
+from collections import deque
+
+from .base import BaseSurrealConnection
+from .http import HTTPConnection
+from .websocket import WebSocketConnection
+from ..types import QueryResponse, RecordResponse, RecordsResponse, DeleteResponse
+from ..exceptions import ConnectionError
+
+
+class ConnectionPool:
+    """
+    Connection pool for SurrealDB connections.
+
+    Manages a pool of reusable connections for improved performance
+    in high-throughput scenarios.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        namespace: str,
+        database: str,
+        size: int = 10,
+        connection_type: str = "http",
+        timeout: float = 30.0,
+        **kwargs: Any,
+    ):
+        """
+        Initialize connection pool.
+
+        Args:
+            url: SurrealDB server URL
+            namespace: Target namespace
+            database: Target database
+            size: Maximum pool size
+            connection_type: "http" or "websocket"
+            timeout: Connection timeout in seconds
+            **kwargs: Additional connection arguments
+        """
+        self.url = url
+        self.namespace = namespace
+        self.database = database
+        self.size = size
+        self.connection_type = connection_type
+        self.timeout = timeout
+        self.kwargs = kwargs
+
+        self._pool: deque[BaseSurrealConnection] = deque()
+        self._in_use: set[BaseSurrealConnection] = set()
+        self._lock = asyncio.Lock()
+        self._closed = False
+        self._credentials: tuple[str, str] | None = None
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+    def _create_connection(self) -> BaseSurrealConnection:
+        """Create a new connection instance."""
+        if self.connection_type == "websocket":
+            return WebSocketConnection(
+                self.url,
+                self.namespace,
+                self.database,
+                timeout=self.timeout,
+                **self.kwargs,
+            )
+        else:
+            return HTTPConnection(
+                self.url,
+                self.namespace,
+                self.database,
+                timeout=self.timeout,
+            )
+
+    async def _init_connection(self, conn: BaseSurrealConnection) -> None:
+        """Initialize a connection."""
+        await conn.connect()
+        if self._credentials:
+            user, password = self._credentials
+            await conn.signin(user, password)
+
+    async def set_credentials(self, user: str, password: str) -> None:
+        """
+        Set credentials for all pool connections.
+
+        Args:
+            user: Username
+            password: Password
+        """
+        self._credentials = (user, password)
+
+        # Re-authenticate existing connections
+        async with self._lock:
+            for conn in self._pool:
+                try:
+                    await conn.signin(user, password)
+                except Exception:
+                    pass
+
+    @asynccontextmanager
+    async def acquire(self) -> AsyncGenerator[BaseSurrealConnection, None]:
+        """
+        Acquire a connection from the pool.
+
+        Usage:
+            async with pool.acquire() as conn:
+                result = await conn.query("SELECT * FROM users")
+
+        Yields:
+            A SurrealDB connection
+        """
+        if self._closed:
+            raise ConnectionError("Pool is closed")
+
+        conn: BaseSurrealConnection | None = None
+
+        async with self._lock:
+            # Try to get an existing connection from pool
+            while self._pool:
+                conn = self._pool.popleft()
+                if conn.is_connected:
+                    break
+                # Connection is dead, discard it
+                try:
+                    await conn.close()
+                except Exception:
+                    pass
+                conn = None
+
+            # Create new connection if needed and pool not at capacity
+            if conn is None:
+                if len(self._in_use) < self.size:
+                    conn = self._create_connection()
+                    await self._init_connection(conn)
+                else:
+                    # Pool at capacity, wait for a connection
+                    pass
+
+            if conn:
+                self._in_use.add(conn)
+
+        if conn is None:
+            # Wait for a connection to become available
+            while conn is None:
+                await asyncio.sleep(0.01)
+                async with self._lock:
+                    if self._pool:
+                        conn = self._pool.popleft()
+                        self._in_use.add(conn)
+
+        try:
+            yield conn
+        finally:
+            async with self._lock:
+                self._in_use.discard(conn)
+                if not self._closed and conn.is_connected:
+                    self._pool.append(conn)
+                else:
+                    try:
+                        await conn.close()
+                    except Exception:
+                        pass
+
+    async def close(self) -> None:
+        """Close all connections in the pool."""
+        self._closed = True
+
+        async with self._lock:
+            # Close pooled connections
+            while self._pool:
+                conn = self._pool.popleft()
+                try:
+                    await conn.close()
+                except Exception:
+                    pass
+
+            # Close in-use connections
+            for conn in self._in_use:
+                try:
+                    await conn.close()
+                except Exception:
+                    pass
+            self._in_use.clear()
+
+    @property
+    def available(self) -> int:
+        """Number of available connections in pool."""
+        return len(self._pool)
+
+    @property
+    def in_use(self) -> int:
+        """Number of connections currently in use."""
+        return len(self._in_use)
+
+    @property
+    def total(self) -> int:
+        """Total number of connections (available + in use)."""
+        return len(self._pool) + len(self._in_use)
+
+    # Convenience methods that acquire a connection
+
+    async def query(self, sql: str, vars: dict[str, Any] | None = None) -> QueryResponse:
+        """Execute a query using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.query(sql, vars)
+
+    async def select(self, thing: str) -> RecordsResponse:
+        """Select records using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.select(thing)
+
+    async def create(self, thing: str, data: dict[str, Any] | None = None) -> RecordResponse:
+        """Create a record using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.create(thing, data)
+
+    async def update(self, thing: str, data: dict[str, Any]) -> RecordsResponse:
+        """Update record(s) using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.update(thing, data)
+
+    async def merge(self, thing: str, data: dict[str, Any]) -> RecordsResponse:
+        """Merge data into record(s) using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.merge(thing, data)
+
+    async def delete(self, thing: str) -> DeleteResponse:
+        """Delete record(s) using a pooled connection."""
+        async with self.acquire() as conn:
+            return await conn.delete(thing)