126 lines
3.6 KiB
Python
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
|