testsurveydb.py

"""
Roger Erismann 2024
Database Functionality Tests

This script demonstrates a series of tests for interacting with a PostgreSQL database hosted in a docker container.
It includes functionality for connecting to the database, creating and modifying a test
table, inserting data, and querying results. The script is designed to verify basic
database operations and ensure the database connection is functioning as expected.

Features:
1. Establish a connection to a PostgreSQL database.
2. Recreate a test table (`test_table`).
3. Insert initial rows of data into the table.
4. Query and display the contents of the table.
5. Add an additional row and verify the update.
6. Safely close the database connection after operations.

Functions:
- `connect_to_db()`: Establish a connection to the database.
- `recreate_table(conn)`: Drop the `test_table` if it exists and create a new table.
- `insert_initial_rows(conn)`: Insert five predefined rows into the table.
- `query_table(conn)`: Query and display all rows in the table.
- `insert_sixth_row(conn)`: Insert an additional row into the table.
- `main()`: Execute the workflow of connecting, modifying, and querying the database.

Dependencies:
- `psycopg2` library for connecting and executing queries in PostgreSQL.
- A PostgreSQL database running on `localhost:5432` with the credentials provided in `DB_CONFIG`.

Ensure that the PostgreSQL server is running and accessible with the provided credentials.

love what you do
"""




import psycopg2

# Database configuration
DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "user": "hammerdirt",
    "password": "hammerdirt",
    "dbname": "surveydb"
}

def connect_to_db():
    """Establish a connection to the PostgreSQL database."""
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        print("Connected to the database!")
        return conn
    except Exception as e:
        print(f"Error connecting to the database: {e}")
        return None

def recreate_table(conn):
    """Drop the table if it exists and create a new one."""
    drop_query = "DROP TABLE IF EXISTS test_table;"
    create_query = """
    CREATE TABLE test_table (
        id SERIAL PRIMARY KEY,
        name VARCHAR(100),
        age INT
    );
    """
    try:
        with conn.cursor() as cursor:
            cursor.execute(drop_query)  # Drop the table if it exists
            cursor.execute(create_query)  # Create a new table
            conn.commit()
        print("Table 'test_table' recreated successfully!")
    except Exception as e:
        print(f"Error recreating table: {e}")
        conn.rollback()

def insert_initial_rows(conn):
    """Insert five rows into the test table."""
    query = """
    INSERT INTO test_table (name, age)
    VALUES 
        ('Alice', 25),
        ('Bob', 30),
        ('Charlie', 35),
        ('Diana', 28),
        ('Eve', 40);
    """
    try:
        with conn.cursor() as cursor:
            cursor.execute(query)
            conn.commit()
        print("Five rows inserted into 'test_table'.")
    except Exception as e:
        print(f"Error inserting rows: {e}")
        conn.rollback()

def query_table(conn):
    """Query the table and return the results."""
    query = "SELECT * FROM test_table;"
    try:
        with conn.cursor() as cursor:
            cursor.execute(query)
            rows = cursor.fetchall()
        print("Current table contents:")
        for row in rows:
            print(row)
    except Exception as e:
        print(f"Error querying table: {e}")
        conn.rollback()

def insert_sixth_row(conn):
    """Insert a sixth row into the table."""
    query = """
    INSERT INTO test_table (name, age)
    VALUES ('Frank', 33);
    """
    try:
        with conn.cursor() as cursor:
            cursor.execute(query)
            conn.commit()
        print("Sixth row inserted into 'test_table'.")
    except Exception as e:
        print(f"Error inserting sixth row: {e}")
        conn.rollback()

def main():
    # Connect to the database
    conn = connect_to_db()
    if not conn:
        return

    try:
        # Drop and recreate the table
        recreate_table(conn)

        # Insert initial data
        insert_initial_rows(conn)

        # Query the table
        query_table(conn)

        # Add a sixth row
        insert_sixth_row(conn)

        # Query again to confirm the sixth row
        query_table(conn)
    finally:
        conn.close()
        print("Database connection closed.")

if __name__ == "__main__":
    main()