sqlite.py

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

from functions import (
    ANIMAL_TABLE,
    setup_sqlite_db,
    db_commit,
    print_total_changes,
    get_cursor,
    setup_table,
    add_data,
    get_all,
)
from argparse import ArgumentParser, Namespace


def parse_args() -> Namespace:
    parser = ArgumentParser()
    parser.add_argument(
        "--mode",
        type=int,
        default=2,
        help="0 to run full db setup, 1 to add additional data, 2 to just run read queries",
    )
    return parser.parse_args()


def main() -> None:

    options = parse_args()
    full = options.mode

    db_con = setup_sqlite_db(db_name="aquarium.db")
    # print_total_changes(connection=db_con)
    cursor = get_cursor(connection=db_con)

    try:
        match full:
            case 0:
                setup_table(cursor=cursor, safe=True, strict=True)
                # setup_table(cursor=cursor, safe=False, strict=True)

                initial_data: dict[str, dict[str, str | int]] = {
                    "Sammy": {"species": "shark", "tank_number": 1},
                    "Jamie": {"species": "cuttlefish", "tank_number": 7},
                }
                add_data(cursor=cursor, data=initial_data)

                db_commit(connection=db_con)

                invalid_data: dict[str, dict[str, str | int]] = {
                    "Spike": {"species": "pufferfish", "tank_number": "two"}
                }
                # add_data(cursor=cursor, data=invalid_data)
            case 1:
                all_data: dict[str, dict[str, str | int]] = {
                    "Spike": {"species": "pufferfish", "tank_number": 2},
                    "Spot": {"species": "turtle", "tank_number": 3},
                    "George": {"species": "jellyfish", "tank_number": 5},
                    "Sophie": {"species": "jellyfish", "tank_number": 5},
                    "Lulu": {"species": "remora", "tank_number": 1},
                    "Ada": {"species": "remora", "tank_number": 1},
                    "Jake": {"species": "remora", "tank_number": 1},
                    "Tag": {"species": "stringray", "tank_number": 1},
                    "Debbie": {"species": "moray eel", "tank_number": 1},
                }
                add_data(cursor=cursor, data=all_data)

                db_commit(connection=db_con)
            case 2:
                # Find all the animals in tank 1
                query = f"SELECT * FROM {ANIMAL_TABLE} WHERE tank_number = 1"
                cursor.execute(query)
                tank1_animals = cursor.fetchall()
                print(f"\nThere are {len(tank1_animals)} animals in tank 1:")
                for idx, row in enumerate(tank1_animals):
                    print(f"\tAnimal #{idx + 1}: {row}")

                # Find the number of animals in each tank
                query = f"SELECT tank_number, COUNT(*) FROM {ANIMAL_TABLE} GROUP BY tank_number"
                cursor.execute(query)
                tank_counts = cursor.fetchall()
                print("\nNumber of animals in each tank:")
                for tank, count in tank_counts:
                    print(f"\tTank {tank}: {count} animals")

                # Find the most common species
                query = f"""
                    SELECT species, COUNT(*) AS count 
                    FROM {ANIMAL_TABLE} 
                    GROUP BY species 
                    ORDER BY count DESC 
                    LIMIT 1
                """
                cursor.execute(query)
                most_common_species = cursor.fetchone()
                print(
                    f"\nMost common species: {most_common_species[0]} ({most_common_species[1]} animals)"
                )

                # List all species in alphabetic order
                query = (
                    f"SELECT DISTINCT species FROM {ANIMAL_TABLE} ORDER BY species ASC"
                )
                cursor.execute(query)
                species_list = cursor.fetchall()
                print("\nList of all species in the aquarium:")
                for idx, species in enumerate(species_list):
                    print(f"\t{idx + 1}. {species[0].title()}")

                # Move an animal to a different tank
                animal_name = "Jamie"
                new_tank = 2
                query = f"UPDATE {ANIMAL_TABLE} SET tank_number = ? WHERE name = ?"
                cursor.execute(query, (new_tank, animal_name))
                db_con.commit()
                print(f"\nMoved {animal_name} to Tank {new_tank}")

        get_all(cursor=cursor)

    finally:
        cursor.close()
        db_con.close()


if __name__ == "__main__":
    main()