sqliter-py 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

sqliter/query/query.py

@@ -1,9 +1,25 @@
-"""Define the 'QueryBuilder' class for building SQL queries."""
+"""Implements the query building and execution logic for SQLiter.
+
+This module defines the QueryBuilder class, which provides a fluent
+interface for constructing SQL queries. It supports operations such
+as filtering, ordering, limiting, and various data retrieval methods,
+allowing for flexible and expressive database queries without writing
+raw SQL.
+"""
 
 from __future__ import annotations
 
 import sqlite3
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+import warnings
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    Optional,
+    Union,
+    overload,
+)
 
 from typing_extensions import LiteralString, Self
 
@@ -28,10 +44,36 @@ FilterValue = Union[
 
 
 class QueryBuilder:
-    """Functions to build and execute queries for a given model."""
+    """Builds and executes database queries for a specific model.
+
+    This class provides methods to construct SQL queries, apply filters,
+    set ordering, and execute the queries against the database.
+
+    Attributes:
+        db (SqliterDB): The database connection object.
+        model_class (type[BaseDBModel]): The Pydantic model class.
+        table_name (str): The name of the database table.
+        filters (list): List of applied filter conditions.
+        _limit (Optional[int]): The LIMIT clause value, if any.
+        _offset (Optional[int]): The OFFSET clause value, if any.
+        _order_by (Optional[str]): The ORDER BY clause, if any.
+        _fields (Optional[list[str]]): List of fields to select, if specified.
+    """
+
+    def __init__(
+        self,
+        db: SqliterDB,
+        model_class: type[BaseDBModel],
+        fields: Optional[list[str]] = None,
+    ) -> None:
+        """Initialize a new QueryBuilder instance.
 
-    def __init__(self, db: SqliterDB, model_class: type[BaseDBModel]) -> None:
-        """Initialize the query builder with the database, model class, etc."""
+        Args:
+            db: The database connection object.
+            model_class: The Pydantic model class for the table.
+            fields: Optional list of field names to select. If None, all fields
+                are selected.
+        """
         self.db = db
         self.model_class = model_class
         self.table_name = model_class.get_table_name()  # Use model_class method
@@ -39,9 +81,48 @@ class QueryBuilder:
         self._limit: Optional[int] = None
         self._offset: Optional[int] = None
         self._order_by: Optional[str] = None
+        self._fields: Optional[list[str]] = fields
+
+        if self._fields:
+            self._validate_fields()
+
+    def _validate_fields(self) -> None:
+        """Validate that the specified fields exist in the model.
+
+        Raises:
+            ValueError: If any specified field is not in the model.
+        """
+        if self._fields is None:
+            return
+        valid_fields = set(self.model_class.model_fields.keys())
+        invalid_fields = set(self._fields) - valid_fields
+        if invalid_fields:
+            err_message = (
+                f"Invalid fields specified: {', '.join(invalid_fields)}"
+            )
+            raise ValueError(err_message)
 
     def filter(self, **conditions: str | float | None) -> QueryBuilder:
-        """Add filter conditions to the query."""
+        """Apply filter conditions to the query.
+
+        This method allows adding one or more filter conditions to the query.
+        Each condition is specified as a keyword argument, where the key is
+        the field name and the value is the condition to apply.
+
+        Args:
+            **conditions: Arbitrary keyword arguments representing filter
+                conditions. The key is the field name, and the value is the
+                condition to apply. Supported operators include equality,
+                comparison, and special operators like __in, __isnull, etc.
+
+        Returns:
+            QueryBuilder: The current QueryBuilder instance for method
+            chaining.
+
+        Examples:
+            >>> query.filter(name="John", age__gt=30)
+            >>> query.filter(status__in=["active", "pending"])
+        """
         valid_fields = self.model_class.model_fields
 
         for field, value in conditions.items():
@@ -53,9 +134,94 @@ class QueryBuilder:
 
         return self
 
+    def fields(self, fields: Optional[list[str]] = None) -> QueryBuilder:
+        """Specify which fields to select in the query.
+
+        Args:
+            fields: List of field names to select. If None, all fields are
+                selected.
+
+        Returns:
+            The QueryBuilder instance for method chaining.
+        """
+        if fields:
+            self._fields = fields
+            self._validate_fields()
+        return self
+
+    def exclude(self, fields: Optional[list[str]] = None) -> QueryBuilder:
+        """Specify which fields to exclude from the query results.
+
+        Args:
+            fields: List of field names to exclude. If None, no fields are
+                excluded.
+
+        Returns:
+            The QueryBuilder instance for method chaining.
+
+        Raises:
+            ValueError: If exclusion results in no fields being selected or if
+                invalid fields are specified.
+        """
+        if fields:
+            all_fields = set(self.model_class.model_fields.keys())
+
+            # Check for invalid fields before subtraction
+            invalid_fields = set(fields) - all_fields
+            if invalid_fields:
+                err = (
+                    "Invalid fields specified for exclusion: "
+                    f"{', '.join(invalid_fields)}"
+                )
+                raise ValueError(err)
+
+            # Subtract the fields specified for exclusion
+            self._fields = list(all_fields - set(fields))
+
+            # Explicit check: raise an error if no fields remain
+            if not self._fields:
+                err = "Exclusion results in no fields being selected."
+                raise ValueError(err)
+
+            # Now validate the remaining fields to ensure they are all valid
+            self._validate_fields()
+
+        return self
+
+    def only(self, field: str) -> QueryBuilder:
+        """Specify a single field to select in the query.
+
+        Args:
+            field: The name of the field to select.
+
+        Returns:
+            The QueryBuilder instance for method chaining.
+
+        Raises:
+            ValueError: If the specified field is invalid.
+        """
+        all_fields = set(self.model_class.model_fields.keys())
+
+        # Validate that the field exists
+        if field not in all_fields:
+            err = f"Invalid field specified: {field}"
+            raise ValueError(err)
+
+        # Set self._fields to just the single field
+        self._fields = [field]
+        return self
+
     def _get_operator_handler(
         self, operator: str
     ) -> Callable[[str, Any, str], None]:
+        """Get the appropriate handler function for the given operator.
+
+        Args:
+            operator: The filter operator string.
+
+        Returns:
+            A callable that handles the specific operator type.
+        """
         handlers = {
             "__isnull": self._handle_null,
             "__notnull": self._handle_null,
@@ -78,12 +244,31 @@ class QueryBuilder:
     def _validate_field(
         self, field_name: str, valid_fields: dict[str, FieldInfo]
     ) -> None:
+        """Validate that a field exists in the model.
+
+        Args:
+            field_name: The name of the field to validate.
+            valid_fields: Dictionary of valid fields from the model.
+
+        Raises:
+            InvalidFilterError: If the field is not in the model.
+        """
         if field_name not in valid_fields:
             raise InvalidFilterError(field_name)
 
     def _handle_equality(
         self, field_name: str, value: FilterValue, operator: str
     ) -> None:
+        """Handle equality filter conditions.
+
+        Args:
+            field_name: The name of the field to filter on.
+            value: The value to compare against.
+            operator: The operator string (usually '__eq').
+
+        This method adds an equality condition to the filters list, handling
+        NULL values separately.
+        """
         if value is None:
             self.filters.append((f"{field_name} IS NULL", None, "__isnull"))
         else:
@@ -92,6 +277,16 @@ class QueryBuilder:
     def _handle_null(
         self, field_name: str, _: FilterValue, operator: str
     ) -> None:
+        """Handle IS NULL and IS NOT NULL filter conditions.
+
+        Args:
+            field_name: The name of the field to filter on. _: Placeholder for
+                unused value parameter.
+            operator: The operator string ('__isnull' or '__notnull').
+
+        This method adds an IS NULL or IS NOT NULL condition to the filters
+        list.
+        """
         condition = (
             f"{field_name} IS NOT NULL"
             if operator == "__notnull"
@@ -102,6 +297,18 @@ class QueryBuilder:
     def _handle_in(
         self, field_name: str, value: FilterValue, operator: str
     ) -> None:
+        """Handle IN and NOT IN filter conditions.
+
+        Args:
+            field_name: The name of the field to filter on.
+            value: A list of values to check against.
+            operator: The operator string ('__in' or '__not_in').
+
+        Raises:
+            TypeError: If the value is not a list.
+
+        This method adds an IN or NOT IN condition to the filters list.
+        """
         if not isinstance(value, list):
             err = f"{field_name} requires a list for '{operator}'"
             raise TypeError(err)
@@ -118,6 +325,19 @@ class QueryBuilder:
     def _handle_like(
         self, field_name: str, value: FilterValue, operator: str
     ) -> None:
+        """Handle LIKE and GLOB filter conditions.
+
+        Args:
+            field_name: The name of the field to filter on.
+            value: The pattern to match against.
+            operator: The operator string (e.g., '__startswith', '__contains').
+
+        Raises:
+            TypeError: If the value is not a string.
+
+        This method adds a LIKE or GLOB condition to the filters list, depending
+        on whether the operation is case-sensitive or not.
+        """
         if not isinstance(value, str):
             err = f"{field_name} requires a string value for '{operator}'"
             raise TypeError(err)
@@ -142,11 +362,29 @@ class QueryBuilder:
     def _handle_comparison(
         self, field_name: str, value: FilterValue, operator: str
     ) -> None:
+        """Handle comparison filter conditions.
+
+        Args:
+            field_name: The name of the field to filter on.
+            value: The value to compare against.
+            operator: The comparison operator string (e.g., '__lt', '__gte').
+
+        This method adds a comparison condition to the filters list.
+        """
         sql_operator = OPERATOR_MAPPING[operator]
         self.filters.append((f"{field_name} {sql_operator} ?", value, operator))
 
     # Helper method for parsing field and operator
     def _parse_field_operator(self, field: str) -> tuple[str, str]:
+        """Parse a field string to separate the field name and operator.
+
+        Args:
+            field: The field string, potentially including an operator.
+
+        Returns:
+            A tuple containing the field name and the operator (or '__eq' if
+            no operator was specified).
+        """
         for operator in OPERATOR_MAPPING:
             if field.endswith(operator):
                 return field[: -len(operator)], operator
@@ -154,7 +392,15 @@ class QueryBuilder:
 
     # Helper method for formatting string operators (like startswith)
     def _format_string_for_operator(self, operator: str, value: str) -> str:
-        # Mapping operators to their corresponding string format
+        """Format a string value based on the specified operator.
+
+        Args:
+            operator: The operator string (e.g., '__startswith', '__contains').
+            value: The original string value.
+
+        Returns:
+            The formatted string value suitable for the given operator.
+        """
         format_map = {
             "__startswith": f"{value}*",
             "__endswith": f"*{value}",
@@ -168,12 +414,29 @@ class QueryBuilder:
         return format_map.get(operator, value)
 
     def limit(self, limit_value: int) -> Self:
-        """Limit the number of results returned by the query."""
+        """Limit the number of results returned by the query.
+
+        Args:
+        limit_value: The maximum number of records to return.
+
+        Returns:
+            The QueryBuilder instance for method chaining.
+        """
         self._limit = limit_value
         return self
 
     def offset(self, offset_value: int) -> Self:
-        """Set an offset value for the query."""
+        """Set an offset value for the query.
+
+        Args:
+            offset_value: The number of records to skip.
+
+        Returns:
+            The QueryBuilder instance for method chaining.
+
+        Raises:
+            InvalidOffsetError: If the offset value is negative.
+        """
         if offset_value < 0:
             raise InvalidOffsetError(offset_value)
         self._offset = offset_value
@@ -182,34 +445,64 @@ class QueryBuilder:
             self._limit = -1
         return self
 
-    def order(self, order_by_field: str, direction: str = "ASC") -> Self:
-        """Order the results by a specific field and optionally direction.
-
-        Currently only supports ordering by a single field, though this will be
-        expanded in the future. You can chain this method to order by multiple
-        fields.
+    def order(
+        self,
+        order_by_field: Optional[str] = None,
+        direction: Optional[str] = None,
+        *,
+        reverse: bool = False,
+    ) -> Self:
+        """Order the query results by the specified field.
 
-        Parameters:
-            order_by_field (str): The field to order by.
-            direction (str, optional): The sorting direction, either 'ASC' or
-                'DESC'. Defaults to 'ASC'.
+        Args:
+            order_by_field: The field to order by [optional].
+            direction: Deprecated. Use 'reverse' instead.
+            reverse: If True, sort in descending order.
 
         Returns:
-            Self: Returns the query object for chaining.
+            The QueryBuilder instance for method chaining.
 
         Raises:
-            InvalidOrderError: If the field or direction is invalid.
+            InvalidOrderError: If the field doesn't exist or if both 'direction'
+                and 'reverse' are specified.
+
+        Warns:
+            DeprecationWarning: If 'direction' is used instead of 'reverse'.
         """
+        if direction:
+            warnings.warn(
+                "'direction' argument is deprecated and will be removed in a "
+                "future version. Use 'reverse' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if order_by_field is None:
+            order_by_field = self.model_class.get_primary_key()
+
         if order_by_field not in self.model_class.model_fields:
             err = f"'{order_by_field}' does not exist in the model fields."
             raise InvalidOrderError(err)
-
-        valid_directions = {"ASC", "DESC"}
-        if direction.upper() not in valid_directions:
-            err = f"'{direction}' is not a valid sorting direction."
+        # Raise an exception if both 'direction' and 'reverse' are specified
+        if direction and reverse:
+            err = (
+                "Cannot specify both 'direction' and 'reverse' as it "
+                "is ambiguous."
+            )
             raise InvalidOrderError(err)
 
-        self._order_by = f'"{order_by_field}" {direction.upper()}'
+        # Determine the sorting direction
+        if reverse:
+            sort_order = "DESC"
+        elif direction:
+            sort_order = direction.upper()
+            if sort_order not in {"ASC", "DESC"}:
+                err = f"'{direction}' is not a valid sorting direction."
+                raise InvalidOrderError(err)
+        else:
+            sort_order = "ASC"
+
+        self._order_by = f'"{order_by_field}" {sort_order}'
         return self
 
     def _execute_query(
@@ -218,15 +511,32 @@ class QueryBuilder:
         fetch_one: bool = False,
         count_only: bool = False,
     ) -> list[tuple[Any, ...]] | Optional[tuple[Any, ...]]:
-        """Helper function to execute the query with filters."""
-        fields = ", ".join(self.model_class.model_fields)
+        """Execute the constructed SQL query.
 
-        # Build the WHERE clause with special handling for None (NULL in SQL)
-        values, where_clause = self._parse_filter()
+        Args:
+            fetch_one: If True, fetch only one result.
+            count_only: If True, return only the count of results.
+
+        Returns:
+            A list of tuples (all results), a single tuple (one result),
+            or None if no results are found.
 
-        select_fields = fields if not count_only else "COUNT(*)"
+        Raises:
+            RecordFetchError: If there's an error executing the query.
+        """
+        if count_only:
+            fields = "COUNT(*)"
+        elif self._fields:
+            fields = ", ".join(f'"{field}"' for field in self._fields)
+        else:
+            fields = ", ".join(
+                f'"{field}"' for field in self.model_class.model_fields
+            )
 
-        sql = f'SELECT {select_fields} FROM "{self.table_name}"'  # noqa: S608 # nosec
+        sql = f'SELECT {fields} FROM "{self.table_name}"'  # noqa: S608 # nosec
+
+        # Build the WHERE clause with special handling for None (NULL in SQL)
+        values, where_clause = self._parse_filter()
 
         if self.filters:
             sql += f" WHERE {where_clause}"
@@ -242,6 +552,11 @@ class QueryBuilder:
             sql += " OFFSET ?"
             values.append(self._offset)
 
+        # Print the raw SQL and values if debug is enabled
+        # Log the SQL if debug is enabled
+        if self.db.debug:
+            self.db._log_sql(sql, values)  # noqa: SLF001
+
         try:
             with self.db.connect() as conn:
                 cursor = conn.cursor()
@@ -251,7 +566,13 @@ class QueryBuilder:
             raise RecordFetchError(self.table_name) from exc
 
     def _parse_filter(self) -> tuple[list[Any], LiteralString]:
-        """Actually parse the filters."""
+        """Parse the filter conditions into SQL clauses and values.
+
+        Returns:
+            A tuple containing:
+            - A list of values to be used in the SQL query.
+            - A string representing the WHERE clause of the SQL query.
+        """
         where_clauses = []
         values = []
         for field, value, operator in self.filters:
@@ -269,68 +590,114 @@ class QueryBuilder:
         where_clause = " AND ".join(where_clauses)
         return values, where_clause
 
-    def fetch_all(self) -> list[BaseDBModel]:
-        """Fetch all results matching the filters."""
-        results = self._execute_query()
+    def _convert_row_to_model(self, row: tuple[Any, ...]) -> BaseDBModel:
+        """Convert a database row to a model instance.
 
-        if not results:
-            return []
+        Args:
+            row: A tuple representing a database row.
 
-        return [
-            self.model_class(
-                **{
-                    field: row[idx]
-                    for idx, field in enumerate(self.model_class.model_fields)
-                }
+        Returns:
+            An instance of the model class populated with the row data.
+        """
+        if self._fields:
+            return self.model_class.model_validate_partial(
+                {field: row[idx] for idx, field in enumerate(self._fields)}
             )
-            for row in results
-        ]
-
-    def fetch_one(self) -> BaseDBModel | None:
-        """Fetch exactly one result."""
-        result = self._execute_query(fetch_one=True)
-        if not result:
-            return None
         return self.model_class(
             **{
-                field: result[idx]
+                field: row[idx]
                 for idx, field in enumerate(self.model_class.model_fields)
             }
         )
 
-    def fetch_first(self) -> BaseDBModel | None:
-        """Fetch the first result of the query."""
-        self._limit = 1
-        result = self._execute_query()
+    @overload
+    def _fetch_result(
+        self, *, fetch_one: Literal[True]
+    ) -> Optional[BaseDBModel]: ...
+
+    @overload
+    def _fetch_result(
+        self, *, fetch_one: Literal[False]
+    ) -> list[BaseDBModel]: ...
+
+    def _fetch_result(
+        self, *, fetch_one: bool = False
+    ) -> Union[list[BaseDBModel], Optional[BaseDBModel]]:
+        """Fetch and convert query results to model instances.
+
+        Args:
+            fetch_one: If True, fetch only one result.
+
+        Returns:
+            A list of model instances, a single model instance, or None if no
+            results are found.
+        """
+        result = self._execute_query(fetch_one=fetch_one)
+
         if not result:
-            return None
-        return self.model_class(
-            **{
-                field: result[0][idx]
-                for idx, field in enumerate(self.model_class.model_fields)
-            }
-        )
+            if fetch_one:
+                return None
+            return []
+
+        if fetch_one:
+            # Ensure we pass a tuple, not a list, to _convert_row_to_model
+            if isinstance(result, list):
+                result = result[
+                    0
+                ]  # Get the first (and only) result if it's wrapped in a list.
+            return self._convert_row_to_model(result)
+
+        return [self._convert_row_to_model(row) for row in result]
+
+    def fetch_all(self) -> list[BaseDBModel]:
+        """Fetch all results of the query.
+
+        Returns:
+            A list of model instances representing all query results.
+        """
+        return self._fetch_result(fetch_one=False)
+
+    def fetch_one(self) -> Optional[BaseDBModel]:
+        """Fetch a single result of the query.
+
+        Returns:
+            A single model instance or None if no result is found.
+        """
+        return self._fetch_result(fetch_one=True)
+
+    def fetch_first(self) -> Optional[BaseDBModel]:
+        """Fetch the first result of the query.
 
-    def fetch_last(self) -> BaseDBModel | None:
-        """Fetch the last result of the query (based on the insertion order)."""
+        Returns:
+            The first model instance or None if no result is found.
+        """
+        self._limit = 1
+        return self._fetch_result(fetch_one=True)
+
+    def fetch_last(self) -> Optional[BaseDBModel]:
+        """Fetch the last result of the query.
+
+        Returns:
+            The last model instance or None if no result is found.
+        """
         self._limit = 1
         self._order_by = "rowid DESC"
-        result = self._execute_query()
-        if not result:
-            return None
-        return self.model_class(
-            **{
-                field: result[0][idx]
-                for idx, field in enumerate(self.model_class.model_fields)
-            }
-        )
+        return self._fetch_result(fetch_one=True)
 
     def count(self) -> int:
-        """Return the count of records matching the filters."""
+        """Count the number of results for the current query.
+
+        Returns:
+            The number of results that match the current query conditions.
+        """
         result = self._execute_query(count_only=True)
 
         return int(result[0][0]) if result else 0
 
     def exists(self) -> bool:
-        """Return True if any record matches the filters."""
+        """Check if any results exist for the current query.
+
+        Returns:
+            True if at least one result exists, False otherwise.
+        """
         return self.count() > 0