twitter_scraper/code/database.py

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
from pymongo import MongoClient


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

def connect_mongo(ip_addr: str = '192.168.1.2', db_name: str = 'twitter', collection_name: str = 'tweet'):
    client = MongoClient(ip_addr)
    db = client[db_name]
    coll = db[collection_name]
    return coll