error handler

Author: brucewzj99

bot/common.py

@@ -34,4 +34,5 @@ class ConversationState(Enum):
         CPF,
         BACKLOG,
         ADD_BACKLOG_ENTRY,
-    ) = range(26)
+        NOTIFICATION,
+    ) = range(27)

bot/database_service/firestore_service.py

@@ -1,6 +1,7 @@
 from bot.database_service.auth import get_db_client
 from datetime import datetime, timedelta
 import pytz
+from bot.error_handler import DatabaseError
 
 
 class FirestoreService:
@@ -14,25 +15,37 @@ def __init__(self, collection_name="users"):
 
     # New user setup
     def new_user_setup(self, telegram_id, sheet_id, telegram_username):
-        user_ref = self.db.collection(self.collection_name).document(str(telegram_id))
-        timestamp = datetime.now(pytz.timezone("Asia/Singapore"))
-        user_ref.set(
-            {
-                "sheet_id": sheet_id,
-                "datetime_created": timestamp,
-                "username": telegram_username,
-                "usage_count": 0,
-                "last_accessed": timestamp,
-                "hourly_accessed": timestamp,
-                "overusage_count": 0,
-            }
-        )
+        try:
+            user_ref = self.db.collection(self.collection_name).document(
+                str(telegram_id)
+            )
+            timestamp = datetime.now(pytz.timezone("Asia/Singapore"))
+            user_ref.set(
+                {
+                    "sheet_id": sheet_id,
+                    "datetime_created": timestamp,
+                    "username": telegram_username,
+                    "usage_count": 0,
+                    "last_accessed": timestamp,
+                    "hourly_accessed": timestamp,
+                    "overusage_count": 0,
+                }
+            )
+        except Exception as e:
+            raise DatabaseError(message="Error setting up new user", extra_info=str(e))
 
     # Check if user exists
     def check_if_user_exists(self, telegram_id):
-        user_ref = self.db.collection(self.collection_name).document(str(telegram_id))
-        user_doc = user_ref.get()
-        return user_doc.exists
+        try:
+            user_ref = self.db.collection(self.collection_name).document(
+                str(telegram_id)
+            )
+            user_doc = user_ref.get()
+            return user_doc.exists
+        except Exception as e:
+            raise DatabaseError(
+                message="Error checking if user exists", extra_info=str(e)
+            )
 
     # Get user sheet id
     def get_user_sheet_id(self, telegram_id, telegram_username):
@@ -76,19 +89,18 @@ def get_user_sheet_id(self, telegram_id, telegram_username):
                 )
                 return user_doc.get("sheet_id")
             except Exception as e:
-                raise e
-        return None
+                raise DatabaseError(
+                    message="Error getting user sheet id", extra_info=str(e)
+                )
+        raise DatabaseError(
+                    message="Error getting user sheet id", extra_info="User does not exist"
+                )
 
     # Get all user IDs
     def get_all_user_id(self):
-        users_ref = self.db.collection(self.collection_name)
-        user_ids = [int(user.id) for user in users_ref.stream()]
-        return user_ids
-
-    # Get all sheet IDs
-    def get_all_sheet_id(self):
-        users_ref = self.db.collection(self.collection_name)
-        sheet_ids = []
-        for user in users_ref.stream():
-            sheet_ids.append(user.get("sheet_id"))
-        return sheet_ids
+        try:
+            users_ref = self.db.collection(self.collection_name)
+            user_ids = [int(user.id) for user in users_ref.stream()]
+            return user_ids
+        except Exception as e:
+            raise DatabaseError(message="Error getting all user ids", extra_info=str(e))

bot/error_handler.py

@@ -0,0 +1,28 @@
+import logging
+
+
+logging.basicConfig(level=logging.ERROR, filename='telegram_bot.log',
+                    format='%(asctime)s - %(levelname)s - %(message)s')
+
+class BaseError(Exception):
+    """Base class for other exceptions"""
+    def __init__(self, error_class, message="An error occurred", extra_info=""):
+        self.message = message
+        self.extra_info = extra_info
+        self.error_class = error_class
+        super().__init__(f"{self.error_class}: {self.message}")
+
+class GoogleSheetError(BaseError):
+    """Exception raised for errors in the Google Sheet service."""
+    def __init__(self, message="", extra_info=""):
+        super().__init__("GoogleSheetError", message, extra_info)
+
+class TelegramBotError(BaseError):
+    """Exception raised for errors in the Telegram bot operations."""
+    def __init__(self, message="", extra_info=""):
+        super().__init__("TelegramBotError", message, extra_info)
+
+class DatabaseError(BaseError):
+    """Exception raised for errors in database operations."""
+    def __init__(self, message="", extra_info=""):
+        super().__init__("DatabaseError", message, extra_info)