Make DB & telegram optional

1 files changed, 36 insertions, 26 deletions

diff --git a/src/db_handler.py b/src/db_handler.py
index a7740f7..cd40112 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -1,31 +1,41 @@
 import os
+import logging
+from dataclasses import dataclass
 
 import psycopg
 import yaml
 
-with open(os.path.join('configs', 'database.yml'), 'r') as file:
-    db_con_params = yaml.safe_load(file.read())
-
-
-def get_comments(url: str) -> list[dict]:
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute(
-                "SELECT id, comment, page_url, visitor_url, nickname, "
-                "(CASE WHEN show_contact THEN contact ELSE NULL END) as contact, created_at "
-                "FROM comment WHERE approved AND public AND page_url = %(url)s ORDER BY created_at DESC;", {'url': url})
-            comments = cur.fetchall()
-            return comments
-
-
-def post_comment(comment: str, page_url: str, visitor_url: str=None, nickname: str=None, contact: str=None,
-                 show_contact: bool=True, public: bool=True):
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute(
-                "INSERT INTO comment (comment, page_url, visitor_url, nickname, contact, show_contact, public)"
-                "VALUES (%(comment)s, %(page_url)s, %(visitor_url)s, %(nickname)s, %(contact)s, %(show_contact)s, "
-                "%(public)s);",
-                {'comment': comment, 'page_url': page_url, 'visitor_url': visitor_url, 'nickname': nickname,
-                 'contact': contact, 'show_contact': show_contact, 'public': public}
-            )
+
+logger = logging.getLogger(__name__) # Set the logger name, to the name of the module
+
+
+@dataclass
+class DBHandler:
+    config: dict
+
+    def conn_factory(self) -> psycopg.Connection:
+        conn: psycopg.Connection = psycopg.connect(**self.config, row_factory=psycopg.rows.dict_row)
+        return conn
+
+    def get_comments(self, url: str) -> list[dict]:
+        with self.conn_factory() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "SELECT id, comment, page_url, visitor_url, nickname, "
+                    "(CASE WHEN show_contact THEN contact ELSE NULL END) as contact, created_at "
+                    "FROM comment WHERE approved AND public AND page_url = %(url)s ORDER BY created_at DESC;", {'url': url})
+                comments = cur.fetchall()
+                return comments
+
+
+    def post_comment(self, comment: str, page_url: str, visitor_url: str=None, nickname: str=None, contact: str=None,
+                     show_contact: bool=True, public: bool=True):
+        with self.conn_factory() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "INSERT INTO comment (comment, page_url, visitor_url, nickname, contact, show_contact, public)"
+                    "VALUES (%(comment)s, %(page_url)s, %(visitor_url)s, %(nickname)s, %(contact)s, %(show_contact)s, "
+                    "%(public)s);",
+                    {'comment': comment, 'page_url': page_url, 'visitor_url': visitor_url, 'nickname': nickname,
+                     'contact': contact, 'show_contact': show_contact, 'public': public}
+                )