Migrate to firestore, reformat getoverall

Author: brucewzj99

bot/firestore.py

@@ -0,0 +1,40 @@
+from bot.firestore_config import db
+
+
+# New user setup
+def new_user_setup(telegram_id, sheet_id):
+    user_ref = db.collection("users").document(str(telegram_id))
+    user_ref.set({"sheet_id": sheet_id})
+
+
+# Check if user exists
+def check_if_user_exists(telegram_id):
+    user_ref = db.collection("users").document(str(telegram_id))
+    user_doc = user_ref.get()
+    return user_doc.exists
+
+
+# Get user sheet id
+def get_user_sheet_id(telegram_id):
+    user_ref = db.collection("users").document(str(telegram_id))
+    user_doc = user_ref.get()
+    if user_doc.exists:
+        return user_doc.get("sheet_id")
+    else:
+        return None
+
+
+# Get all user IDs
+def get_all_user_id():
+    users_ref = db.collection("users")
+    user_ids = [int(user.id) for user in users_ref.stream()]
+    return user_ids
+
+
+# Get all sheet IDs
+def get_all_sheet_id():
+    users_ref = db.collection("users")
+    sheet_ids = []
+    for user in users_ref.stream():
+        sheet_ids.append(user.get("sheet_id"))
+    return sheet_ids

bot/firestore_config.py

@@ -0,0 +1,12 @@
+from firebase_admin import firestore
+import firebase_admin
+from firebase_admin import credentials
+import os
+import json
+
+firebase_json = json.loads(os.environ["FIREBASE_JSON"])
+
+cred = credentials.Certificate(firebase_json)
+firebase_admin.initialize_app(cred)
+
+db = firestore.client()

bot/telegram_bot.py

@@ -16,7 +16,7 @@
 from bot.common import EntryType
 from bot.common import ConversationState as CS
 import bot.google_sheet as gs
-import bot.firebase as db
+import bot.firestore as db
 
 timezone = pytz.timezone("Asia/Singapore")
 
@@ -37,6 +37,7 @@ def check_date_format(date_string):
     pattern = r"\b\d{1,2}\s(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b"
     return bool(re.fullmatch(pattern, date_string, re.IGNORECASE))
 
+
 def check_month_format(month_string):
     pattern = r"^(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)$"
     return bool(re.fullmatch(pattern, month_string, re.IGNORECASE))
@@ -623,6 +624,7 @@ def get_day_transaction(update, context):
         update.message.reply_text(ERROR_TEXT)
         return ConversationHandler.END
 
+
 def get_overall(update, context):
     context.user_data.clear()
     telegram_id = update.effective_user.id
@@ -634,6 +636,7 @@ def get_overall(update, context):
         update.message.reply_text(ERROR_TEXT)
         return ConversationHandler.END
 
+
 def handle_get_transaction(update, context):
     sheet_id = context.user_data["sheet_id"]
     reply = update.message.text
@@ -662,28 +665,42 @@ def handle_get_transaction(update, context):
     except Exception as e:
         update.message.reply_text(ERROR_TEXT)
 
+
 def handle_get_overall(update, context):
     sheet_id = context.user_data["sheet_id"]
     month = update.message.text
-    msg = ""
     try:
         if check_month_format(month):
             values = gs.get_overall(sheet_id, month)
-            msg  = "```\n"
-            for row in values:
-                if len(row) == 3:
-                    msg += "{:<20} {:<10} {:<10}\n".format(*row)
-                elif len(row) == 2:
-                    msg += "{:<20} {:<10}\n".format(*row)
-            msg += "```"
-            update.message.reply_text(msg, parse_mode='Markdown')
+
+            final_income = values[0]
+            msg = f"--- Final Income ---\n`{final_income[2]}`\n\n"
+
+            msg += "--- Spending ---\n"
+            max_len = (
+                max(len(row[0]) for row in values[1:-2]) + 1
+            )  # +1 for some extra space
+            for row in values[1:-2]:
+                if row[1].startswith("-"):  # if value is negative
+                    msg += f"`{row[0].ljust(max_len)}{row[1]}`\n"
+                else:
+                    msg += f"`{row[0].ljust(max_len)} {row[1]}`\n"  # else keep the original space
+
+            total_spent = values[-2]
+            msg += f"\n--- Total Spent ---\n`{total_spent[1]}`\n"
+
+            overall = values[-1]
+            msg += f"\n--- Overall ---\n`{overall[2]}`\n"
+
+            update.message.reply_text(msg, parse_mode="Markdown")
             return ConversationHandler.END
         else:
             update.message.reply_text(GET_OVERALL_TEXT)
             return CS.HANDLE_GET_OVERALL
     except Exception as e:
         update.message.reply_text(ERROR_TEXT)
 
+
 def add_income(update, context):
     context.user_data.clear()
     telegram_id = update.effective_user.id
@@ -807,7 +824,7 @@ def setup_handlers(dispatcher):
             CommandHandler("addothers", add_others),
             CommandHandler("addincome", add_income),
             CommandHandler("getdaytransaction", get_day_transaction),
-            CommandHandler("getoverall", get_overall)
+            CommandHandler("getoverall", get_overall),
         ],
         states={
             CS.SET_UP: [MessageHandler(Filters.text & ~Filters.command, set_up)],