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

sqliter/sqliter.py

@@ -1,21 +1,34 @@
-"""This is the main module for the sqliter package."""
+"""Core module for SQLiter, providing the main database interaction class.
+
+This module defines the SqliterDB class, which serves as the primary
+interface for all database operations in SQLiter. It handles connection
+management, table creation, and CRUD operations, bridging the gap between
+Pydantic models and SQLite database interactions.
+"""
 
 from __future__ import annotations
 
+import logging
 import sqlite3
-from typing import TYPE_CHECKING, Optional
+import time
+from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
 from typing_extensions import Self
 
 from sqliter.exceptions import (
     DatabaseConnectionError,
+    InvalidIndexError,
     RecordDeletionError,
     RecordFetchError,
     RecordInsertionError,
     RecordNotFoundError,
     RecordUpdateError,
+    SqlExecutionError,
     TableCreationError,
+    TableDeletionError,
 )
+from sqliter.helpers import infer_sqlite_type
+from sqliter.model.unique import Unique
 from sqliter.query.query import QueryBuilder
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -23,20 +36,52 @@ if TYPE_CHECKING:  # pragma: no cover
 
     from sqliter.model.model import BaseDBModel
 
+T = TypeVar("T", bound="BaseDBModel")
+
 
 class SqliterDB:
-    """Class to manage SQLite database interactions."""
+    """Main class for interacting with SQLite databases.
+
+    This class provides methods for connecting to a SQLite database,
+    creating tables, and performing CRUD operations.
+
+    Arguements:
+        db_filename (str): The filename of the SQLite database.
+        auto_commit (bool): Whether to automatically commit transactions.
+        debug (bool): Whether to enable debug logging.
+        logger (Optional[logging.Logger]): Custom logger for debug output.
+    """
+
+    MEMORY_DB = ":memory:"
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         db_filename: Optional[str] = None,
         *,
         memory: bool = False,
         auto_commit: bool = True,
+        debug: bool = False,
+        logger: Optional[logging.Logger] = None,
+        reset: bool = False,
+        return_local_time: bool = True,
     ) -> None:
-        """Initialize the class and options."""
+        """Initialize a new SqliterDB instance.
+
+        Args:
+            db_filename: The filename of the SQLite database.
+            memory: If True, create an in-memory database.
+            auto_commit: Whether to automatically commit transactions.
+            debug: Whether to enable debug logging.
+            logger: Custom logger for debug output.
+            reset: Whether to reset the database on initialization. This will
+                basically drop all existing tables.
+            return_local_time: Whether to return local time for datetime fields.
+
+        Raises:
+            ValueError: If no filename is provided for a non-memory database.
+        """
         if memory:
-            self.db_filename = ":memory:"
+            self.db_filename = self.MEMORY_DB
         elif db_filename:
             self.db_filename = db_filename
         else:
@@ -46,10 +91,149 @@ class SqliterDB:
             )
             raise ValueError(err)
         self.auto_commit = auto_commit
+        self.debug = debug
+        self.logger = logger
         self.conn: Optional[sqlite3.Connection] = None
+        self.reset = reset
+        self.return_local_time = return_local_time
+
+        self._in_transaction = False
+
+        if self.debug:
+            self._setup_logger()
+
+        if self.reset:
+            self._reset_database()
+
+    @property
+    def filename(self) -> Optional[str]:
+        """Returns the filename of the current database or None if in-memory."""
+        return None if self.db_filename == self.MEMORY_DB else self.db_filename
+
+    @property
+    def is_memory(self) -> bool:
+        """Returns True if the database is in-memory."""
+        return self.db_filename == self.MEMORY_DB
+
+    @property
+    def is_autocommit(self) -> bool:
+        """Returns True if auto-commit is enabled."""
+        return self.auto_commit
+
+    @property
+    def is_connected(self) -> bool:
+        """Returns True if the database is connected, False otherwise."""
+        return self.conn is not None
+
+    @property
+    def table_names(self) -> list[str]:
+        """Returns a list of all table names in the database.
+
+        Temporarily connects to the database if not connected and restores
+        the connection state afterward.
+        """
+        was_connected = self.is_connected
+        if not was_connected:
+            self.connect()
+
+        if self.conn is None:
+            err_msg = "Failed to establish a database connection."
+            raise DatabaseConnectionError(err_msg)
+
+        cursor = self.conn.cursor()
+        cursor.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' "
+            "AND name NOT LIKE 'sqlite_%';"
+        )
+        tables = [row[0] for row in cursor.fetchall()]
+
+        # Restore the connection state
+        if not was_connected:
+            self.close()
+
+        return tables
+
+    def _reset_database(self) -> None:
+        """Drop all user-created tables in the database."""
+        with self.connect() as conn:
+            cursor = conn.cursor()
+
+            # Get all table names, excluding SQLite system tables
+            cursor.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' "
+                "AND name NOT LIKE 'sqlite_%';"
+            )
+            tables = cursor.fetchall()
+
+            # Drop each user-created table
+            for table in tables:
+                cursor.execute(f"DROP TABLE IF EXISTS {table[0]}")
+
+            conn.commit()
+
+        if self.debug and self.logger:
+            self.logger.debug(
+                "Database reset: %s user-created tables dropped.", len(tables)
+            )
+
+    def _setup_logger(self) -> None:
+        """Set up the logger for debug output.
+
+        This method configures a logger for the SqliterDB instance, either
+        using an existing logger or creating a new one specifically for
+        SQLiter.
+        """
+        # Check if the root logger is already configured
+        root_logger = logging.getLogger()
+
+        if root_logger.hasHandlers():
+            # If the root logger has handlers, use it without modifying the root
+            # configuration
+            self.logger = root_logger.getChild("sqliter")
+        else:
+            # If no root logger is configured, set up a new logger specific to
+            # SqliterDB
+            self.logger = logging.getLogger("sqliter")
+
+            handler = logging.StreamHandler()  # Output to console
+            formatter = logging.Formatter(
+                "%(levelname)-8s%(message)s"
+            )  # Custom format
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+
+            self.logger.setLevel(logging.DEBUG)
+            self.logger.propagate = False
+
+    def _log_sql(self, sql: str, values: list[Any]) -> None:
+        """Log the SQL query and its values if debug mode is enabled.
+
+        The values are inserted into the SQL query string to replace the
+        placeholders.
+
+        Args:
+            sql: The SQL query string.
+            values: The list of values to be inserted into the query.
+        """
+        if self.debug and self.logger:
+            formatted_sql = sql
+            for value in values:
+                if isinstance(value, str):
+                    formatted_sql = formatted_sql.replace("?", f"'{value}'", 1)
+                else:
+                    formatted_sql = formatted_sql.replace("?", str(value), 1)
+
+            self.logger.debug("Executing SQL: %s", formatted_sql)
 
     def connect(self) -> sqlite3.Connection:
-        """Create or return a connection to the SQLite database."""
+        """Establish a connection to the SQLite database.
+
+        Returns:
+            The SQLite connection object.
+
+        Raises:
+            DatabaseConnectionError: If unable to connect to the database.
+        """
         if not self.conn:
             try:
                 self.conn = sqlite3.connect(self.db_filename)
@@ -58,41 +242,77 @@ class SqliterDB:
         return self.conn
 
     def close(self) -> None:
-        """Close the connection to the SQLite database."""
+        """Close the database connection.
+
+        This method commits any pending changes if auto_commit is True,
+        then closes the connection. If the connection is already closed or does
+        not exist, this method silently does nothing.
+        """
         if self.conn:
             self._maybe_commit()
             self.conn.close()
             self.conn = None
 
     def commit(self) -> None:
-        """Commit any pending transactions."""
+        """Commit the current transaction.
+
+        This method explicitly commits any pending changes to the database.
+        """
         if self.conn:
             self.conn.commit()
 
-    def create_table(self, model_class: type[BaseDBModel]) -> None:
-        """Create a table based on the Pydantic model."""
+    def create_table(
+        self,
+        model_class: type[BaseDBModel],
+        *,
+        exists_ok: bool = True,
+        force: bool = False,
+    ) -> None:
+        """Create a table in the database based on the given model class.
+
+        Args:
+            model_class: The Pydantic model class representing the table.
+            exists_ok: If True, do not raise an error if the table already
+                exists. Default is True which is the original behavior.
+            force: If True, drop the table if it exists before creating.
+                Defaults to False.
+
+        Raises:
+            TableCreationError: If there's an error creating the table.
+            ValueError: If the primary key field is not found in the model.
+        """
         table_name = model_class.get_table_name()
         primary_key = model_class.get_primary_key()
-        create_pk = model_class.should_create_pk()
 
-        fields = ", ".join(
-            f"{field_name} TEXT" for field_name in model_class.model_fields
-        )
+        if force:
+            drop_table_sql = f"DROP TABLE IF EXISTS {table_name}"
+            self._execute_sql(drop_table_sql)
+
+        fields = [f'"{primary_key}" INTEGER PRIMARY KEY AUTOINCREMENT']
 
-        if create_pk:
-            create_table_sql = f"""
-                CREATE TABLE IF NOT EXISTS {table_name} (
-                    {primary_key} INTEGER PRIMARY KEY AUTOINCREMENT,
-                    {fields}
+        # Add remaining fields
+        for field_name, field_info in model_class.model_fields.items():
+            if field_name != primary_key:
+                sqlite_type = infer_sqlite_type(field_info.annotation)
+                unique_constraint = (
+                    "UNIQUE" if isinstance(field_info, Unique) else ""
                 )
-            """
-        else:
-            create_table_sql = f"""
-                CREATE TABLE IF NOT EXISTS {table_name} (
-                    {fields},
-                    PRIMARY KEY ({primary_key})
+                fields.append(
+                    f"{field_name} {sqlite_type} {unique_constraint}".strip()
                 )
-            """
+
+        create_str = (
+            "CREATE TABLE IF NOT EXISTS" if exists_ok else "CREATE TABLE"
+        )
+
+        create_table_sql = f"""
+        {create_str} {table_name} (
+            {", ".join(fields)}
+        )
+        """
+
+        if self.debug:
+            self._log_sql(create_table_sql, [])
 
         try:
             with self.connect() as conn:
@@ -102,17 +322,166 @@ class SqliterDB:
         except sqlite3.Error as exc:
             raise TableCreationError(table_name) from exc
 
+        # Create regular indexes
+        if hasattr(model_class.Meta, "indexes"):
+            self._create_indexes(
+                model_class, model_class.Meta.indexes, unique=False
+            )
+
+        # Create unique indexes
+        if hasattr(model_class.Meta, "unique_indexes"):
+            self._create_indexes(
+                model_class, model_class.Meta.unique_indexes, unique=True
+            )
+
+    def _create_indexes(
+        self,
+        model_class: type[BaseDBModel],
+        indexes: list[Union[str, tuple[str]]],
+        *,
+        unique: bool = False,
+    ) -> None:
+        """Helper method to create regular or unique indexes.
+
+        Args:
+            model_class: The model class defining the table.
+            indexes: List of fields or tuples of fields to create indexes for.
+            unique: If True, creates UNIQUE indexes; otherwise, creates regular
+                indexes.
+
+        Raises:
+            InvalidIndexError: If any fields specified for indexing do not exist
+                in the model.
+        """
+        valid_fields = set(
+            model_class.model_fields.keys()
+        )  # Get valid fields from the model
+
+        for index in indexes:
+            # Handle multiple fields in tuple form
+            fields = list(index) if isinstance(index, tuple) else [index]
+
+            # Check if all fields exist in the model
+            invalid_fields = [
+                field for field in fields if field not in valid_fields
+            ]
+            if invalid_fields:
+                raise InvalidIndexError(invalid_fields, model_class.__name__)
+
+            # Build the SQL string
+            index_name = "_".join(fields)
+            index_postfix = "_unique" if unique else ""
+            index_type = " UNIQUE " if unique else " "
+
+            create_index_sql = (
+                f"CREATE{index_type}INDEX IF NOT EXISTS "
+                f"idx_{model_class.get_table_name()}"
+                f"_{index_name}{index_postfix} "
+                f"ON {model_class.get_table_name()} ({', '.join(fields)})"
+            )
+            self._execute_sql(create_index_sql)
+
+    def _execute_sql(self, sql: str) -> None:
+        """Execute an SQL statement.
+
+        Args:
+            sql: The SQL statement to execute.
+
+        Raises:
+            SqlExecutionError: If the SQL execution fails.
+        """
+        if self.debug:
+            self._log_sql(sql, [])
+
+        try:
+            with self.connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute(sql)
+                conn.commit()
+        except (sqlite3.Error, sqlite3.Warning) as exc:
+            raise SqlExecutionError(sql) from exc
+
+    def drop_table(self, model_class: type[BaseDBModel]) -> None:
+        """Drop the table associated with the given model class.
+
+        Args:
+            model_class: The model class for which to drop the table.
+
+        Raises:
+            TableDeletionError: If there's an error dropping the table.
+        """
+        table_name = model_class.get_table_name()
+        drop_table_sql = f"DROP TABLE IF EXISTS {table_name}"
+
+        if self.debug:
+            self._log_sql(drop_table_sql, [])
+
+        try:
+            with self.connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute(drop_table_sql)
+                self.commit()
+        except sqlite3.Error as exc:
+            raise TableDeletionError(table_name) from exc
+
     def _maybe_commit(self) -> None:
-        """Commit changes if auto_commit is True."""
-        if self.auto_commit and self.conn:
+        """Commit changes if auto_commit is enabled.
+
+        This method is called after operations that modify the database,
+        committing changes only if auto_commit is set to True.
+        """
+        if not self._in_transaction and 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."""
+    def insert(
+        self, model_instance: T, *, timestamp_override: bool = False
+    ) -> T:
+        """Insert a new record into the database.
+
+        Args:
+            model_instance: The instance of the model class to insert.
+            timestamp_override: If True, override the created_at and updated_at
+                timestamps with provided values. Default is False. If the values
+                are not provided, they will be set to the current time as
+                normal. Without this flag, the timestamps will always be set to
+                the current time, even if provided.
+
+        Returns:
+            The updated model instance with the primary key (pk) set.
+
+        Raises:
+            RecordInsertionError: If an error occurs during the insertion.
+        """
         model_class = type(model_instance)
         table_name = model_class.get_table_name()
 
+        # Always set created_at and updated_at timestamps
+        current_timestamp = int(time.time())
+
+        # Handle the case where timestamp_override is False
+        if not timestamp_override:
+            # Always override both timestamps with the current time
+            model_instance.created_at = current_timestamp
+            model_instance.updated_at = current_timestamp
+        else:
+            # Respect provided values, but set to current time if they are 0
+            if model_instance.created_at == 0:
+                model_instance.created_at = current_timestamp
+            if model_instance.updated_at == 0:
+                model_instance.updated_at = current_timestamp
+
+        # Get the data from the model
         data = model_instance.model_dump()
+
+        # Serialize the data
+        for field_name, value in list(data.items()):
+            data[field_name] = model_instance.serialize_field(value)
+
+        # remove the primary key field if it exists, otherwise we'll get
+        # TypeErrors as multiple primary keys will exist
+        if data.get("pk", None) == 0:
+            data.pop("pk")
+
         fields = ", ".join(data.keys())
         placeholders = ", ".join(
             ["?" if value is not None else "NULL" for value in data.values()]
@@ -129,13 +498,34 @@ class SqliterDB:
                 cursor = conn.cursor()
                 cursor.execute(insert_sql, values)
                 self._maybe_commit()
+
         except sqlite3.Error as exc:
             raise RecordInsertionError(table_name) from exc
+        else:
+            data.pop("pk", None)
+            # Deserialize each field before creating the model instance
+            deserialized_data = {}
+            for field_name, value in data.items():
+                deserialized_data[field_name] = model_class.deserialize_field(
+                    field_name, value, return_local_time=self.return_local_time
+                )
+            return model_class(pk=cursor.lastrowid, **deserialized_data)
 
     def get(
-        self, model_class: type[BaseDBModel], primary_key_value: str
+        self, model_class: type[BaseDBModel], primary_key_value: int
     ) -> BaseDBModel | None:
-        """Retrieve a record by its PK and return a Pydantic instance."""
+        """Retrieve a single record from the database by its primary key.
+
+        Args:
+            model_class: The Pydantic model class representing the table.
+            primary_key_value: The value of the primary key to look up.
+
+        Returns:
+            An instance of the model class if found, None otherwise.
+
+        Raises:
+            RecordFetchError: If there's an error fetching the record.
+        """
         table_name = model_class.get_table_name()
         primary_key = model_class.get_primary_key()
 
@@ -156,29 +546,51 @@ class SqliterDB:
                     field: result[idx]
                     for idx, field in enumerate(model_class.model_fields)
                 }
-                return model_class(**result_dict)
+                # Deserialize each field before creating the model instance
+                deserialized_data = {}
+                for field_name, value in result_dict.items():
+                    deserialized_data[field_name] = (
+                        model_class.deserialize_field(
+                            field_name,
+                            value,
+                            return_local_time=self.return_local_time,
+                        )
+                    )
+                return model_class(**deserialized_data)
         except sqlite3.Error as exc:
             raise RecordFetchError(table_name) from exc
         else:
             return None
 
     def update(self, model_instance: BaseDBModel) -> None:
-        """Update an existing record using the Pydantic model."""
+        """Update an existing record in the database.
+
+        Args:
+            model_instance: An instance of a Pydantic model to be updated.
+
+        Raises:
+            RecordUpdateError: If there's an error updating the record or if it
+                is not found.
+        """
         model_class = type(model_instance)
         table_name = model_class.get_table_name()
         primary_key = model_class.get_primary_key()
 
-        fields = ", ".join(
-            f"{field} = ?"
-            for field in model_class.model_fields
-            if field != primary_key
-        )
-        values = tuple(
-            getattr(model_instance, field)
-            for field in model_class.model_fields
-            if field != primary_key
-        )
-        primary_key_value = getattr(model_instance, primary_key)
+        # Set updated_at timestamp
+        current_timestamp = int(time.time())
+        model_instance.updated_at = current_timestamp
+
+        # Get the data and serialize any datetime/date fields
+        data = model_instance.model_dump()
+        for field_name, value in list(data.items()):
+            data[field_name] = model_instance.serialize_field(value)
+
+        # Remove the primary key from the update data
+        primary_key_value = data.pop(primary_key)
+
+        # Create the SQL using the processed data
+        fields = ", ".join(f"{field} = ?" for field in data)
+        values = tuple(data.values())
 
         update_sql = f"""
             UPDATE {table_name}
@@ -203,7 +615,17 @@ class SqliterDB:
     def delete(
         self, model_class: type[BaseDBModel], primary_key_value: str
     ) -> None:
-        """Delete a record by its primary key."""
+        """Delete a record from the database by its primary key.
+
+        Args:
+            model_class: The Pydantic model class representing the table.
+            primary_key_value: The value of the primary key of the record to
+                delete.
+
+        Raises:
+            RecordDeletionError: If there's an error deleting the record.
+            RecordNotFoundError: If the record to delete is not found.
+        """
         table_name = model_class.get_table_name()
         primary_key = model_class.get_primary_key()
 
@@ -228,18 +650,15 @@ class SqliterDB:
         fields: Optional[list[str]] = None,
         exclude: Optional[list[str]] = None,
     ) -> QueryBuilder:
-        """Start a query for the given model.
+        """Create a QueryBuilder instance for selecting records.
 
         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.
+            model_class: The Pydantic model class representing the table.
+            fields: Optional list of fields to include in the query.
+            exclude: Optional list of fields to exclude from the query.
 
         Returns:
-            QueryBuilder: An instance of QueryBuilder for the given model and
-            fields.
+            A QueryBuilder instance for further query construction.
         """
         query_builder = QueryBuilder(self, model_class, fields)
 
@@ -251,8 +670,20 @@ class SqliterDB:
 
     # --- Context manager methods ---
     def __enter__(self) -> Self:
-        """Enter the runtime context for the 'with' statement."""
+        """Enter the runtime context for the SqliterDB instance.
+
+        This method is called when entering a 'with' statement. It ensures
+        that a database connection is established.
+
+        Note that this method should never be called explicitly, but will be
+        called by the 'with' statement when entering the context.
+
+        Returns:
+            The SqliterDB instance.
+
+        """
         self.connect()
+        self._in_transaction = True
         return self
 
     def __exit__(
@@ -261,7 +692,24 @@ class SqliterDB:
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> None:
-        """Exit the runtime context and close the connection."""
+        """Exit the runtime context for the SqliterDB instance.
+
+        This method is called when exiting a 'with' statement. It handles
+        committing or rolling back transactions based on whether an exception
+        occurred, and closes the database connection.
+
+        Args:
+            exc_type: The type of the exception that caused the context to be
+                exited, or None if no exception was raised.
+            exc_value: The instance of the exception that caused the context
+                to be exited, or None if no exception was raised.
+            traceback: A traceback object encoding the stack trace, or None
+                if no exception was raised.
+
+        Note that this method should never be called explicitly, but will be
+        called by the 'with' statement when exiting the context.
+
+        """
         if self.conn:
             try:
                 if exc_type:
@@ -273,3 +721,4 @@ class SqliterDB:
                 # Close the connection and reset the instance variable
                 self.conn.close()
                 self.conn = None
+                self._in_transaction = False