ADD: Custom CallBack Context and Storage

custom_callback_contex.py

@@ -0,0 +1,11 @@
+from telegram.ext import CallbackContext as BaseCallbackContext
+
+from storage.dynamodb_storage import DynamoDBStorage
+
+
+class CustomCallbackContext(BaseCallbackContext):
+    def __init__(self, application, chat_id=None, user_id=None):
+        super().__init__(application=application, chat_id=chat_id, user_id=user_id)
+
+    def storage(self) -> DynamoDBStorage:
+        return self.bot_data["storage"]

deploy_lambda.sh

@@ -23,7 +23,7 @@ then
 fi
 
 rm -f $ZIP_FILE
-zip -r $ZIP_FILE *.py html/
+zip -r $ZIP_FILE *.py html/ storage/
 
 # Append the contents of the 'dependencies' directory at the root of the zip file
 if [ -d "dependencies" ]; then

lambda_main.py

@@ -7,13 +7,15 @@
 import json
 import traceback
 from telegram import Update
+from storage.dynamodb_storage import DynamoDBStorage
 
 if logging.getLogger().hasHandlers():
     logging.getLogger().setLevel(logging.INFO)
 else:
     logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 
+storage = DynamoDBStorage()
 
 def lambda_handler(event, context):
     logger.info(f"Event body type: {type(event)}")
@@ -36,7 +38,7 @@ def handle_spotify_event(event):
 
     if not state_encoded or not code:
         return {"statusCode": 400, "body": "Missing required parameters"}
-    handle_spotify_auth(state_encoded, code)
+    handle_spotify_auth(state_encoded, code, storage)
     html_content = load_html_file("index.html")
     return {
         "statusCode": 200,
@@ -47,7 +49,8 @@ def handle_spotify_event(event):
 
 async def main(event, context):
     TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
-    application = build_application(TOKEN)
+
+    application = build_application(TOKEN, storage)
 
     # Convert the incoming event to a Telegram Update object
     if isinstance(event["body"], str):

polling_main.py

@@ -5,6 +5,8 @@
 from bot import handle_spotify_auth
 import logging
 from flask import Response
+from storage.dynamodb_storage import DynamoDBStorage
+
 
 webserver = Flask(__name__)
 
@@ -24,7 +26,7 @@ def callback():
     code = request.args.get("code")
     state = request.args.get("state")
     logger.info(f"honey we got the code: {code}, and the state: {state}")
-    handle_spotify_auth(state, code)
+    handle_spotify_auth(state, code, storage)
     return Response(html_content, status=200, content_type="text/html")
 
 
@@ -34,17 +36,14 @@ def run_flask_app():
 
 if __name__ == "__main__":
     TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
-    application = build_application(TOKEN)
+
+    # Initialize storage
+    storage = DynamoDBStorage()
+
+    application = build_application(TOKEN, storage)
 
     # Start the Flask app in a separate thread
     threading.Thread(target=run_flask_app, daemon=True).start()
 
     # Start polling
     application.run_polling()
-    # Main thread does something totally different.
-    # i = 0
-    # while True:
-    #     time.sleep(1)
-    #     i += 1
-    #     sys.stdout.write(f"{i} {SPINNER[i % len(SPINNER)]}\r")
-    #     sys.stdout.flush()

requirements.txt

@@ -10,13 +10,13 @@ docutils==0.16
 exceptiongroup==1.2.0
 h11==0.14.0
 httpcore==1.0.2
-httpx==0.25.2
+httpx~=0.26.0
 idna==3.6
 jmespath==1.0.1
 pyasn1==0.5.1
 python-dateutil==2.8.2
 python-dotenv==1.0.0
-python-telegram-bot==20.7
+python-telegram-bot>=20.8
 PyYAML==6.0.1
 redis==5.0.1
 requests==2.31.0

storage/dynamodb_storage.py

@@ -0,0 +1,158 @@
+import boto3
+import logging
+from .storage_interface import BotState, Storage
+import os
+
+
+class DynamoDBStorage(Storage):
+    def __init__(self):
+        dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
+        self.bot_table = dynamodb.Table(os.getenv("BOT_TABLE"))
+        self.credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE"))
+
+    def save_current_state(self, chat_id, state_key: BotState):
+        try:
+            user_id = self.get_user_id_from_chat_id(chat_id)
+            if user_id is None:
+                user_id = self.get_user_id_from_channel_credentials(chat_id)
+
+            if state_key is BotState.NO_STATE:
+                response = self.bot_table.update_item(
+                    Key={"chat_id": str(chat_id)},
+                    UpdateExpression="REMOVE current_state",
+                    ReturnValues="UPDATED_NEW",
+                )
+            else:
+                response = self.bot_table.update_item(
+                    Key={"chat_id": str(chat_id)},
+                    UpdateExpression="SET current_state = :state, user_id = :uid",
+                    ExpressionAttributeValues={
+                        ":state": state_key.value,
+                        ":uid": user_id,
+                    },
+                    ReturnValues="UPDATED_NEW",
+                )
+            logging.info(f"Updated state in DynamoDB: {response}")
+            return response
+        except Exception as e:
+            logging.error(f"Error saving current state to DynamoDB: {e}")
+
+    def get_current_state(self, chat_id):
+        try:
+            response = self.bot_table.get_item(Key={"chat_id": str(chat_id)})
+            if "Item" in response and "current_state" in response["Item"]:
+                state_value = response["Item"]["current_state"]
+                return BotState(state_value)
+            return None
+        except Exception as e:
+            logging.error(f"Error retrieving current state from DynamoDB: {e}")
+            return None
+
+    def save_playlist_to_dynamodb(self, chat_id, playlist_id):
+        try:
+            self.bot_table.put_item(
+                Item={
+                    "chat_id": str(chat_id),
+                    "playlist_id": playlist_id,
+                    "user_id": self.get_user_id_from_chat_id(chat_id),
+                }
+            )
+        except Exception as e:
+            logging.error(f"Error saving to DynamoDB: {e}")
+
+    def get_playlist_from_dynamodb(self, chat_id):
+        try:
+            response = self.bot_table.get_item(Key={"chat_id": str(chat_id)})
+            if "Item" in response:
+                playlist_id = response["Item"].get("playlist_id")
+                return playlist_id
+            return None
+        except Exception as e:
+            logging.error(f"Error retrieving from DynamoDB: {e}")
+            return None
+
+    def get_user_id_from_chat_id(self, chat_id):
+        try:
+            response = self.bot_table.get_item(Key={"chat_id": str(chat_id)})
+            if "Item" in response and "user_id" in response["Item"]:
+                return response["Item"]["user_id"]
+            else:
+                logging.info(
+                    f"get_user_id_from_chat_id: No user_id found for chat_id: {chat_id}: {response}"
+                )
+                return None
+        except Exception as e:
+            logging.error(
+                f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}"
+            )
+            return None
+
+    def get_user_id_from_channel_credentials(self, chat_id):
+        try:
+            response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)})
+            if "Item" in response and "user_id" in response["Item"]:
+                return response["Item"]["user_id"]
+            else:
+                logging.info(
+                    f"get_user_id_from_channel_credentials: No user_id found for chat_id: {chat_id}: {response}"
+                )
+                return None
+        except Exception as e:
+            logging.error(
+                f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}"
+            )
+            return None
+
+    def get_cached_token(self, chat_id):
+        try:
+            response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)})
+            if "Item" in response:
+                return response["Item"]
+            return None
+        except Exception as e:
+            logging.error(f"Error retrieving cached token from DynamoDB: {e}")
+            raise
+
+    def save_token_to_cache(self, chat_id, user_id, token_info):
+        try:
+            self.credentials_table.put_item(
+                Item={
+                    "chat_id": str(chat_id),
+                    "user_id": user_id,
+                    **token_info,
+                }
+            )
+            self.bot_table.update_item(
+                Key={"chat_id": str(chat_id)},
+                UpdateExpression="SET user_id = :uid",
+                ExpressionAttributeValues={":uid": str(user_id)},
+            )
+        except Exception as e:
+            logging.error(f"Error saving token to DynamoDB: {e}")
+            raise
+
+    def delete_item(self, table, chat_id):
+        try:
+            if table == "credentials_table":
+                response = self.credentials_table.delete_item(Key={"chat_id": str(chat_id)})
+                logging.info(f"Deleted item from credentials_table for chat_id {chat_id}: {response}")
+            elif table == "bot_table":
+                response = self.bot_table.delete_item(Key={"chat_id": str(chat_id)})
+                logging.info(f"Deleted item from bot_table for chat_id {chat_id}: {response}")
+        except Exception as e:
+            logging.error(f"Error deleting item from DynamoDB: {e}")
+
+    def check_item_exists(self, table, chat_id):
+        try:
+            if table == "credentials_table":
+                response = self.credentials_table.get_item(Key={"chat_id": str(chat_id)})
+                logging.info(f"Checked item in credentials_table for chat_id {chat_id}: {response}")
+                return "Item" in response
+            elif table == "bot_table":
+                response = self.bot_table.get_item(Key={"chat_id": str(chat_id)})
+                logging.info(f"Checked item in bot_table for chat_id {chat_id}: {response}")
+                return "Item" in response
+            return False
+        except Exception as e:
+            logging.error(f"Error checking item in DynamoDB: {e}")
+            return False

storage/in_memory_storage.py

@@ -0,0 +1,116 @@
+import logging
+from .storage_interface import BotState, Storage
+
+
+class InMemoryStorage(Storage):
+    def __init__(self):
+        self.bot_table = {}
+        self.credentials_table = {}
+
+    def save_current_state(self, chat_id, state_key: BotState):
+        try:
+            user_id = self.get_user_id_from_chat_id(chat_id)
+            if user_id is None:
+                user_id = self.get_user_id_from_channel_credentials(chat_id)
+
+            if state_key is BotState.NO_STATE:
+                if str(chat_id) in self.bot_table:
+                    self.bot_table[str(chat_id)].pop("current_state", None)
+            else:
+                if str(chat_id) not in self.bot_table:
+                    self.bot_table[str(chat_id)] = {}
+                self.bot_table[str(chat_id)]["current_state"] = state_key.value
+                self.bot_table[str(chat_id)]["user_id"] = user_id
+
+            logging.info(
+                f"Updated state in InMemoryStorage: {self.bot_table[str(chat_id)]}"
+            )
+        except Exception as e:
+            logging.error(f"Error saving current state to InMemoryStorage: {e}")
+
+    def get_current_state(self, chat_id):
+        try:
+            item = self.bot_table.get(str(chat_id), {})
+            if "current_state" in item:
+                state_value = item["current_state"]
+                return BotState(state_value)
+            return None
+        except Exception as e:
+            logging.error(f"Error retrieving current state from InMemoryStorage: {e}")
+            return None
+
+    def save_playlist_to_dynamodb(self, chat_id, playlist_id):
+        try:
+            if str(chat_id) not in self.bot_table:
+                self.bot_table[str(chat_id)] = {}
+            self.bot_table[str(chat_id)]["playlist_id"] = playlist_id
+            self.bot_table[str(chat_id)]["user_id"] = self.get_user_id_from_chat_id(
+                chat_id
+            )
+        except Exception as e:
+            logging.error(f"Error saving to InMemoryStorage: {e}")
+
+    def get_playlist_from_dynamodb(self, chat_id):
+        try:
+            item = self.bot_table.get(str(chat_id), {})
+            return item.get("playlist_id")
+        except Exception as e:
+            logging.error(f"Error retrieving from InMemoryStorage: {e}")
+            return None
+
+    def get_user_id_from_chat_id(self, chat_id):
+        try:
+            item = self.bot_table.get(str(chat_id), {})
+            return item.get("user_id")
+        except Exception as e:
+            logging.error(
+                f"Error retrieving user ID from InMemoryStorage for chat_id: {chat_id}, error: {e}"
+            )
+            return None
+
+    def get_user_id_from_channel_credentials(self, chat_id):
+        try:
+            item = self.credentials_table.get(str(chat_id), {})
+            return item.get("user_id")
+        except Exception as e:
+            logging.error(
+                f"Error retrieving user ID from InMemoryStorage for chat_id: {chat_id}, error: {e}"
+            )
+            return None
+
+    def get_cached_token(self, chat_id):
+        try:
+            return self.credentials_table.get(str(chat_id))
+        except Exception as e:
+            logging.error(f"Error retrieving cached token from InMemoryStorage: {e}")
+            raise
+
+    def save_token_to_cache(self, chat_id, user_id, token_info):
+        try:
+            self.credentials_table[str(chat_id)] = {
+                "chat_id": str(chat_id),
+                "user_id": user_id,
+                **token_info,
+            }
+            if str(chat_id) not in self.bot_table:
+                self.bot_table[str(chat_id)] = {}
+            self.bot_table[str(chat_id)]["user_id"] = str(user_id)
+        except Exception as e:
+            logging.error(f"Error saving token to InMemoryStorage: {e}")
+            raise
+
+    def delete_item(self, table, chat_id):
+        try:
+            table.pop(str(chat_id), None)
+        except Exception as e:
+            logging.error(f"Error deleting item from InMemoryStorage: {e}")
+            raise
+
+    def check_item_exists(self, table, chat_id):
+        try:
+            exists = str(chat_id) in table
+            logging.info(f"check_item_exists: {exists}")
+            return exists
+        except Exception as e:
+            logging.error(f"Error checking item in InMemoryStorage: {e}")
+            raise