126 lines
3.6 KiB
Python

import hashlib
import psycopg2
import sys
import os
import pandas as pd
def connect_db():
"""Connect to database
Returns:
psycopg2 connector: psycopg2 postgresql connector
"""
conn = None
try:
conn = psycopg2.connect(
host=os.environ['DB_PATH'],
database=os.environ['DB_NAME'],
user=os.environ['DB_USERNAME'],
password=os.environ['DB_PASSWORD'],
)
except (Exception, psycopg2.DatabaseError) as error:
print(error)
sys.exit(1)
return conn
def get_watchlist(username : str):
"""Read list of tickers/descriptions from database
Args:
username (str): database table prefix
Returns:
Pandas DataFrame: it has two columns - first column is ticker, second column is description
"""
if username:
table_name = f"{username + '_watch_list'}"
else: # username is None, use default table
table_name = "stock_watch_list"
QUERY1 = f'''CREATE TABLE IF NOT EXISTS {table_name}
(
tick character varying(5) NOT NULL,
description text,
PRIMARY KEY (tick)
);'''
QUERY2 = f"INSERT INTO {table_name} SELECT 'SPY', 'SPDR S&P 500 ETF Trust' WHERE NOT EXISTS (SELECT NULL FROM {table_name});"
QUERY3 = f"SELECT * FROM {table_name};"
with connect_db() as conn:
with conn.cursor() as curs:
curs.execute(QUERY1)
curs.execute(QUERY2)
curs.execute(QUERY3)
tuples_list = curs.fetchall()
df = pd.DataFrame(tuples_list)
return df
def remove_from_db(username, tick):
"""Remove a row from database table using ticker as key
Args:
username (str): database table prefix
tick (str): ticker
"""
if username:
table_name = f"{username + '_watch_list'}"
else: # username is None, use default table
table_name = "stock_watch_list"
QUERY = f"DELETE FROM {table_name} WHERE tick = '{tick}';"
with connect_db() as conn:
with conn.cursor() as curs:
curs.execute(QUERY)
def insert_into_db(username : str, tick : str, name : str):
"""Insert ticker and description into database
Args:
username (str): database table prefix - each user has its own list of tickers
tick (str): stock or mutual fund ticker
name (str): company name for stock, series ID for mutual fund
"""
if username:
table_name = f"{username + '_watch_list'}"
else: # username is None, use default table
table_name = "stock_watch_list"
QUERY1 = f'''CREATE TABLE IF NOT EXISTS {table_name}
(
tick character varying(5) NOT NULL,
description text,
PRIMARY KEY (tick)
);'''
QUERY2 = f"INSERT INTO {table_name} SELECT 'SPY', 'SPDR S&P 500 ETF Trust' WHERE NOT EXISTS (SELECT NULL FROM {table_name});"
QUERY3 = f"INSERT INTO {table_name} VALUES ('{tick}', '{name}') ON CONFLICT DO NOTHING;"
with connect_db() as conn:
with conn.cursor() as curs:
curs.execute(QUERY1)
curs.execute(QUERY2)
curs.execute(QUERY3)
def hash_password(password : str):
"""Generate hash from string using sha256
Args:
password (str): any text
Returns:
str: hash string
"""
# Encode the password as bytes
password_bytes = password.encode('utf-8')
# Use SHA-256 hash function to create a hash object
hash_object = hashlib.sha256(password_bytes)
# Get the hexadecimal representation of the hash
password_hash = hash_object.hexdigest()
return password_hash