diff --git a/database.py b/database.py index 982626f..f8e63ae 100644 --- a/database.py +++ b/database.py @@ -1,161 +1,272 @@ +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 = sqlite3.connect(user_db_file_location) + _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 = sqlite3.connect(user_db_file_location) + _conn = get_connection(user_db_file_location) _c = _conn.cursor() - _c.execute("SELECT pw FROM users WHERE id = '" + id + "';") - result = _c.fetchone()[0] == hashlib.sha256(pw.encode()).hexdigest() - + _c.execute( + f"SELECT pw FROM users WHERE id = {placeholder()};", + (id.upper(),) + ) + + row = _c.fetchone() _conn.close() - return result + if row is None: + return False + + return row[0] == hashlib.sha256(pw.encode()).hexdigest() + def delete_user_from_db(id): - _conn = sqlite3.connect(user_db_file_location) + user_id = id.upper() + + _conn = get_connection(user_db_file_location) _c = _conn.cursor() - _c.execute("DELETE FROM users WHERE id = ?;", (id)) + _c.execute( + f"DELETE FROM users WHERE id = {placeholder()};", + (user_id,) + ) _conn.commit() _conn.close() - # when we delete a user FROM database USERS, we also need to delete all his or her notes data FROM database NOTES - _conn = sqlite3.connect(note_db_file_location) + # 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("DELETE FROM notes WHERE user = ?;", (id)) + _c.execute( + f'DELETE FROM notes WHERE "user" = {placeholder()};', + (user_id,) + ) _conn.commit() _conn.close() - # when we delete a user FROM database USERS, we also need to - # [1] delete all his or her images FROM image pool (done in app.py) - # [2] delete all his or her images records FROM database IMAGES - _conn = sqlite3.connect(image_db_file_location) + # 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("DELETE FROM images WHERE owner = ?;", (id)) + _c.execute( + f"DELETE FROM images WHERE owner = {placeholder()};", + (user_id,) + ) _conn.commit() _conn.close() + def add_user(id, pw): - _conn = sqlite3.connect(user_db_file_location) + _conn = get_connection(user_db_file_location) _c = _conn.cursor() - _c.execute("INSERT INTO users values(?, ?)", (id.upper(), hashlib.sha256(pw.encode()).hexdigest())) - + _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 = sqlite3.connect(note_db_file_location) + _conn = get_connection(note_db_file_location) _c = _conn.cursor() - command = "SELECT note_id, timestamp, note FROM notes WHERE user = '" + id.upper() + "';" - _c.execute(command) + _c.execute( + f'SELECT note_id, timestamp, note FROM notes WHERE "user" = {placeholder()};', + (id.upper(),) + ) result = _c.fetchall() - _conn.commit() _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 which is being operated. - _conn = sqlite3.connect(note_db_file_location) + # 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() - command = "SELECT user FROM notes WHERE note_id = '" + note_id + "';" - _c.execute(command) - result = _c.fetchone()[0] + _c.execute( + f'SELECT "user" FROM notes WHERE note_id = {placeholder()};', + (note_id,) + ) - _conn.commit() + row = _c.fetchone() _conn.close() - return result + if row is None: + return None + + return row[0] + def write_note_into_db(id, note_to_write): - _conn = sqlite3.connect(note_db_file_location) + _conn = get_connection(note_db_file_location) _c = _conn.cursor() current_timestamp = str(datetime.datetime.now()) - _c.execute("INSERT INTO notes values(?, ?, ?, ?)", (id.upper(), current_timestamp, note_to_write, hashlib.sha1((id.upper() + current_timestamp).encode()).hexdigest())) + 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 = sqlite3.connect(note_db_file_location) + _conn = get_connection(note_db_file_location) _c = _conn.cursor() - _c.execute("DELETE FROM notes WHERE note_id = ?;", (note_id)) + _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 = sqlite3.connect(image_db_file_location) + _conn = get_connection(image_db_file_location) _c = _conn.cursor() - _c.execute("INSERT INTO images VALUES (?, ?, ?, ?)", (uid, owner, image_name, timestamp)) + _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 = sqlite3.connect(image_db_file_location) + _conn = get_connection(image_db_file_location) _c = _conn.cursor() - command = "SELECT uid, timestamp, name FROM images WHERE owner = '{0}'".format(owner) - _c.execute(command) + _c.execute( + f"SELECT uid, timestamp, name FROM images WHERE owner = {placeholder()};", + (owner,) + ) result = _c.fetchall() - _conn.commit() _conn.close() return result + def match_user_id_with_image_uid(image_uid): - # Given the note id, confirm if the current user is the owner of the note which is being operated. - _conn = sqlite3.connect(image_db_file_location) + # 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() - command = "SELECT owner FROM images WHERE uid = '" + image_uid + "';" - _c.execute(command) - result = _c.fetchone()[0] + _c.execute( + f"SELECT owner FROM images WHERE uid = {placeholder()};", + (image_uid,) + ) - _conn.commit() + row = _c.fetchone() _conn.close() - return result + if row is None: + return None + + return row[0] + def delete_image_from_db(image_uid): - _conn = sqlite3.connect(image_db_file_location) + _conn = get_connection(image_db_file_location) _c = _conn.cursor() - _c.execute("DELETE FROM images WHERE uid = ?;", (image_uid)) + _c.execute( + f"DELETE FROM images WHERE uid = {placeholder()};", + (image_uid,) + ) _conn.commit() _conn.close() - - - - - if __name__ == "__main__": print(list_users()) + diff --git a/manifests/prod/postgres.yaml b/manifests/prod/postgres.yaml index 3b9cf54..6278f8d 100644 --- a/manifests/prod/postgres.yaml +++ b/manifests/prod/postgres.yaml @@ -26,11 +26,11 @@ spec: apiVersion: v1 kind: Service metadata: - name: postgres + name: postgres-service namespace: prod spec: selector: - app: postgres-service + app: postgres ports: - port: 5432 targetPort: 5432 diff --git a/requirements.txt b/requirements.txt index bad38d9..f5f276f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ itsdangerous==2.2.0 Jinja2==3.1.6 MarkupSafe==3.0.3 Werkzeug==3.1.8 +psycopg2-binary==2.9.11