import os import sys import psycopg2 from psycopg2 import sql from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT import uuid from datetime import datetime import pytz import pdb tz = pytz.timezone('Europe/Madrid') class Database(): name = 'followers database library' def __init__(self, config_file=None, followers_db=None, followers_db_user=None, followers_db_user_password=None): self.config_file = "config/db_config.txt" self.followers_db = self.__get_parameter("followers_db", self.config_file) self.followers_db_user = self.__get_parameter("followers_db_user", self.config_file) self.followers_db_user_password = self.__get_parameter("followers_db_user_password", self.config_file) db_setup = self.__check_dbsetup(self) if not db_setup: self.followers_db = input("\nFollowers database name: ") self.followers_db_user = input("\nFollowers database user: ") self.followers_db_user_password = input("\nFollowers database user password: ") self.__createdb(self) self.__create_config(self) self.__write_config(self) def find_follower(self, follower, account): now = datetime.now() found = False select_sql = "select from followers where follower=(%s) and account=(%s)" conn = None try: conn = psycopg2.connect(database=self.followers_db, user=self.followers_db_user, password = self.followers_db_user_password, host="/var/run/postgresql", port="6432") cur = conn.cursor() print(f'Searching follower {follower}...') cur.execute(select_sql, (follower, account)) row = cur.fetchone() if row != None: print(f'follower {follower} is already in the database.') found = True else: print(f'follower {follower} not found! Adding it...') self.save_follower(follower, account) cur.close() except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() def save_follower(self, follower, account): now = datetime.now() insert_sql = "INSERT INTO followers(follower, account, updated_at) VALUES(%s,%s,%s) ON CONFLICT DO NOTHING" conn = None try: conn = psycopg2.connect(database=self.followers_db, user=self.followers_db_user, password = self.followers_db_user_password, host="/var/run/postgresql", port="6432") cur = conn.cursor() print(f'Writing follower {follower}...') cur.execute(insert_sql, (follower, account, now)) conn.commit() cur.close() except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() def last_total(self): now = datetime.now() last_local = 0 last_remote = 0 select_sql = "select local, remote from total order by updated_at desc limit 1" conn = None try: conn = psycopg2.connect(database=self.followers_db, user=self.followers_db_user, password = self.followers_db_user_password, host="/var/run/postgresql", port="6432") cur = conn.cursor() cur.execute(select_sql) row = cur.fetchone() if row != None: last_local = row[0] last_remote = row[1] cur.close() except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() return last_local, last_remote def save_total(self, local, local_incr, remote, remote_incr): now = datetime.now() insert_sql = "INSERT INTO total(local, local_incr, remote, remote_incr, updated_at) VALUES(%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING" conn = None try: conn = psycopg2.connect(database=self.followers_db, user=self.followers_db_user, password = self.followers_db_user_password, host="/var/run/postgresql", port="6432") cur = conn.cursor() print(f'Writing totals {local} & {remote}...') cur.execute(insert_sql, (local, local_incr, remote, remote_incr, now)) conn.commit() cur.close() except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() def totals(self): local_followers = 0 remote_followers = 0 local_sql = "select count(*) from followers where account not like '%@%'" remote_sql = "select count(*) from followers where account like '%@%'" conn = None try: conn = psycopg2.connect(database=self.followers_db, user=self.followers_db_user, password = self.followers_db_user_password, host="/var/run/postgresql", port="6432") cur = conn.cursor() cur.execute(local_sql) row = cur.fetchone() if row != None: local_followers = row[0] cur.execute(remote_sql) row = cur.fetchone() if row != None: remote_followers = row[0] cur.close() except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() return local_followers, remote_followers @staticmethod def __check_dbsetup(self): db_setup = False try: conn = None conn = psycopg2.connect(database = self.followers_db, user = self.followers_db_user, password = self.followers_db_user_password, host = "/var/run/postgresql", port = "6432") db_setup = True except (Exception, psycopg2.DatabaseError) as error: print(error) return db_setup @staticmethod def __createdb(self): conn = None try: conn = psycopg2.connect(dbname='postgres', user=self.followers_db_user, host='', password=self.followers_db_user_password) conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) cur = conn.cursor() print(f"Creating database {self.followers_db}. Please wait...") cur.execute(sql.SQL("CREATE DATABASE {}").format( sql.Identifier(self.followers_db)) ) print(f"Database {self.followers_db} created!\n") self.__dbtables_schemes(self) except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() @staticmethod def __dbtables_schemes(self): db = self.followers_db table = "followers" sql = "create table "+table+" (follower varchar(50) PRIMARY KEY, account varchar(100), updated_at timestamptz)" self.__create_table(self, table, sql) table = "total" sql = "create table "+table+" (local int, local_incr int, remote int, remote_incr int, updated_at timestamptz)" self.__create_table(self, table, sql) @staticmethod def __create_table(self, table, sql): conn = None try: conn = psycopg2.connect(database = self.followers_db, user = self.followers_db_user, password = self.followers_db_user_password, host = "/var/run/postgresql", port = "6432") cur = conn.cursor() print(f"Creating table {table}") cur.execute(sql) conn.commit() print(f"Table {table} created!\n") except (Exception, psycopg2.DatabaseError) as error: print(error) finally: if conn is not None: conn.close() def __get_parameter(self, parameter, config_file): if not os.path.isfile(config_file): print(f"File {config_file} not found..") return with open( config_file ) as f: for line in f: if line.startswith( parameter ): return line.replace(parameter + ":", "").strip() print(f"{config_file} Missing parameter {parameter}") sys.exit(0) @staticmethod def __create_config(self): if not os.path.exists('config'): os.makedirs('config') if not os.path.exists(self.config_file): print(self.config_file + " created!") with open(self.config_file, 'w'): pass @staticmethod def __write_config(self): with open(self.config_file, 'a') as the_file: the_file.write(f'followers_db: {self.followers_db}\nfollowers_db_user: {self.followers_db_user}\nfollowers_db_user_password: {self.followers_db_user_password}') print(f"adding parameters to {self.config_file}\n")