db.py

import psycopg2
import psycopg2.extras
from pprint import pprint
import os

# We are deploying to Heroku!!!
class DatabaseConnection:
    def __init__(self):

        

        if os.getenv('DB_NAME') == 'test_db':
            self.db_name = 'test_db'
        else:
            self.db_name = 'learn_db'

        try:
            self.connection = psycopg2.connect(
                dbname=self.db_name, user='postgres', host='localhost', password='kengo1234', port=5432
            )
            
            self.connection.autocommit = True
            self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)

            print('Connected to the database successfully.')
            # print(self.db_name)

            create_users_table = "CREATE TABLE IF NOT EXISTS users (userId SERIAL NOT NULL PRIMARY KEY, username TEXT NOT NULL, email TEXT NOT NULL, password TEXT NOT NULL);"

            self.cursor.execute(create_users_table)
        except Exception as e:
            pprint(e)
            pprint('Failed to connect to the database.')

    def register_user(self, username, email, password):
        """
        Register a user
        """
        reg_user = f"INSERT INTO users(username, email, password) VALUES('{username}', '{email}', '{password}');"
        pprint(reg_user)
        self.cursor.execute(reg_user)
    
    def check_username(self, username):
        """
        Check if a username already exists. Usernames are unique
        """
        query = f"SELECT * FROM users WHERE username='{username}';"
        pprint(query)
        self.cursor.execute(query)
        user = self.cursor.fetchone()
        return user
    
    def check_email(self, email):
        """
        Check if a email already exists. Emails are unique
        """
        query = f"SELECT * FROM users WHERE email='{email}';"
        pprint(query)
        self.cursor.execute(query)
        user = self.cursor.fetchone()
        return user

    def login(self, username):
        """
        Get login credentials for comarison to login
        """
        query = f"SELECT * FROM users WHERE username='{username}';"
        pprint(query)
        self.cursor.execute(query)
        user = self.cursor.fetchone()
        pprint(user)
        return user

    def drop_table(self, table_name):
        """
        Drop tables after tests
        """
        drop = f"DROP TABLE {table_name};"
        self.cursor.execute(drop)


if __name__ == '__main__':
    db_name = DatabaseConnection()