add optional 'fields' arg to 'Select()'

Signed-off-by: Grant Ramsay <seapagan@gmail.com>

Author: seapagan

sqliter/model/model.py

@@ -2,14 +2,23 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Any, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
+
+T = TypeVar("T", bound="BaseDBModel")
 
 
 class BaseDBModel(BaseModel):
     """Custom base model for database models."""
 
+    model_config = ConfigDict(
+        extra="ignore",
+        populate_by_name=True,
+        validate_assignment=False,
+        from_attributes=True,
+    )
+
     class Meta:
         """Configure the base model with default options."""
 
@@ -19,6 +28,15 @@ class Meta:
             None  # Table name, defaults to class name if not set
         )
 
+    @classmethod
+    def model_validate_partial(cls: type[T], obj: dict[str, Any]) -> T:
+        """Validate a partial model object.
+
+        This would be in the case that we are only returning a subset of the
+        fields.
+        """
+        return cls.model_validate(obj, strict=False)
+
     @classmethod
     def get_table_name(cls) -> str:
         """Get the table name from the Meta, or default to the classname."""

sqliter/query/query.py

@@ -30,15 +30,45 @@
 class QueryBuilder:
     """Functions to build and execute queries for a given model."""
 
-    def __init__(self, db: SqliterDB, model_class: type[BaseDBModel]) -> None:
-        """Initialize the query builder with the database, model class, etc."""
+    def __init__(
+        self,
+        db: SqliterDB,
+        model_class: type[BaseDBModel],
+        fields: Optional[list[str]] = None,
+    ) -> None:
+        """Initialize the query builder.
+
+        Pass the database, model class, and optional fields.
+
+        Args:
+            db: The SqliterDB instance.
+            model_class: The model class to query.
+            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
         self.filters: list[tuple[str, Any, str]] = []
         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."""
+        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."""
@@ -219,15 +249,20 @@ def _execute_query(
         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)
+        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 {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()
 
-        select_fields = fields if not count_only else "COUNT(*)"
-
-        sql = f'SELECT {select_fields} FROM "{self.table_name}"'  # noqa: S608 # nosec
-
         if self.filters:
             sql += f" WHERE {where_clause}"
 
@@ -276,9 +311,17 @@ def fetch_all(self) -> list[BaseDBModel]:
         if not results:
             return []
 
+        if self._fields:
+            return [
+                self.model_class.model_validate_partial(
+                    {field: row[idx] for idx, field in enumerate(self._fields)}
+                )
+                for row in results
+            ]
+
         return [
-            self.model_class(
-                **{
+            self.model_class.model_validate(
+                {
                     field: row[idx]
                     for idx, field in enumerate(self.model_class.model_fields)
                 }

sqliter/sqliter.py

@@ -206,9 +206,21 @@ def delete(
         except sqlite3.Error as exc:
             raise RecordDeletionError(table_name) from exc
 
-    def select(self, model_class: type[BaseDBModel]) -> QueryBuilder:
-        """Start a query for the given model."""
-        return QueryBuilder(self, model_class)
+    def select(
+        self, model_class: type[BaseDBModel], fields: Optional[list[str]] = None
+    ) -> QueryBuilder:
+        """Start a query for the given model.
+
+        Args:
+            model_class: The model class to query.
+            fields: Optional list of field names to select.
+                If None, all fields are selected.
+
+        Returns:
+            QueryBuilder: An instance of QueryBuilder for the given model and
+                fields.
+        """
+        return QueryBuilder(self, model_class, fields)
 
     # --- Context manager methods ---
     def __enter__(self) -> Self: