sqliter-py 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

sqliter/sqliter.py

@@ -27,9 +27,24 @@ if TYPE_CHECKING:  # pragma: no cover
 class SqliterDB:
     """Class to manage SQLite database interactions."""
 
-    def __init__(self, db_filename: str, *, auto_commit: bool = False) -> None:
+    def __init__(
+        self,
+        db_filename: Optional[str] = None,
+        *,
+        memory: bool = False,
+        auto_commit: bool = True,
+    ) -> None:
         """Initialize the class and options."""
-        self.db_filename = db_filename
+        if memory:
+            self.db_filename = ":memory:"
+        elif db_filename:
+            self.db_filename = db_filename
+        else:
+            err = (
+                "Database name must be provided if not using an in-memory "
+                "database."
+            )
+            raise ValueError(err)
         self.auto_commit = auto_commit
         self.conn: Optional[sqlite3.Connection] = None
 
@@ -42,17 +57,29 @@ class SqliterDB:
                 raise DatabaseConnectionError(self.db_filename) from exc
         return self.conn
 
+    def close(self) -> None:
+        """Close the connection to the SQLite database."""
+        if self.conn:
+            self._maybe_commit()
+            self.conn.close()
+            self.conn = None
+
+    def commit(self) -> None:
+        """Commit any pending transactions."""
+        if self.conn:
+            self.conn.commit()
+
     def create_table(self, model_class: type[BaseDBModel]) -> None:
         """Create a table based on the Pydantic model."""
         table_name = model_class.get_table_name()
         primary_key = model_class.get_primary_key()
-        create_id = model_class.should_create_id()
+        create_pk = model_class.should_create_pk()
 
         fields = ", ".join(
             f"{field_name} TEXT" for field_name in model_class.model_fields
         )
 
-        if create_id:
+        if create_pk:
             create_table_sql = f"""
                 CREATE TABLE IF NOT EXISTS {table_name} (
                     {primary_key} INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -75,32 +102,33 @@ class SqliterDB:
         except sqlite3.Error as exc:
             raise TableCreationError(table_name) from exc
 
-    def _maybe_commit(self, conn: sqlite3.Connection) -> None:
+    def _maybe_commit(self) -> None:
         """Commit changes if auto_commit is True."""
-        if self.auto_commit:
-            conn.commit()
+        if self.auto_commit and self.conn:
+            self.conn.commit()
 
     def insert(self, model_instance: BaseDBModel) -> None:
         """Insert a new record into the table defined by the Pydantic model."""
         model_class = type(model_instance)
         table_name = model_class.get_table_name()
 
-        fields = ", ".join(model_class.model_fields)
-        placeholders = ", ".join(["?"] * len(model_class.model_fields))
-        values = tuple(
-            getattr(model_instance, field) for field in model_class.model_fields
+        data = model_instance.model_dump()
+        fields = ", ".join(data.keys())
+        placeholders = ", ".join(
+            ["?" if value is not None else "NULL" for value in data.values()]
         )
+        values = tuple(value for value in data.values() if value is not None)
 
         insert_sql = f"""
         INSERT INTO {table_name} ({fields})
         VALUES ({placeholders})
-    """  # noqa: S608
+        """  # noqa: S608
 
         try:
             with self.connect() as conn:
                 cursor = conn.cursor()
                 cursor.execute(insert_sql, values)
-                self._maybe_commit(conn)
+                self._maybe_commit()
         except sqlite3.Error as exc:
             raise RecordInsertionError(table_name) from exc
 
@@ -167,7 +195,7 @@ class SqliterDB:
                 if cursor.rowcount == 0:
                     raise RecordNotFoundError(primary_key_value)
 
-                self._maybe_commit(conn)
+                self._maybe_commit()
 
         except sqlite3.Error as exc:
             raise RecordUpdateError(table_name) from exc
@@ -190,13 +218,36 @@ class SqliterDB:
 
                 if cursor.rowcount == 0:
                     raise RecordNotFoundError(primary_key_value)
-                self._maybe_commit(conn)
+                self._maybe_commit()
         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,
+        exclude: 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.
+            exclude: Optional list of field names to exclude from the query
+                output.
+
+        Returns:
+            QueryBuilder: An instance of QueryBuilder for the given model and
+            fields.
+        """
+        query_builder = QueryBuilder(self, model_class, fields)
+
+        # If exclude is provided, apply the exclude method
+        if exclude:
+            query_builder.exclude(exclude)
+
+        return query_builder
 
     # --- Context manager methods ---
     def __enter__(self) -> Self:
@@ -216,7 +267,8 @@ class SqliterDB:
                 if exc_type:
                     # Roll back the transaction if there was an exception
                     self.conn.rollback()
-                self._maybe_commit(self.conn)
+                else:
+                    self.conn.commit()
             finally:
                 # Close the connection and reset the instance variable
                 self.conn.close()