devopsexam/database.py

import os
import sqlite3
import hashlib
import datetime
import psycopg2

user_db_file_location = "database_file/users.db"
note_db_file_location = "database_file/notes.db"
image_db_file_location = "database_file/images.db"

_schema_initialized = False


def use_postgres():
    return os.getenv("DATABASE_URL") is not None


def placeholder():
    return "%s" if use_postgres() else "?"


def get_connection(sqlite_db_file=None):
    global _schema_initialized

    database_url = os.getenv("DATABASE_URL")

    if use_postgres():
        conn = psycopg2.connect(database_url)

        if not _schema_initialized:
            init_postgres_schema(conn)
            _schema_initialized = True

        return conn

    return sqlite3.connect(sqlite_db_file)


def init_postgres_schema(conn):
    c = conn.cursor()

    c.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id TEXT PRIMARY KEY,
            pw TEXT NOT NULL
        );
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS notes (
            "user" TEXT NOT NULL,
            timestamp TEXT NOT NULL,
            note TEXT NOT NULL,
            note_id TEXT PRIMARY KEY
        );
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS images (
            uid TEXT PRIMARY KEY,
            owner TEXT NOT NULL,
            name TEXT NOT NULL,
            timestamp TEXT NOT NULL
        );
    """)

    conn.commit()
    c.close()


def list_users():
    _conn = get_connection(user_db_file_location)
    _c = _conn.cursor()

    _c.execute("SELECT id FROM users;")
    result = [x[0] for x in _c.fetchall()]

    _conn.close()

    return result


def verify(id, pw):
    _conn = get_connection(user_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"SELECT pw FROM users WHERE id = {placeholder()};",
        (id.upper(),)
    )

    row = _c.fetchone()
    _conn.close()

    if row is None:
        return False

    return row[0] == hashlib.sha256(pw.encode()).hexdigest()


def delete_user_from_db(id):
    user_id = id.upper()

    _conn = get_connection(user_db_file_location)
    _c = _conn.cursor()
    _c.execute(
        f"DELETE FROM users WHERE id = {placeholder()};",
        (user_id,)
    )
    _conn.commit()
    _conn.close()

    # when we delete a user from USERS, delete all notes owned by the user
    _conn = get_connection(note_db_file_location)
    _c = _conn.cursor()
    _c.execute(
        f'DELETE FROM notes WHERE "user" = {placeholder()};',
        (user_id,)
    )
    _conn.commit()
    _conn.close()

    # when we delete a user from USERS, delete all image records owned by the user
    _conn = get_connection(image_db_file_location)
    _c = _conn.cursor()
    _c.execute(
        f"DELETE FROM images WHERE owner = {placeholder()};",
        (user_id,)
    )
    _conn.commit()
    _conn.close()


def add_user(id, pw):
    _conn = get_connection(user_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"INSERT INTO users VALUES ({placeholder()}, {placeholder()});",
        (id.upper(), hashlib.sha256(pw.encode()).hexdigest())
    )

    _conn.commit()
    _conn.close()


def read_note_from_db(id):
    _conn = get_connection(note_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f'SELECT note_id, timestamp, note FROM notes WHERE "user" = {placeholder()};',
        (id.upper(),)
    )
    result = _c.fetchall()

    _conn.close()

    return result


def match_user_id_with_note_id(note_id):
    # Given the note id, confirm if the current user is the owner of the note being operated.
    _conn = get_connection(note_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f'SELECT "user" FROM notes WHERE note_id = {placeholder()};',
        (note_id,)
    )

    row = _c.fetchone()
    _conn.close()

    if row is None:
        return None

    return row[0]


def write_note_into_db(id, note_to_write):
    _conn = get_connection(note_db_file_location)
    _c = _conn.cursor()

    current_timestamp = str(datetime.datetime.now())
    note_id = hashlib.sha1((id.upper() + current_timestamp).encode()).hexdigest()

    _c.execute(
        f"INSERT INTO notes VALUES ({placeholder()}, {placeholder()}, {placeholder()}, {placeholder()});",
        (id.upper(), current_timestamp, note_to_write, note_id)
    )

    _conn.commit()
    _conn.close()


def delete_note_from_db(note_id):
    _conn = get_connection(note_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"DELETE FROM notes WHERE note_id = {placeholder()};",
        (note_id,)
    )

    _conn.commit()
    _conn.close()


def image_upload_record(uid, owner, image_name, timestamp):
    _conn = get_connection(image_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"INSERT INTO images VALUES ({placeholder()}, {placeholder()}, {placeholder()}, {placeholder()});",
        (uid, owner, image_name, timestamp)
    )

    _conn.commit()
    _conn.close()


def list_images_for_user(owner):
    _conn = get_connection(image_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"SELECT uid, timestamp, name FROM images WHERE owner = {placeholder()};",
        (owner,)
    )
    result = _c.fetchall()

    _conn.close()

    return result


def match_user_id_with_image_uid(image_uid):
    # Given the image uid, confirm if the current user is the owner of the image being operated.
    _conn = get_connection(image_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"SELECT owner FROM images WHERE uid = {placeholder()};",
        (image_uid,)
    )

    row = _c.fetchone()
    _conn.close()

    if row is None:
        return None

    return row[0]


def delete_image_from_db(image_uid):
    _conn = get_connection(image_db_file_location)
    _c = _conn.cursor()

    _c.execute(
        f"DELETE FROM images WHERE uid = {placeholder()};",
        (image_uid,)
    )

    _conn.commit()
    _conn.close()


if __name__ == "__main__":
    print(list_users())