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

surreal_orm/query_set.py

@@ -1,14 +1,20 @@
 from .constants import LOOKUP_OPERATORS
 from .enum import OrderBy
 from .utils import remove_quotes_for_variables
-from surrealdb import QueryResponse, Table, AsyncSurrealDB
-from surrealdb.errors import SurrealDbError
 from . import BaseSurrealModel, SurrealDBConnectionManager
-from typing import Self, Any, cast
+from .aggregations import Aggregation
+from typing import Self, Any, Sequence, cast
 from pydantic_core import ValidationError
 
 import logging
 
+
+class SurrealDbError(Exception):
+    """Error from SurrealDB operations."""
+
+    pass
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -55,8 +61,14 @@ class QuerySet:
         self._limit: int | None = None
         self._offset: int | None = None
         self._order_by: str | None = None
-        self._model_table: str = getattr(model, "_table_name", model.__name__)
+        self._model_table: str = model.__name__
         self._variables: dict = {}
+        self._group_by_fields: list[str] = []
+        self._annotations: dict[str, Aggregation] = {}
+        # Relation query options
+        self._select_related: list[str] = []
+        self._prefetch_related: list[str] = []
+        self._traversal_path: str | None = None
 
     def select(self, *fields: str) -> Self:
         """
@@ -201,7 +213,7 @@ class QuerySet:
         self._offset = value
         return self
 
-    def order_by(self, field_name: str, type: OrderBy = OrderBy.ASC) -> Self:
+    def order_by(self, field_name: str, order_type: OrderBy = OrderBy.ASC) -> Self:
         """
         Set the field and direction to order the results by.
 
@@ -220,9 +232,248 @@ class QuerySet:
             queryset.order_by('name', OrderBy.DESC)
             ```
         """
-        self._order_by = f"{field_name} {type}"
+        self._order_by = f"{field_name} {order_type}"
+        return self
+
+    def values(self, *fields: str) -> Self:
+        """
+        Specify the fields to group by for aggregation queries.
+
+        This method is used in conjunction with `annotate()` to perform GROUP BY operations.
+        The specified fields become the grouping keys, and aggregation functions are applied
+        to each group.
+
+        Args:
+            *fields (str): Variable length argument list of field names to group by.
+
+        Returns:
+            Self: The current instance of QuerySet to allow method chaining.
+
+        Example:
+            ```python
+            # Group orders by status and calculate statistics
+            stats = await Order.objects().values("status").annotate(
+                count=Count(),
+                total=Sum("amount"),
+            )
+            # Result: [{"status": "paid", "count": 42, "total": 5000}, ...]
+            ```
+        """
+        self._group_by_fields = list(fields)
+        return self
+
+    def annotate(self, **aggregations: Aggregation) -> Self:
+        """
+        Add aggregation functions to compute values for each group.
+
+        This method is used in conjunction with `values()` to perform GROUP BY operations.
+        Each keyword argument should be an Aggregation instance (Count, Sum, Avg, Min, Max).
+
+        Args:
+            **aggregations (Aggregation): Keyword arguments where keys are alias names
+                and values are Aggregation instances.
+
+        Returns:
+            Self: The current instance of QuerySet to allow method chaining.
+
+        Example:
+            ```python
+            from surreal_orm.aggregations import Count, Sum, Avg
+
+            # Calculate statistics per status
+            stats = await Order.objects().values("status").annotate(
+                count=Count(),
+                total=Sum("amount"),
+                avg_amount=Avg("amount"),
+            )
+            ```
+        """
+        self._annotations = aggregations
+        return self
+
+    # ==================== Relation Query Methods ====================
+
+    def select_related(self, *relations: str) -> Self:
+        """
+        Eagerly load related objects in the same query.
+
+        This method optimizes queries by loading related objects alongside
+        the main query results, avoiding N+1 query problems for forward relations.
+
+        Note: SurrealDB handles this through graph traversal syntax.
+        The actual loading happens during query execution.
+
+        Args:
+            *relations: Names of relations to load eagerly.
+
+        Returns:
+            Self: The current instance of QuerySet to allow method chaining.
+
+        Example:
+            ```python
+            # Load posts with their authors in one query
+            posts = await Post.objects().select_related("author").all()
+            for post in posts:
+                print(post.author.name)  # No additional query
+            ```
+        """
+        self._select_related = list(relations)
+        return self
+
+    def prefetch_related(self, *relations: str) -> Self:
+        """
+        Prefetch related objects using separate optimized queries.
+
+        This method reduces N+1 query problems by fetching related objects
+        in batches after the main query completes. This is more efficient
+        than select_related for many-to-many and reverse relations.
+
+        Args:
+            *relations: Names of relations to prefetch.
+
+        Returns:
+            Self: The current instance of QuerySet to allow method chaining.
+
+        Example:
+            ```python
+            # Load users and prefetch their followers
+            users = await User.objects().prefetch_related("followers", "posts").all()
+            for user in users:
+                print(user.followers)  # Already loaded
+                print(user.posts)      # Already loaded
+            ```
+        """
+        self._prefetch_related = list(relations)
+        return self
+
+    def traverse(self, path: str) -> Self:
+        """
+        Add a graph traversal path to the query.
+
+        This method allows querying across graph relations using
+        SurrealDB's traversal syntax.
+
+        Args:
+            path: Graph traversal path (e.g., "->follows->users->likes->posts")
+
+        Returns:
+            Self: The current instance of QuerySet to allow method chaining.
+
+        Example:
+            ```python
+            # Get all posts liked by users that alice follows
+            posts = await User.objects().filter(id="alice").traverse(
+                "->follows->users->likes->posts"
+            ).all()
+            ```
+        """
+        self._traversal_path = path
         return self
 
+    async def graph_query(self, traversal: str, **variables: Any) -> list[dict[str, Any]]:
+        """
+        Execute a raw graph traversal query.
+
+        This method provides direct access to SurrealDB's graph capabilities
+        for complex traversal patterns that can't be expressed through
+        the standard QuerySet API.
+
+        Args:
+            traversal: Graph traversal expression (e.g., "->follows->User")
+            **variables: Variables to bind in the query
+
+        Returns:
+            list[dict[str, Any]]: Raw query results as dictionaries
+
+        Example:
+            ```python
+            # Find users that alice follows
+            result = await User.objects().filter(id="alice").graph_query("->follows->User")
+
+            # Multi-hop traversal
+            result = await User.objects().filter(id="alice").graph_query(
+                "->follows->User->follows->User"
+            )
+            ```
+        """
+        # Parse the traversal to determine edge and direction
+        # Expected format: "->edge->Table" or "<-edge<-Table"
+        # For now, support simple single-hop traversals
+
+        # Check if we have an id filter to use as starting point
+        source_id = None
+        for field_name, lookup_name, value in self._filters:
+            if field_name == "id" and lookup_name == "exact":
+                source_id = value
+                break
+
+        client = await SurrealDBConnectionManager.get_client()
+
+        if source_id:
+            # Use specific record as starting point
+            source_thing = f"{self._model_table}:{source_id}"
+
+            # Parse the traversal to get edge and direction
+            # Simple pattern: ->edge->Table or <-edge<-Table
+            if traversal.startswith("->"):
+                # Outgoing: get targets where source is 'in'
+                parts = traversal.split("->")
+                if len(parts) >= 3:
+                    edge = parts[1]
+                    query = f"SELECT out FROM {edge} WHERE in = {source_thing} FETCH out;"
+                    result = await client.query(query, {**self._variables, **variables})
+                    records: list[dict[str, Any]] = []
+                    for row in result.all_records or []:
+                        if isinstance(row.get("out"), dict):
+                            records.append(row["out"])
+                    return records
+            elif traversal.startswith("<-"):
+                # Incoming: get sources where target is 'out'
+                parts = traversal.split("<-")
+                if len(parts) >= 3:
+                    edge = parts[1]
+                    query = f"SELECT in FROM {edge} WHERE out = {source_thing} FETCH in;"
+                    result = await client.query(query, {**self._variables, **variables})
+                    records = []
+                    for row in result.all_records or []:
+                        if isinstance(row.get("in"), dict):
+                            records.append(row["in"])
+                    return records
+
+        # Fallback: return empty for unsupported patterns
+        return []
+
+    async def _execute_annotate(self) -> list[dict[str, Any]]:
+        """
+        Execute the GROUP BY query with annotations.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries containing grouped results.
+        """
+        # Build SELECT clause with group fields and aggregations
+        select_parts: list[str] = list(self._group_by_fields)
+
+        for alias, aggregation in self._annotations.items():
+            select_parts.append(aggregation.to_surql(alias))
+
+        select_clause = ", ".join(select_parts)
+
+        # Build WHERE clause
+        where_clause = self._compile_where_clause()
+
+        # Build GROUP BY clause
+        if self._group_by_fields:
+            group_clause = f" GROUP BY {', '.join(self._group_by_fields)}"
+        else:
+            group_clause = " GROUP ALL"
+
+        query = f"SELECT {select_clause} FROM {self._model_table}{where_clause}{group_clause};"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        return cast(list[dict[str, Any]], result.all_records)
+
     def _compile_query(self) -> str:
         """
         Compile the QuerySet parameters into a SQL query string.
@@ -261,6 +512,10 @@ class QuerySet:
         if where_clauses:
             query += " WHERE " + " AND ".join(where_clauses)
 
+        # Append ORDER BY if set (must come before LIMIT/START in SurrealQL)
+        if self._order_by:
+            query += f" ORDER BY {self._order_by}"
+
         # Append LIMIT if set
         if self._limit is not None:
             query += f" LIMIT {self._limit}"
@@ -269,10 +524,6 @@ class QuerySet:
         if self._offset is not None:
             query += f" START {self._offset}"
 
-        # Append ORDER BY if set
-        if self._order_by:
-            query += f" ORDER BY {self._order_by}"
-
         query += ";"
         return query
 
@@ -284,28 +535,41 @@ class QuerySet:
         the results. If the data conforms to the model schema, it returns a list of model instances;
         otherwise, it returns a list of dictionaries.
 
+        When `annotate()` has been called, this returns the aggregated results as dictionaries
+        instead of model instances.
+
         Returns:
             list[BaseSurrealModel] | list[dict]: A list of model instances if validation is successful,
-            otherwise a list of dictionaries representing the raw data.
+            otherwise a list of dictionaries representing the raw data. For annotated queries,
+            always returns a list of dictionaries.
 
         Raises:
             SurrealDbError: If there is an issue executing the query.
 
         Example:
             ```python
+            # Regular query
             results = await queryset.exec()
+
+            # Aggregation query
+            stats = await Order.objects().values("status").annotate(
+                count=Count(),
+                total=Sum("amount"),
+            ).exec()
             ```
         """
-        data: dict[str, Any] = {"result": []}
+        # If annotations are set, execute as GROUP BY query
+        if self._annotations:
+            return await self._execute_annotate()
+
         query = self._compile_query()
         results = await self._execute_query(query)
         try:
-            data = cast(dict, results[0])
-
-            return self.model.from_db(data["result"])
+            # surrealdb SDK 1.0.8 returns records directly, not wrapped in {"result": ...}
+            return self.model.from_db(cast(dict | list | None, results))
         except ValidationError as e:
             logger.info(f"Pydantic invalid format for the class, returning dict value: {e}")
-            return data["result"]
+            return results
 
     async def first(self) -> Any:
         """
@@ -331,7 +595,7 @@ class QuerySet:
         if results:
             return results[0]
 
-        raise SurrealDbError("No result found.")
+        raise self.model.DoesNotExist("Query returned no results.")
 
     async def get(self, id_item: Any = None) -> Any:
         """
@@ -357,15 +621,18 @@ class QuerySet:
         """
         if id_item:
             client = await SurrealDBConnectionManager.get_client()
-            data = await client.select(f"{self._model_table}:{id_item}")
-            return self.model.from_db(data)
+            result = await client.select(f"{self._model_table}:{id_item}")
+            # SDK returns RecordsResponse
+            if result.is_empty:
+                raise self.model.DoesNotExist("Record not found.")
+            return self.model.from_db(cast(dict | list | None, result.first))
         else:
             result = await self.exec()
             if len(result) > 1:
                 raise SurrealDbError("More than one result found.")
 
             if len(result) == 0:
-                raise SurrealDbError("No result found.")
+                raise self.model.DoesNotExist("Record not found.")
             return result[0]
 
     async def all(self) -> Any:
@@ -386,10 +653,168 @@ class QuerySet:
             ```
         """
         client = await SurrealDBConnectionManager.get_client()
-        results = await client.select(Table(self._model_table))
-        return self.model.from_db(results)
+        result = await client.select(self._model_table)
+        return self.model.from_db(cast(dict | list | None, result.records))
+
+    # ==================== Aggregation Methods ====================
+
+    def _compile_where_clause(self) -> str:
+        """
+        Compile the WHERE clause from filters.
+
+        Returns:
+            str: The WHERE clause string (including WHERE keyword) or empty string.
+        """
+        if not self._filters:
+            return ""
+
+        where_clauses = []
+        for field_name, lookup_name, value in self._filters:
+            op = LOOKUP_OPERATORS.get(lookup_name, "=")
+            if lookup_name == "in":
+                formatted_values = ", ".join(repr(v) for v in value)
+                where_clauses.append(f"{field_name} {op} [{formatted_values}]")
+            else:
+                where_clauses.append(f"{field_name} {op} {repr(value)}")
+
+        return " WHERE " + " AND ".join(where_clauses)
+
+    async def count(self) -> int:
+        """
+        Count the number of records matching the current filters.
+
+        Returns:
+            int: The number of matching records.
+
+        Example:
+            ```python
+            total = await User.objects().count()
+            active = await User.objects().filter(active=True).count()
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        query = f"SELECT count() FROM {self._model_table}{where_clause} GROUP ALL;"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        if result.all_records:
+            record = result.all_records[0]
+            if isinstance(record, dict) and "count" in record:
+                return int(record["count"])
+        return 0
+
+    async def sum(self, field: str) -> float | int:
+        """
+        Calculate the sum of a numeric field.
+
+        Args:
+            field: The field name to sum.
+
+        Returns:
+            float | int: The sum of the field values, or 0 if no records match.
+
+        Example:
+            ```python
+            total = await Order.objects().filter(status="paid").sum("amount")
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        query = f"SELECT math::sum({field}) AS total FROM {self._model_table}{where_clause} GROUP ALL;"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        if result.all_records:
+            record = result.all_records[0]
+            if isinstance(record, dict) and "total" in record:
+                value = record["total"]
+                return value if value is not None else 0
+        return 0
+
+    async def avg(self, field: str) -> float | None:
+        """
+        Calculate the average of a numeric field.
+
+        Args:
+            field: The field name to average.
+
+        Returns:
+            float | None: The average value, or None if no records match.
+
+        Example:
+            ```python
+            avg_age = await User.objects().filter(active=True).avg("age")
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        query = f"SELECT math::mean({field}) AS average FROM {self._model_table}{where_clause} GROUP ALL;"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        if result.all_records:
+            record = result.all_records[0]
+            if isinstance(record, dict) and "average" in record:
+                value = record["average"]
+                return float(value) if value is not None else None
+        return None
+
+    async def min(self, field: str) -> Any:
+        """
+        Get the minimum value of a field.
+
+        Args:
+            field: The field name to find the minimum of.
+
+        Returns:
+            Any: The minimum value, or None if no records match.
+
+        Example:
+            ```python
+            min_price = await Product.objects().min("price")
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        query = f"SELECT math::min({field}) AS minimum FROM {self._model_table}{where_clause} GROUP ALL;"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        if result.all_records:
+            record = result.all_records[0]
+            if isinstance(record, dict) and "minimum" in record:
+                return record["minimum"]
+        return None
+
+    async def max(self, field: str) -> Any:
+        """
+        Get the maximum value of a field.
+
+        Args:
+            field: The field name to find the maximum of.
 
-    async def _execute_query(self, query: str) -> list[QueryResponse]:
+        Returns:
+            Any: The maximum value, or None if no records match.
+
+        Example:
+            ```python
+            max_price = await Product.objects().max("price")
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        query = f"SELECT math::max({field}) AS maximum FROM {self._model_table}{where_clause} GROUP ALL;"
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+
+        if result.all_records:
+            record = result.all_records[0]
+            if isinstance(record, dict) and "maximum" in record:
+                return record["maximum"]
+        return None
+
+    async def _execute_query(self, query: str) -> list[Any]:
         """
         Execute the given SQL query using the SurrealDB client.
 
@@ -400,7 +825,7 @@ class QuerySet:
             query (str): The SQL query string to execute.
 
         Returns:
-            list[QueryResponse]: A list of `QueryResponse` objects containing the query results.
+            list[Any]: A list of query response objects containing the query results.
 
         Raises:
             SurrealDbError: If there is an issue executing the query.
@@ -413,7 +838,7 @@ class QuerySet:
         client = await SurrealDBConnectionManager.get_client()
         return await self._run_query_on_client(client, query)
 
-    async def _run_query_on_client(self, client: AsyncSurrealDB, query: str) -> list[QueryResponse]:
+    async def _run_query_on_client(self, client: Any, query: str) -> list[Any]:
         """
         Run the SQL query on the provided SurrealDB client.
 
@@ -421,11 +846,11 @@ class QuerySet:
         and returns the raw query responses.
 
         Args:
-            client (AsyncSurrealDB): The active SurrealDB client instance.
+            client: The active SurrealDB client instance.
             query (str): The SQL query string to execute.
 
         Returns:
-            list[QueryResponse]: A list of `QueryResponse` objects containing the query results.
+            list[Any]: A list of query response objects containing the query results.
 
         Raises:
             SurrealDbError: If there is an issue executing the query.
@@ -435,7 +860,9 @@ class QuerySet:
             results = await self._run_query_on_client(client, "SELECT * FROM users;")
             ```
         """
-        return await client.query(remove_quotes_for_variables(query), self._variables)  # type: ignore
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+        # SDK returns QueryResponse, extract all records
+        return cast(list[Any], result.all_records)
 
     async def delete_table(self) -> bool:
         """
@@ -456,7 +883,7 @@ class QuerySet:
             ```
         """
         client = await SurrealDBConnectionManager.get_client()
-        await client.delete(Table(self._model_table))
+        await client.delete(self._model_table)
         return True
 
     async def query(self, query: str, variables: dict[str, Any] = {}) -> Any:
@@ -484,9 +911,157 @@ class QuerySet:
             results = await queryset.query(custom_query, variables={'status': 'active'})
             ```
         """
-        if f"FROM {self.model.__name__}" not in query:
-            raise SurrealDbError(f"The query must include 'FROM {self.model.__name__}' to reference the correct table.")
+        if f"FROM {self._model_table}" not in query:
+            raise SurrealDbError(f"The query must include 'FROM {self._model_table}' to reference the correct table.")
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), variables)
+        # SDK returns QueryResponse, extract all records
+        return self.model.from_db(cast(dict | list | None, result.all_records))
+
+    # ==================== Bulk Operations ====================
+
+    async def bulk_create(
+        self,
+        instances: Sequence[BaseSurrealModel],
+        atomic: bool = False,
+        batch_size: int | None = None,
+    ) -> list[BaseSurrealModel]:
+        """
+        Create multiple model instances in the database efficiently.
+
+        Args:
+            instances: A sequence of model instances to create.
+            atomic: If True, all creates are wrapped in a transaction.
+                    If any fails, all are rolled back.
+            batch_size: If specified, instances are created in batches of this size.
+                        Useful for very large datasets to avoid memory issues.
+
+        Returns:
+            list[BaseSurrealModel]: The created instances.
+
+        Example:
+            ```python
+            users = [User(name=f"User{i}") for i in range(1000)]
+
+            # Simple bulk create
+            created = await User.objects().bulk_create(users)
+
+            # Atomic bulk create
+            created = await User.objects().bulk_create(users, atomic=True)
+
+            # With batch size
+            created = await User.objects().bulk_create(users, batch_size=100)
+            ```
+        """
+        if not instances:
+            return []
+
+        created: list[BaseSurrealModel] = []
+
+        if atomic:
+            # Use transaction for atomicity
+            async with await SurrealDBConnectionManager.transaction() as tx:
+                for instance in instances:
+                    await instance.save(tx=tx)
+                    created.append(instance)
+        elif batch_size:
+            # Process in batches
+            for i in range(0, len(instances), batch_size):
+                batch = instances[i : i + batch_size]
+                for instance in batch:
+                    await instance.save()
+                    created.append(instance)
+        else:
+            # Simple sequential create
+            for instance in instances:
+                await instance.save()
+                created.append(instance)
+
+        return created
+
+    async def bulk_update(
+        self,
+        data: dict[str, Any],
+        atomic: bool = False,
+    ) -> int:
+        """
+        Update all records matching the current filters.
+
+        Args:
+            data: A dictionary of field names and values to update.
+            atomic: If True, all updates are wrapped in a transaction.
+
+        Returns:
+            int: The number of records updated.
+
+        Example:
+            ```python
+            # Update all matching records
+            updated = await User.objects().filter(
+                last_login__lt="2025-01-01"
+            ).bulk_update({"status": "inactive"})
+
+            # Atomic update
+            updated = await User.objects().filter(role="guest").bulk_update(
+                {"verified": True},
+                atomic=True
+            )
+            ```
+        """
+        where_clause = self._compile_where_clause()
+
+        # Build SET clause
+        set_parts = []
+        for field, value in data.items():
+            set_parts.append(f"{field} = {repr(value)}")
+        set_clause = ", ".join(set_parts)
+
+        query = f"UPDATE {self._model_table} SET {set_clause}{where_clause};"
+
+        if atomic:
+            # For atomic operations, count first then update in transaction
+            current_count = await self.count()
+            async with await SurrealDBConnectionManager.transaction() as tx:
+                await tx.query(remove_quotes_for_variables(query), self._variables)
+            return current_count
+
+        client = await SurrealDBConnectionManager.get_client()
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+        return len(result.all_records)
+
+    async def bulk_delete(self, atomic: bool = False) -> int:
+        """
+        Delete all records matching the current filters.
+
+        Args:
+            atomic: If True, all deletes are wrapped in a transaction.
+
+        Returns:
+            int: The number of records deleted.
+
+        Example:
+            ```python
+            # Delete all matching records
+            deleted = await User.objects().filter(status="deleted").bulk_delete()
+
+            # Atomic delete
+            deleted = await Order.objects().filter(
+                created_at__lt="2024-01-01"
+            ).bulk_delete(atomic=True)
+            ```
+        """
+        where_clause = self._compile_where_clause()
+        # Use RETURN BEFORE to get deleted records count
+        query = f"DELETE FROM {self._model_table}{where_clause} RETURN BEFORE;"
+
+        if atomic:
+            # For atomic operations, count first then delete in transaction
+            current_count = await self.count()
+            delete_query = f"DELETE FROM {self._model_table}{where_clause};"
+            async with await SurrealDBConnectionManager.transaction() as tx:
+                await tx.query(remove_quotes_for_variables(delete_query), self._variables)
+            return current_count
+
         client = await SurrealDBConnectionManager.get_client()
-        results = await client.query(remove_quotes_for_variables(query), variables)
-        data = cast(dict, results[0])
-        return self.model.from_db(data["result"])
+        result = await client.query(remove_quotes_for_variables(query), self._variables)
+        return len(result.all_records)