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

surreal_orm/connection_manager.py

@@ -1,23 +1,31 @@
 from typing import Any
-from surrealdb import AsyncSurrealDB
-from surrealdb.errors import SurrealDbConnectionError
 import logging
 
+from surreal_sdk import HTTPConnection
+from surreal_sdk.exceptions import SurrealDBError
+from surreal_sdk.transaction import HTTPTransaction
+
 logger = logging.getLogger(__name__)
 
 
+class SurrealDbConnectionError(Exception):
+    """Connection error for SurrealDB."""
+
+    pass
+
+
 class SurrealDBConnectionManager:
     __url: str | None = None
     __user: str | None = None
     __password: str | None = None
     __namespace: str | None = None
     __database: str | None = None
-    __client: AsyncSurrealDB | None = None
+    __client: HTTPConnection | None = None
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await SurrealDBConnectionManager.close_connection()
 
-    async def __aenter__(self) -> AsyncSurrealDB:
+    async def __aenter__(self) -> HTTPConnection:
         return await SurrealDBConnectionManager.get_client()
 
     @classmethod
@@ -57,31 +65,43 @@ class SurrealDBConnectionManager:
         return all([cls.__url, cls.__user, cls.__password, cls.__namespace, cls.__database])
 
     @classmethod
-    async def get_client(cls) -> AsyncSurrealDB:
+    async def get_client(cls) -> HTTPConnection:
         """
-        Connect to the SurrealDB instance.
+        Connect to the SurrealDB instance using the custom SDK.
 
-        :return: The SurrealDB instance.
+        :return: The HTTPConnection instance.
         """
 
-        if cls.__client is not None:
+        if cls.__client is not None and cls.__client.is_connected:
             return cls.__client
 
         if not cls.is_connection_set():
             raise ValueError("Connection not been set.")
 
-        # Établir la connexion
+        # Establish connection
         try:
-            _client = AsyncSurrealDB(cls.get_connection_string())
-            await _client.connect()  # type: ignore
-            await _client.use(cls.__namespace, cls.__database)  # type: ignore
-            await _client.sign_in(cls.__user, cls.__password)  # type: ignore
+            url = cls.get_connection_string()
+            assert url is not None  # Already validated by is_connection_set()
+            assert cls.__namespace is not None
+            assert cls.__database is not None
+            assert cls.__user is not None
+            assert cls.__password is not None
+
+            _client = HTTPConnection(url, cls.__namespace, cls.__database)
+            await _client.connect()
+            await _client.signin(cls.__user, cls.__password)
 
             cls.__client = _client
             return cls.__client
+        except SurrealDBError as e:
+            logger.warning(f"Can't get connection: {e}")
+            if cls.__client is not None:  # pragma: no cover
+                await cls.__client.close()
+                cls.__client = None
+            raise SurrealDbConnectionError(f"Can't connect to the database: {e}")
         except Exception as e:
             logger.warning(f"Can't get connection: {e}")
-            if isinstance(cls.__client, AsyncSurrealDB):  # pragma: no cover
+            if cls.__client is not None:  # pragma: no cover
                 await cls.__client.close()
                 cls.__client = None
             raise SurrealDbConnectionError("Can't connect to the database.")
@@ -96,17 +116,19 @@ class SurrealDBConnectionManager:
         if cls.__client is None:
             return
 
-        await cls.__client.close()
+        try:
+            await cls.__client.close()
+        except NotImplementedError:
+            # close() is not implemented for HTTP connections in surrealdb SDK 1.0.8
+            pass
         cls.__client = None
 
     @classmethod
-    async def reconnect(cls) -> AsyncSurrealDB | None:
+    async def reconnect(cls) -> HTTPConnection | None:
         """
         Reconnect to the SurrealDB instance.
         """
-        # Fermer la connexion
         await cls.close_connection()
-        # Établir la connexion
         return await cls.get_client()
 
     @classmethod
@@ -299,3 +321,21 @@ class SurrealDBConnectionManager:
         """
 
         return cls.__client is not None
+
+    @classmethod
+    async def transaction(cls) -> HTTPTransaction:
+        """
+        Create a transaction context manager for atomic operations.
+
+        Usage:
+            async with SurrealDBConnectionManager.transaction() as tx:
+                user = User(name="Alice")
+                await user.save(tx=tx)
+                order = Order(user_id=user.id)
+                await order.save(tx=tx)
+                # Auto-commit on success, auto-rollback on exception
+
+        :return: HTTPTransaction context manager
+        """
+        client = await cls.get_client()
+        return client.transaction()

surreal_orm/fields/__init__.py

@@ -0,0 +1,36 @@
+"""
+Custom field types for SurrealDB ORM.
+
+This module provides specialized field types that leverage SurrealDB's
+built-in functions for encryption, validation, and other operations.
+"""
+
+from .encrypted import Encrypted, EncryptedField, EncryptedFieldInfo
+from .relation import (
+    ForeignKey,
+    ManyToMany,
+    Relation,
+    RelationInfo,
+    get_relation_info,
+    is_foreign_key,
+    is_graph_relation,
+    is_many_to_many,
+    is_relation_field,
+)
+
+__all__ = [
+    # Encrypted fields
+    "Encrypted",
+    "EncryptedField",
+    "EncryptedFieldInfo",
+    # Relation fields
+    "ForeignKey",
+    "ManyToMany",
+    "Relation",
+    "RelationInfo",
+    "get_relation_info",
+    "is_foreign_key",
+    "is_graph_relation",
+    "is_many_to_many",
+    "is_relation_field",
+]

surreal_orm/fields/encrypted.py

@@ -0,0 +1,166 @@
+"""
+Encrypted field type for password and sensitive data storage.
+
+This module provides the Encrypted[T] type that automatically generates
+the appropriate SurrealDB crypto functions in schema definitions.
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, Any, get_args, get_origin
+
+from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import CoreSchema, core_schema
+
+from ..types import EncryptionAlgorithm
+
+if TYPE_CHECKING:
+    pass
+
+
+@dataclass
+class EncryptedFieldInfo:
+    """
+    Metadata for encrypted fields used during schema generation.
+
+    Attributes:
+        algorithm: The encryption algorithm to use (default: argon2)
+        compare_function: SurrealDB function for password comparison
+        generate_function: SurrealDB function for password generation
+    """
+
+    algorithm: EncryptionAlgorithm = EncryptionAlgorithm.ARGON2
+
+    @property
+    def compare_function(self) -> str:
+        """Get the SurrealDB compare function for this algorithm."""
+        return f"crypto::{self.algorithm}::compare"
+
+    @property
+    def generate_function(self) -> str:
+        """Get the SurrealDB generate function for this algorithm."""
+        return f"crypto::{self.algorithm}::generate"
+
+
+class _EncryptedMarker:
+    """
+    Marker class for encrypted fields.
+
+    This is used internally to mark fields as encrypted in Pydantic's
+    type annotation system.
+    """
+
+    algorithm: EncryptionAlgorithm
+
+    def __init__(self, algorithm: EncryptionAlgorithm = EncryptionAlgorithm.ARGON2):
+        self.algorithm = algorithm
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls,
+        source_type: Any,
+        handler: GetCoreSchemaHandler,
+    ) -> CoreSchema:
+        """Build the Pydantic core schema for validation."""
+        # Get algorithm from marker if available
+        algorithm = EncryptionAlgorithm.ARGON2
+        args = get_args(source_type)
+        for arg in args:
+            if isinstance(arg, _EncryptedMarker):
+                algorithm = arg.algorithm
+                break
+
+        return core_schema.str_schema(
+            metadata={
+                "encrypted": True,
+                "algorithm": str(algorithm),
+                "surreal_type": "string",
+                "generate_function": f"crypto::{algorithm}::generate",
+                "compare_function": f"crypto::{algorithm}::compare",
+            }
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        core_schema_obj: CoreSchema,
+        handler: GetJsonSchemaHandler,
+    ) -> JsonSchemaValue:
+        """Generate JSON schema for OpenAPI documentation."""
+        return {"type": "string", "format": "password"}
+
+
+# Type alias for Encrypted[str] using Annotated
+Encrypted = Annotated[str, _EncryptedMarker()]
+
+
+def EncryptedField(
+    algorithm: EncryptionAlgorithm = EncryptionAlgorithm.ARGON2,
+) -> Any:
+    """
+    Create an encrypted field type with a specific algorithm.
+
+    Usage:
+        class User(BaseSurrealModel):
+            password: EncryptedField(EncryptionAlgorithm.BCRYPT)
+
+    Args:
+        algorithm: The encryption algorithm to use
+
+    Returns:
+        Annotated type with encryption marker
+    """
+    return Annotated[str, _EncryptedMarker(algorithm)]
+
+
+def is_encrypted_field(field_type: Any) -> bool:
+    """
+    Check if a field type is an Encrypted type.
+
+    Args:
+        field_type: The type annotation to check
+
+    Returns:
+        True if the field is an Encrypted type
+    """
+    # Check for Annotated type with _EncryptedMarker
+    origin = get_origin(field_type)
+
+    if origin is Annotated:
+        args = get_args(field_type)
+        for arg in args:
+            if isinstance(arg, _EncryptedMarker):
+                return True
+
+    # Also check for _EncryptedMarker class directly
+    if isinstance(field_type, type) and issubclass(field_type, _EncryptedMarker):
+        return True
+
+    return False
+
+
+def get_encryption_info(field_type: Any) -> EncryptedFieldInfo | None:
+    """
+    Extract encryption information from a field type.
+
+    Args:
+        field_type: The type annotation to extract from
+
+    Returns:
+        EncryptedFieldInfo if the field is encrypted, None otherwise
+    """
+    if not is_encrypted_field(field_type):
+        return None
+
+    # Get algorithm from Annotated args
+    algorithm = EncryptionAlgorithm.ARGON2
+    origin = get_origin(field_type)
+
+    if origin is Annotated:
+        args = get_args(field_type)
+        for arg in args:
+            if isinstance(arg, _EncryptedMarker):
+                algorithm = arg.algorithm
+                break
+
+    return EncryptedFieldInfo(algorithm=algorithm)