import logging import os import traceback import urllib.parse import pandas as pd import psycopg2 from dotenv import load_dotenv from sqlalchemy import create_engine class Database(object): def __init__(self): # generate uri engine = os.getenv('DATABASE_ENGINE') username = os.getenv('DATABASE_USERNAME') password = os.getenv('DATABASE_PASSWORD') password = urllib.parse.quote_plus(password) # escape special characters host = os.getenv('DATABASE_HOST') port = os.getenv('DATABASE_PORT') database = os.getenv('DATABASE_DATABASE') self._uri = f'{engine}://{username}:{password}@{host}:{port}/{database}' # prepare connection to database self._engine = create_engine(self._uri) self.con = None def __enter__(self): self.connect() return self def __exit__(self, exc_type, exc_value, tb): if not exc_type is None: traceback.print_exception(exc_type, exc_value, tb) self.disconnect() def connect(self): try: con = self._engine.connect() except psycopg2.OperationalError as e: logging.warning(e) else: self.con = con def disconnect(self): self.con.close() def execute(self, query, var=None): assert self.con is not None, 'connection to database could not be established' try: if var is None: self.con.execute(query) else: self.con.execute(query, var) except Exception as e: print(e) raise def get(self, query): df = pd.read_sql_query(query, self.con) return df class Schema(Database): def __init__(self, name): super(Schema, self).__init__() self._name = name def __enter__(self): self.connect() if not self.exists(): self.create() return self def exists(self): query = ( 'SELECT schema_name ' 'FROM information_schema.schemata ' 'ORDER BY schema_name ' ';' ) df = self.get(query) schema_list = df['schema_name'].tolist() return self._name in schema_list def create(self): query = f'CREATE SCHEMA {self._name};' self.execute(query) def purge(self): query = f'DROP SCHEMA IF EXISTS {self._name} CASCADE;' self.execute(query) self.create() class Table(Database): def __init__(self, name): super(Table, self).__init__() self._name = name def __enter__(self): self.connect() if not self.exists(): self.create() return self def exists(self): query = ( 'SELECT table_name ' 'FROM information_schema.tables ' "WHERE table_schema='public' " 'ORDER BY table_name ' ) df = self.get(query) table_list = df['table_name'].tolist() return self._name in table_list def create(self): # placeholder for childrens' methods pass def purge(self): query = f'DROP TABLE IF EXISTS {self._name} CASCADE;' self.execute(query) self.create() class Data_table(Table): def __init__(self): super(Data_table, self).__init__('data_table') def create(self): query = ( f'CREATE TABLE {self._name} ' '(' 'id SERIAL PRIMARY KEY, ' 'username VARCHAR(150) NOT NULL, ' 'time TIMESTAMPTZ NOT NULL, ' 'content TEXT NOT NULL, ' 'UNIQUE (username, time, content) ' ')' ';' ) self.execute(query) def insert(self, val_dict): query = ( f'INSERT INTO {self._name} (username, time, content) ' 'VALUES (%(username)s, %(time)s, %(content)s) ' 'ON CONFLICT (username, time, content) DO NOTHING ' ';' ) self.execute(query, val_dict) def insert(username, timestamp, content): val_dict = { 'username': username, 'time': timestamp, 'content': content } with Data_table() as dt: dt.insert(val_dict) def insert_list(username_list, timestamp_list, content_list): with Data_table() as dt: for username, timestamp, content in zip(username_list, timestamp_list, content_list): try: val_dict = { 'username': username, 'time': timestamp, 'content': content } dt.insert(val_dict) except Exception as e: logging.warning(e) def insert_translated_list(content_list, id_list): with Data_table() as dt: for id, content in zip(id_list, content_list): try: val_dict = { 'id': id, 'content_en': content } query = ( f'UPDATE {dt._name} ' ' SET content_en = (%(content_en)s) ' 'WHERE id = (%(id)s) ' ) dt.execute(query, val_dict) except Exception as e: print(f'error when updating database: {e}') def get_content(): with Data_table() as dt: query = f'SELECT * FROM {dt._name} ORDER BY time DESC ' df = pd.read_sql_query(query, dt.con) return df['content'] def get_all(): with Data_table() as dt: query = f'SELECT * FROM {dt._name} ORDER BY time DESC ' df = pd.read_sql_query(query, dt.con) return df def insert_translated(dt, val_dict): assert 'id' in val_dict.keys() assert 'content_en' in val_dict.keys() query = ( f'UPDATE {dt._name} ' ' SET content_en = (%(content_en)s) ' 'WHERE id = (%(id)s) ' ) dt.execute(query, val_dict) def get_untranslated(dt): query = ( f'SELECT * FROM {dt._name} ' 'WHERE content_en IS NULL ' 'ORDER BY random() ' 'LIMIT 1 ' ) df = pd.read_sql_query(query, dt.con) if len(df)==0: return None val_dict = { 'id': int(df['id'][0]), 'content': str(df['content'][0]) } return val_dict