moved files to code folder

2022-12-17 01:02:13 +01:00
parent fc5ebd1dde
commit bc73952e76
5 changed files with 0 additions and 0 deletions
@@ -0,0 +1,159 @@
+#!/Users/brian/Projects/twitter_scraper/venv/bin/python
+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(15) 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)