Coverage for database.py: 100%

1import sys 

2 

3import psycopg 

4from psycopg import Error 

5 

6 

7class Database: 

8 def __init__(self, conn_string, db_name): 

9 self.conn_string = conn_string 

10 self.db_name = db_name 

11 

12 def get_db_connection(self): 

13 try: 

14 conn = psycopg.connect(self.conn_string, autocommit=True) 

15 return conn 

16 except psycopg.Error as e: 

17 print( 

18 f"Error {e} on connection string {self.conn_string}", file=sys.stderr) 

19 sys.exit(1) 

20 

21 def execute_query(self, query, fetch=True): 

22 conn = self.get_db_connection() 

23 if conn is None: 

24 return None 

25 

26 try: 

27 with conn.cursor() as cur: 

28 cur.execute(query) 

29 if fetch: 

30 results = cur.fetchall() 

31 return results 

32 except Error as e: 

33 print( 

34 f"Error {e} with query '{query}' on host '{self.conn_string}'", file=sys.stderr) 

35 return None 

36 finally: 

37 conn.close() 

38 

39 def execute_query_rollback_on_error(self, queries: str): 

40 with psycopg.connect(self.conn_string) as conn: 

41 with conn.cursor() as cur: 

42 try: 

43 cur.execute(queries) 

44 conn.commit() 

45 

46 except Exception as e: 

47 print(f"An error occurred: {e}", file=sys.stderr) 

48 conn.rollback()