functions.py

# Loosely following the DigitalOcean SQLite tutorial:
# https://www.digitalocean.com/community/tutorials/how-to-use-the-sqlite3-module-in-python-3

# Python's sqlite3 module comes pre-installed with Python
import sqlite3
from sqlite3 import Connection, Cursor

ANIMAL_TABLE = "animals"


def setup_sqlite_db(db_name: str) -> Connection:
    """The connect function creates a `Connection` object. A connection is
    used to interact with the database.
    """
    print(f"Connecting to database {db_name}...")
    connection = sqlite3.connect(db_name)
    return connection


def db_commit(connection: Connection) -> None:
    """SQLite doesn't persist changes to the database until the `.commit()` method
    is called. As queries are executed (`INSERT`, `CREATE_TABLE`, etc.) the changes
    are stored in memory but not written to disk until the `.commit()` call.
    """
    connection.commit()


def print_total_changes(connection: Connection) -> None:
    """The `total_changes` property on the `Connection` class gives the number
    of database rows that have been changed by this connection.
    """
    print(f"Total changes: {connection.total_changes}")


def get_cursor(connection: Connection) -> Cursor:
    """The cursor function creates a `Cursor` object. A cursor allows us to
    send SQL statements to the database.
    """
    print("Grabbing cursor...")
    return connection.cursor()


def setup_table(
    cursor: Cursor,
    table_name: str = ANIMAL_TABLE,
    safe: bool = True,
    strict: bool = False,
) -> None:
    """The `.execute()` method on the `Cursor` object allows us to execute
    SQL statements against the database. Here we setup our `animals` table
    with columns for the animal's `id`, `name`, `species`, and the `tank_number` to
    indicate which tank they are in.

    Note: Technically in SQLite an `INTEGER PRIMARY KEY` will have its value automatically
    filled in regardless with an unused integer. The `AUTOINCREMENT` keyword just changes
    the `id` assignment algorithm to prevent the reuse of integers over the lifetime of the
    database. In other words, the `AUTOINCREMENT` keyword is to prevent the reuse of row IDs
    from previously deleted rows. This change actually imposes some extra CPU, memory, and
    disk overhead and should generally be avoided unless absolutely necessary. For our small
    example, we'll leave it and the performance difference will be negligible.
    """
    column_schema = "(id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, species TEXT, tank_number INTEGER)"
    safe_str = " IF NOT EXISTS " if safe else " "
    strict_str = " STRICT" if strict else ""
    query = f"CREATE TABLE{safe_str}{table_name} {column_schema}{strict_str};"

    print(f"Table creation query:\n\t{query}")
    cursor.execute(query)


def add_data(
    cursor: Cursor, data: dict[str, dict[str, str | int]], table_name: str = ANIMAL_TABLE
) -> None:
    """We can use the `.execute()` method to execute SQL statements which insert data
    into a table in the database. In each query we pass in the values we want to create
    a new database row out of. We don't include the `id` value because we have the `id`
    set as an autoincrementing value that is automatically assigned in ascending order.
    The first row will have an id of `1`, second row will have an id of `2`, and so on.
    """
    for name, info in data.items():
        species = info["species"]
        tank_number = info["tank_number"]
        print(
            f"Adding data for row: ({name} ({type(name)}), "
            f"{species} ({type(species)}), {tank_number} ({type(tank_number)}))"
        )

        # Add to database
        cursor.execute(
            f"INSERT INTO {table_name} (name, species, tank_number) VALUES (?, ?, ?)",
            (name, species, tank_number),
        )


def get_all(cursor: Cursor, table_name: str = ANIMAL_TABLE) -> None:
    """We can write an SQL query to return all data in a table. In this case there
    are two ways we can iterate over the found rows:

    1. We can use the `.fetchall()` method to read all records into memory
    and then returns the list. Can be categorized as eager execution.
    2. We can also diretly iterate over the cursor with a `for` loop. This
    method only reads each row into memory as we reach that iteration of the
    loop. Can be categorized as lazy loading.

    Here we use the direct `for` loop method.
    """
    print("\nAll Animals:")
    cursor.execute(f"SELECT * FROM {table_name};")
    for idx, row in enumerate(cursor):
        print(f"\tRow: {idx + 1}: {row}")