From 174a36f824cf9d515e41fa1e814bfcf27f2756da Mon Sep 17 00:00:00 2001 From: Alex Huddleston Date: Sat, 4 Apr 2020 19:12:28 -0500 Subject: [PATCH] Refactored database integration code to be more modular, updated template config. --- config/template_config.json | 5 ++ covid_bot.py | 7 +-- covid_report.py | 4 +- docker-compose.yml | 4 ++ lib/config_lib.py | 50 +++++++++++++++++ lib/covidBot.py | 8 +-- lib/covidData.py | 26 ++------- lib/covid_data_lib.py | 104 ++++++++++++++++-------------------- lib/parse_data.py | 71 +++++++++--------------- 9 files changed, 145 insertions(+), 134 deletions(-) create mode 100644 lib/config_lib.py diff --git a/config/template_config.json b/config/template_config.json index 25f2b6c..f4d1c8c 100644 --- a/config/template_config.json +++ b/config/template_config.json @@ -1,5 +1,10 @@ { "api_key": "https://discordapp.com/developers/docs/intro", + "postgres_db": "covid_prod", + "postgres_user": "covidbot", + "postgres_host": "postgres", + "postgres_pass": "config/pgdb.key", + "postgres_port": 5432, "report_channel_id": 0, "report_timezones":"US/Central", "report_times": [1200, 1600, 0], diff --git a/covid_bot.py b/covid_bot.py index 735a632..c9461dc 100644 --- a/covid_bot.py +++ b/covid_bot.py @@ -1,11 +1,12 @@ -#!/bin/python +#!/bin/python3 + from discord.ext import commands from lib.covidBot import add_commands -from lib.parse_data import import_config +from lib.config_lib import import_config if(__name__ == '__main__'): - bot = commands.Bot(command_prefix='!') config_dict = import_config() + bot = commands.Bot(command_prefix='!') add_commands(bot) bot.run(config_dict['api_key']) diff --git a/covid_report.py b/covid_report.py index 9a625bf..677be26 100644 --- a/covid_report.py +++ b/covid_report.py @@ -1,8 +1,8 @@ -#!/bin/python +#!/bin/python3 from discord import Client from lib.covid_report_lib import background_task -from lib.parse_data import import_config +from lib.config_lib import import_config if(__name__ == '__main__'): diff --git a/docker-compose.yml b/docker-compose.yml index 238cc16..f86d184 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,6 +15,8 @@ services: volumes: - ./lib:/home/covidbot/lib - ./config:/home/covidbot/config + environment: + BOT_CONFIG_PATH: config/config.json covidreport: build: context: . @@ -29,6 +31,8 @@ services: volumes: - ./lib:/home/covidreport/lib - ./config:/home/covidreport/config + environment: + BOT_CONFIG_PATH: config/config.json postgres: image: postgres:alpine networks: diff --git a/lib/config_lib.py b/lib/config_lib.py new file mode 100644 index 0000000..46ac4ae --- /dev/null +++ b/lib/config_lib.py @@ -0,0 +1,50 @@ +#!/usr/bin/python3 + +from os import environ +from os.path import exists +from json import load +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + + +def import_config(): + try: + path = environ['BOT_CONFIG_PATH'] + except: + path = 'config/config.json' + if(exists(path)): + try: + with open(path) as config_file: + config_dict = load(config_file) + except Exception as e: + print(f'There was some issue opening and loading the config.\n{e}') + exit(1) + else: + print('Didn\'t find the config file.') + exit(1) + return config_dict + + +def get_engine(): + config_data = import_config() + try: + with open(config_data['postgres_pass']) as pgdb_pass: + engine = create_engine( + f"postgresql+psycopg2://{config_data['postgres_user']}:{pgdb_pass.readline().strip()}@{config_data['postgres_host']}:{config_data['postgres_port']}/{config_data['postgres_db']}") + return engine + except Exception as e: + print( + f'There was an issue opening the config file for the postgres password.\n{e}') + exit(1) + + +def create_session(Base): + try: + engine = get_engine() + Base.metadata.bind = engine + covidDataSession = sessionmaker(bind=engine) + except Exception as e: + print( + f'There was an issue creating a session for the database.\n{e}') + exit(1) + return covidDataSession() diff --git a/lib/covidBot.py b/lib/covidBot.py index 37d725c..0faa8e8 100644 --- a/lib/covidBot.py +++ b/lib/covidBot.py @@ -1,6 +1,5 @@ from discord.ext import commands -from re import match -from lib.parse_data import * +from lib.parse_data import get_covid_data, get_top_data @commands.command() @@ -25,7 +24,10 @@ async def top(ctx, arg='5'): except Exception as e: await ctx.send(f'{arg} isn\'t a number.') return - await ctx.send(get_top_data(num)) + try: + await ctx.send(get_top_data(num)) + except Exception as e: + await ctx.send(f'{e}') def add_commands(bot): diff --git a/lib/covidData.py b/lib/covidData.py index 68e2a95..9640cfd 100755 --- a/lib/covidData.py +++ b/lib/covidData.py @@ -1,7 +1,6 @@ -from sqlalchemy import Column, Integer, String, create_engine +from sqlalchemy import Column, Integer, String from sqlalchemy.ext.declarative import declarative_base -import json -from os.path import exists +from lib.config_lib import get_engine Base = declarative_base() @@ -20,22 +19,5 @@ class covidData(Base): total_cases_per_one_mil = Column(String(32)) -path = 'config/config.json' -if(exists(path)): - try: - with open(path) as config_file: - config_data = json.load(config_file) - except Exception as e: - print(f'There was some issue opening and loading the config.\n{e}') - exit(1) -else: - print('Didn\'t find the config file.') - exit(1) -try: - with open(config_data['postgres_pass']) as pgdb_pass: - engine = create_engine( - f"postgresql+psycopg2://{config_data['postgres_user']}:{pgdb_pass.readline().strip()}@{config_data['postgres_host']}:{config_data['postgres_port']}/{config_data['postgres_db']}") - Base.metadata.create_all(engine) -except Exception as e: - print( - f'There was an issue opening the config file for the postgres password.\n{e}') +engine = get_engine() +Base.metadata.create_all(engine) diff --git a/lib/covid_data_lib.py b/lib/covid_data_lib.py index 995e7de..684dd4f 100644 --- a/lib/covid_data_lib.py +++ b/lib/covid_data_lib.py @@ -1,74 +1,62 @@ -from sqlalchemy import create_engine -from sqlalchemy.orm import sessionmaker -from lib.covidData import covidData, Base +#!/usr/bin/python3 + +from lib.covidData import covidData -def init_database(config_data): - try: - with open(config_data['postgres_pass']) as pgdb_pass: - engine = create_engine( - f"postgresql+psycopg2://{config_data['postgres_user']}:{pgdb_pass.readline().strip()}@{config_data['postgres_host']}:{config_data['postgres_port']}/{config_data['postgres_db']}") - Base.metadata.bind = engine - covidDataSession = sessionmaker(bind=engine) - return covidDataSession() - except Exception as e: - print( - f'There was an issue opening the config file for the postgres password.\n{e}') - exit(1) +def set_data_dict(selection, import_data): + columns = [i for i in covidData.__dict__.keys() if i[:1] != '_'] + formatted_import_data = [d.strip().replace( + ',', '').replace('+', '') for d in import_data] + output_dict = { + 'selection': selection.upper(), + 'selection_original': selection + } + for i in range(2, len(columns)): + if(formatted_import_data[i - 1] and not formatted_import_data[i - 1] == ''): + if(i == len(columns) - 1): + output_dict[columns[i]] = formatted_import_data[i - 1] + else: + output_dict[columns[i]] = int(formatted_import_data[i - 1]) + return output_dict + + +def get_data_dict(query_data): + columns = [i for i in covidData.__dict__.keys() if i[:1] != '_'] + output_dict = { + 'selection': query_data.selection, + 'selection_original': query_data.selection_original + } + for i in range(2, len(columns)): + if(i == len(columns) - 1): + output_dict[columns[i]] = query_data.__getattribute__( + columns[i]) + else: + output_dict[columns[i]] = int( + query_data.__getattribute__(columns[i])) + return output_dict def set_data(session, selection, import_data): - new_data = covidData(selection=selection.upper(), - selection_original=selection) - for n in range(1, 8): - data = import_data[n].strip().replace( - ',', '').replace('+', '').replace('+', '') - if(data and not data == ''): - if(n == 1): - new_data.total_cases = int(data) - if(n == 2): - new_data.new_cases = int(data) - if(n == 3): - new_data.total_deaths = int(data) - if(n == 4): - new_data.new_deaths = int(data) - if(n == 5): - new_data.total_recovered = int(data) - if(n == 6): - new_data.active_cases = int(data) - if(n == 7): - new_data.serious_critical = int(data) - else: - if(n == 1): - new_data.total_cases = 0 - if(n == 2): - new_data.new_cases = 0 - if(n == 3): - new_data.total_deaths = 0 - if(n == 4): - new_data.new_deaths = 0 - if(n == 5): - new_data.total_recovered = 0 - if(n == 6): - new_data.active_cases = 0 - if(n == 7): - new_data.serious_critical = 0 - new_data.total_cases_per_one_mil = import_data[8].strip() + new_data = covidData() + formatted_data_dict = set_data_dict(selection, import_data) + try: + for key in formatted_data_dict: + new_data.__setattr__(key, formatted_data_dict[key]) + except Exception as e: + print(e) session.merge(new_data) session.commit() -def get_formatted_data(session, selection): +def get_data(session, selection): print('Formatting data.') - columns = sorted([i for i in covidData.__dict__.keys() if i[:1] != '_']) - columns = [' '.join([d.capitalize() for d in c.replace( - 'per', '/').split('_')]) for c in columns] all_data_query = session.query(covidData).filter( covidData.selection == selection).all() - return (columns, all_data_query) + return get_data_dict(all_data_query[0]) def get_top_n_rows(session, num): print(f'Getting top {num} rows.') - top_n_rows = session.query(covidData).order_by(covidData.total_cases.desc()).limit(num).all() - return top_n_rows \ No newline at end of file + top_n_rows = session.query(covidData).order_by( + covidData.total_cases.desc()).limit(num).all() + return [get_data_dict(n) for n in top_n_rows] diff --git a/lib/parse_data.py b/lib/parse_data.py index 5908578..ac1dcec 100755 --- a/lib/parse_data.py +++ b/lib/parse_data.py @@ -1,11 +1,11 @@ -#!/usr/bin/python +#!/usr/bin/python3 import requests from bs4 import BeautifulSoup -import json -from os.path import exists from inspect import getmembers, isroutine -from lib.covid_data_lib import init_database, set_data, get_formatted_data, get_top_n_rows +from lib.config_lib import create_session +from lib.covid_data_lib import set_data, get_data, get_top_n_rows +from lib.covidData import Base # San Antonio url sa_data_url = 'https://www.sanantonio.gov/health/news/alerts/coronavirus' @@ -24,26 +24,11 @@ def format_parse_int(num): return output[::-1] -def import_config(path='config/config.json'): - if(exists(path)): - try: - with open(path) as config_file: - config_dict = json.load(config_file) - except Exception as e: - print(f'There was some issue opening and loading the config.\n{e}') - exit(1) - else: - print('Didn\'t find the config file.') - exit(1) - return config_dict - - def update_data(): try: print('Creating session.') - session = init_database(import_config()) + session = create_session(Base) except Exception as e: - session.rollback() print(f'There was an error trying to create a database session:\n{e}') data_html = requests.get('https://www.worldometers.info/coronavirus/') if(data_html.status_code == '200' or data_html.status_code == 200): @@ -60,44 +45,38 @@ def update_data(): session.close() -def format_covid_data(columns, data): - output = '' - get_values = [attr for attr in getmembers(data, lambda a:not( - isroutine(a))) if not(attr[0].startswith('__') and attr[0].endswith('__')) and not attr[0].startswith('_') and not attr[0] == 'metadata'] - output += f'{columns[4]}: {get_values[4][1]}\n' - output += f'{columns[6]}: {format_parse_int(int(get_values[6][1]))}\n' - output += f'{columns[1]}: {format_parse_int(int(get_values[1][1]))}\n' - output += f'{columns[8]}: {format_parse_int(int(get_values[8][1]))}\n' - output += f'{columns[2]}: {format_parse_int(int(get_values[2][1]))}\n' - output += f'{columns[9]}: {format_parse_int(int(get_values[9][1]))}\n' - output += f'{columns[0]}: {format_parse_int(int(get_values[0][1]))}\n' - output += f'{columns[5]}: {format_parse_int(int(get_values[5][1]))}\n' - output += f'{columns[7]}: {get_values[7][1]}\n' - return output - - def get_covid_data(selection): print('Updating data.') - columns, all_data_query = get_formatted_data( - init_database(import_config()), selection) - output = '' - for data in all_data_query: - output += format_covid_data(columns, data) + session = create_session(Base) + data_dict = get_data(session, selection) + session.close() + output = f'Selection: {data_dict["selection_original"]}\n' + for key in data_dict: + temp_key = ' '.join([d.capitalize() for d in key.replace( + 'per', '/').split('_')]) + if(key == 'selection' or key == 'selection_original'): + pass + elif(key == 'total_cases_per_one_mil'): + output += f'{temp_key}: {data_dict[key]}\n' + else: + output += f'{temp_key}: {format_parse_int(data_dict[key])}\n' return output def get_top_data(number): - top_n_rows = get_top_n_rows(init_database(import_config()), number + 1) + session = create_session(Base) + top_n_rows = get_top_n_rows(session, number + 2) + session.close() output = '' count = 0 for row in top_n_rows: - if(not count == 0): - output += f'# {count}\n{row.selection_original}: {format_parse_int(int(row.total_cases))}' - if(not count == number): + if(count > 1): + output += f'# {count - 1}\n{row["selection_original"]}: {format_parse_int(row["total_cases"])}' + if(not count == number + 1): output += '\n' count += 1 return output if(__name__ == '__main__'): - print(get_covid_data()) + print(get_covid_data('TOTAL'))