From d2d4b8e3bef545de122553c678a548b7b42e3d09 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Mon, 20 May 2024 14:02:22 -0700 Subject: [PATCH 1/8] ADD: Custom CallBack Context and Storage --- bot.py | 254 +++++++++++++++-------------------- custom_callback_contex.py | 14 ++ polling_main.py | 25 ++-- storage/dynamodb_init.py | 26 ++++ storage/dynamodb_storage.py | 111 +++++++++++++++ storage/in_memory_storage.py | 26 ++++ storage/storage_interface.py | 27 ++++ 7 files changed, 328 insertions(+), 155 deletions(-) create mode 100644 custom_callback_contex.py create mode 100644 storage/dynamodb_init.py create mode 100644 storage/dynamodb_storage.py create mode 100644 storage/in_memory_storage.py create mode 100644 storage/storage_interface.py diff --git a/bot.py b/bot.py index 7dbdae8..3210bd9 100644 --- a/bot.py +++ b/bot.py @@ -16,10 +16,13 @@ filters, CallbackContext, Defaults, + ContextTypes, ) from spotipy.oauth2 import SpotifyOAuth, CacheHandler from spotipy.exceptions import SpotifyException import urllib.parse +from custom_callback_contex import CustomCallbackContext +from storage.dynamodb_init import bot_table, credentials_table if logging.getLogger().hasHandlers(): @@ -39,10 +42,10 @@ SPOTIFY_REDIRECT_URI = os.getenv("SPOTIFY_REDIRECT_URI") spotify_link_pattern = r"https://open\.spotify\.com/track/([a-zA-Z0-9]+)" # Dynamodb -dynamodb = boto3.resource("dynamodb", region_name="us-east-1") +# dynamodb = boto3.resource("dynamodb", region_name="us-east-1") -bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) -credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) +# bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) +# credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) # Enums for bot states @@ -81,123 +84,23 @@ def handle_spotify_auth(state, code): # return {"statusCode": 200, "body": json.dumps(event, indent=2)} -# ----------------------------------------------- -# DynamoDB Utility Functions -# ----------------------------------------------- -def save_current_state(chat_id, state_key: BotState): - # Saves the current bot state for a given chat_id in DynamoDB. - try: - user_id = get_user_id_from_chat_id(chat_id) - if user_id is None: - user_id = get_user_id_from_channel_credentials(chat_id) - - if state_key is BotState.NO_STATE: - response = bot_table.update_item( - Key={"chat_id": str(chat_id)}, - UpdateExpression="REMOVE current_state", - ReturnValues="UPDATED_NEW", - ) - else: - response = 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(chat_id): - # Retrieves the current bot state for a given chat_id from DynamoDB. - try: - response = 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(chat_id, playlist_id): - try: - bot_table.put_item( - Item={ - "chat_id": str(chat_id), - "playlist_id": playlist_id, - "user_id": 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(chat_id): - try: - response = 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(chat_id): - try: - response = 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(chat_id): - try: - response = 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: - print( - f"get_user_id_from_channel_credentials: No user_id found for chat_id: {chat_id}: {response}" - ) - return None - except Exception as e: - print( - f"Error retrieving user ID from DynamoDB for chat_id: {chat_id}, error: {e}" - ) - return None - - class DynamoCredentialsCache(CacheHandler): """ A cache handler that stores OAuth credentials in a Dynamo bot_table called 'ChannelCredentials' that has a primary key of chat_id. """ - def __init__(self, chat_id, user_id): + def __init__(self, chat_id, user_id, credentials_table, bot_table): self.chat_id = chat_id self.user_id = user_id + self.credentials_table = credentials_table + self.bot_table = bot_table def get_cached_token(self): try: - response = credentials_table.get_item(Key={"chat_id": str(self.chat_id)}) + response = self.credentials_table.get_item( + Key={"chat_id": str(self.chat_id)} + ) if "Item" in response: return response["Item"] return None @@ -207,14 +110,14 @@ def get_cached_token(self): def save_token_to_cache(self, token_info): try: - credentials_table.put_item( + self.credentials_table.put_item( Item={ "chat_id": str(self.chat_id), "user_id": self.user_id, **token_info, } ) - bot_table.update_item( + self.bot_table.update_item( Key={"chat_id": str(self.chat_id)}, UpdateExpression="SET user_id = :uid", ExpressionAttributeValues={":uid": str(self.user_id)}, @@ -232,7 +135,9 @@ def get_sp_oauth(chat_id, user_id): SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, SPOTIFY_REDIRECT_URI, - cache_handler=DynamoCredentialsCache(chat_id, user_id), + cache_handler=DynamoCredentialsCache( + chat_id, user_id, credentials_table, bot_table + ), scope="playlist-modify-public ugc-image-upload", ) @@ -272,23 +177,23 @@ def create_spotify_playlist(playlist_name, sp_oauth): # ----------------------------------------------- # Telegram Message Handlers # ----------------------------------------------- -async def handle_playlist_image(update: Update, context: CallbackContext) -> None: +async def handle_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_table = get_user_id_from_chat_id(chat_id) + user_id_table = context.storage.get_user_id_from_chat_id(chat_id) if user_id_table != str(user_id): await update.message.reply_text( "You are not authorized for this playlist process." ) return - current_state = get_current_state(chat_id) + current_state = context.storage.get_current_state(chat_id) if ( current_state == BotState.AWAITING_PLAYLIST_IMAGE or current_state == BotState.CHANGING_PLAYLIST_IMAGE ): photo = update.message.photo[-1] await update.message.reply_text("Processing your image, please wait...") - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if not playlist_id: await update.message.reply_text("No playlist found for this chat.") return @@ -320,7 +225,7 @@ async def upload_image(): ), ) - save_current_state(chat_id, BotState.NO_STATE) + context.storage.save_current_state(chat_id, BotState.NO_STATE) except asyncio.TimeoutError: logging.exception("Error Timeout: TimeoutError") await update.message.reply_text("Image upload timed out. Please try again.") @@ -329,13 +234,15 @@ async def upload_image(): await update.message.reply_text(f"An error occurred: {e}") -async def handle_playlist_name(update: Update, context: CallbackContext) -> None: +async def handle_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_bot_table = get_user_id_from_chat_id(chat_id) - user_id_credentials_table = get_user_id_from_channel_credentials(chat_id) + user_id_bot_table = context.storage.get_user_id_from_chat_id(chat_id) + user_id_credentials_table = context.storage.get_user_id_from_channel_credentials( + chat_id + ) - current_state = get_current_state(chat_id) + current_state = context.storage.get_current_state(chat_id) sp_oauth = get_sp_oauth(chat_id, user_id) if current_state == BotState.CHANGING_PLAYLIST_NAME: @@ -344,7 +251,7 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None "Please click on authorize link before entering playlist name." ) return - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if playlist_id: new_name = update.message.text.strip() if change_spotify_playlist_name(playlist_id, new_name, sp_oauth): @@ -353,7 +260,7 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None await update.message.reply_text( "Failed to change the playlist name. Please try again later." ) - save_current_state(chat_id, BotState.NO_STATE) + context.storage.save_current_state(chat_id, BotState.NO_STATE) else: await update.message.reply_text( "Make sure you have a playlist created before changing the name." @@ -368,8 +275,10 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None playlist_name = update.message.text.strip() try: playlist_id = create_spotify_playlist(playlist_name, sp_oauth) - save_playlist_to_dynamodb(chat_id, playlist_id) - save_current_state(chat_id, BotState.AWAITING_PLAYLIST_IMAGE) + context.storage.save_playlist_to_dynamodb(chat_id, playlist_id) + context.storage.save_current_state( + chat_id, BotState.AWAITING_PLAYLIST_IMAGE + ) await update.message.reply_text( f"Created new playlist: {playlist_name}. " "Now please send me a cool image to set as your playlist cover 😎." @@ -384,19 +293,19 @@ async def handle_playlist_name(update: Update, context: CallbackContext) -> None return -async def handle_spotify_links(update: Update, context: CallbackContext) -> None: +async def handle_spotify_links(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id message_text = update.message.text match = re.search(spotify_link_pattern, message_text) - current_state = get_current_state(chat_id) + current_state = context.storage.get_current_state(chat_id) if current_state == BotState.CREATING_PLAYLIST: await update.message.reply_text( "You are in the process of creating a playlist. " "Please wait until it's done before sending links." ) return - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if match and playlist_id: track_id = match.group(1) sp_oauth = get_sp_oauth(chat_id, user_id) @@ -416,31 +325,31 @@ async def handle_spotify_links(update: Update, context: CallbackContext) -> None # ----------------------------------------------- # Telegram Command Handlers # ----------------------------------------------- -async def change_playlist_image(update: Update, context: CallbackContext) -> None: +async def change_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please send the new image for your playlist:") - save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) + context.storage.save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." ) -async def change_playlist_name(update: Update, context: CallbackContext) -> None: +async def change_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please enter the new name for your playlist:") - save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) + context.storage.save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." ) -async def create_playlist(update: Update, context: CallbackContext) -> bool: +async def create_playlist(update: Update, context: CustomCallbackContext) -> bool: chat_id = update.effective_chat.id user_id = update.effective_user.id @@ -448,9 +357,9 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: state_encoded = json.dumps(state_info) state_url_safe = urllib.parse.quote(state_encoded) - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) - current_state = get_current_state(chat_id) + current_state = context.storage.get_current_state(chat_id) if playlist_id: await update.message.reply_text( @@ -463,7 +372,7 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: "You are already in the process of creating a playlist." ) return False - save_current_state(chat_id, BotState.CREATING_PLAYLIST) + context.storage.save_current_state(chat_id, BotState.CREATING_PLAYLIST) sp_oauth = get_sp_oauth(chat_id, user_id) token_info = sp_oauth.cache_handler.get_cached_token() if sp_oauth.validate_token(token_info) is None: @@ -476,7 +385,7 @@ async def create_playlist(update: Update, context: CallbackContext) -> bool: return True -async def help_command(update: Update, context: CallbackContext) -> None: +async def help_command(update: Update, context: CustomCallbackContext) -> None: help_text = ( "Here are the commands you can use:\n" "/start - Start interacting with the bot\n" @@ -492,7 +401,7 @@ async def help_command(update: Update, context: CallbackContext) -> None: await update.message.reply_text(help_text) -async def reset_playlist(update: Update, context: CallbackContext) -> None: +async def reset_playlist(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id # Delete the playlist entry from DynamoDB try: @@ -507,9 +416,9 @@ async def reset_playlist(update: Update, context: CallbackContext) -> None: ) -async def send_playlist_link(update: Update, context: CallbackContext) -> None: +async def send_playlist_link(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) if playlist_id: playlist_url = f"https://open.spotify.com/playlist/{playlist_id}" await update.message.reply_text( @@ -523,18 +432,57 @@ async def send_playlist_link(update: Update, context: CallbackContext) -> None: await update.message.reply_text("No playlist found for this chat.") -async def start(update: Update, context: CallbackContext) -> None: +async def start(update: Update, context: CustomCallbackContext) -> None: await update.message.reply_text( "Hiya! I'm your Spotify Skunk bot 🦨. /createplaylist " "to add songs to your playlist!" ) -async def unlink_credentials(update: Update, context: CallbackContext) -> None: +async def unlink_credentials(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id + # Log the storage object and tables + # logging.info(f"Storage object: {context.storage}") + # logging.info(f"Credentials table: {context.storage.credentials_table}") + # logging.info(f"Bot table: {context.storage.bot_table}") try: - credentials_table.delete_item(Key={"chat_id": str(chat_id)}) - bot_table.delete_item(Key={"chat_id": str(chat_id)}) + # Check if the item exists in the credentials table + credentials_response = context.storage.credentials_table.get_item( + Key={"chat_id": str(chat_id)} + ) + logging.info(f"Credentials get_item response: {credentials_response}") + if "Item" not in credentials_response: + logging.error(f"Credentials not found for chat_id {chat_id}") + await update.message.reply_text( + "Failed to unlink your Spotify credentials: credentials not found." + ) + return + + # Check if the item exists in the bot table + bot_response = context.storage.bot_table.get_item(Key={"chat_id": str(chat_id)}) + logging.info(f"Bot get_item response: {bot_response}") + if "Item" not in bot_response: + logging.error(f"Bot data not found for chat_id {chat_id}") + await update.message.reply_text( + "Failed to unlink your Spotify credentials: bot data not found." + ) + return + + # Log the deletion attempt + logging.info( + f"Attempting to delete item from credentials table for chat_id {chat_id}" + ) + credentials_delete_response = context.storage.credentials_table.delete_item( + Key={"chat_id": str(chat_id)} + ) + logging.info(f"Credentials delete_item response: {credentials_delete_response}") + + logging.info(f"Attempting to delete item from bot table for chat_id {chat_id}") + bot_delete_response = context.storage.bot_table.delete_item( + Key={"chat_id": str(chat_id)} + ) + logging.info(f"Bot delete_item response: {bot_delete_response}") + await update.message.reply_text( "Your Spotify credentials have been unlinked successfully." ) @@ -545,9 +493,21 @@ async def unlink_credentials(update: Update, context: CallbackContext) -> None: ) -def build_application(token): +def build_application(token, storage): logger.info(f"token: {token}") application = Application.builder().token(token).defaults(defaults).build() + + # Set the custom context for each handler + context_types = ContextTypes(context=CustomCallbackContext) + application = ( + Application.builder() + .token(token) + .context_types(context_types) + .defaults(defaults) + .build() + ) + application.bot_data["storage"] = storage + register_handlers(application) return application diff --git a/custom_callback_contex.py b/custom_callback_contex.py new file mode 100644 index 0000000..65d8f86 --- /dev/null +++ b/custom_callback_contex.py @@ -0,0 +1,14 @@ +from telegram.ext import CallbackContext as BaseCallbackContext + + +class CustomCallbackContext(BaseCallbackContext): + def __init__(self, application, chat_id=None, user_id=None, storage=None): + super().__init__(application=application, chat_id=chat_id, user_id=user_id) + self.storage = storage + + @classmethod + def from_update(cls, update, application): + """Override from_update to set storage.""" + context = super().from_update(update, application) + context.storage = application.bot_data.get("storage") + return context diff --git a/polling_main.py b/polling_main.py index e36a74d..b835ff2 100644 --- a/polling_main.py +++ b/polling_main.py @@ -5,6 +5,10 @@ from bot import handle_spotify_auth import logging from flask import Response +from storage.dynamodb_storage import DynamoDBStorage +from storage.in_memory_storage import InMemoryStorage +from storage.dynamodb_init import bot_table, credentials_table +import boto3 webserver = Flask(__name__) @@ -34,17 +38,22 @@ def run_flask_app(): if __name__ == "__main__": TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") - application = build_application(TOKEN) + + # Initialize storage + dynamodb = boto3.resource("dynamodb", region_name="us-east-1") + # bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) + # credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) + dynamodb_storage = DynamoDBStorage(bot_table, credentials_table) + in_memory_storage = InMemoryStorage() + + # Use DynamoDB storage for production, in-memory storage for testing + storage = dynamodb_storage + + 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() + diff --git a/storage/dynamodb_init.py b/storage/dynamodb_init.py new file mode 100644 index 0000000..76ce7e2 --- /dev/null +++ b/storage/dynamodb_init.py @@ -0,0 +1,26 @@ +import os +import boto3 +import logging + +# Initialize logging +if logging.getLogger().hasHandlers(): + logging.getLogger().setLevel(logging.INFO) +else: + logging.basicConfig(level=logging.INFO) +logger = logging.getLogger() + +# Initialize the DynamoDB resource and tables +dynamodb = boto3.resource("dynamodb", region_name="us-east-1") + +bot_table_name = os.getenv("BOT_TABLE") +credentials_table_name = os.getenv("CREDENTIALS_TABLE") + +logger.info( + f"Initializing DynamoDB tables: BOT_TABLE={bot_table_name}, CREDENTIALS_TABLE={credentials_table_name}" +) + +bot_table = dynamodb.Table(bot_table_name) +credentials_table = dynamodb.Table(credentials_table_name) + + +__all__ = ["bot_table", "credentials_table"] diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py new file mode 100644 index 0000000..57dbb5f --- /dev/null +++ b/storage/dynamodb_storage.py @@ -0,0 +1,111 @@ +import boto3 +import logging +from .storage_interface import Storage +from enum import Enum + + +class BotState(Enum): + AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" + CHANGING_PLAYLIST_NAME = "changing_playlist_name" + CREATING_PLAYLIST = "creating_playlist" + CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" + NO_STATE = None + + +class DynamoDBStorage(Storage): + def __init__(self, bot_table, credentials_table): + self.bot_table = bot_table + self.credentials_table = 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 diff --git a/storage/in_memory_storage.py b/storage/in_memory_storage.py new file mode 100644 index 0000000..80ea66b --- /dev/null +++ b/storage/in_memory_storage.py @@ -0,0 +1,26 @@ +from .storage_interface import Storage + + +class InMemoryStorage(Storage): + def __init__(self): + self.state = {} + self.playlists = {} + self.users = {} + + def save_current_state(self, chat_id, state_key): + self.state[chat_id] = state_key + + def get_current_state(self, chat_id): + return self.state.get(chat_id, None) + + def save_playlist_to_dynamodb(self, chat_id, playlist_id): + self.playlists[chat_id] = playlist_id + + def get_playlist_from_dynamodb(self, chat_id): + return self.playlists.get(chat_id, None) + + def get_user_id_from_chat_id(self, chat_id): + return self.users.get(chat_id, None) + + def get_user_id_from_channel_credentials(self, chat_id): + return self.users.get(chat_id, None) diff --git a/storage/storage_interface.py b/storage/storage_interface.py new file mode 100644 index 0000000..4296fa7 --- /dev/null +++ b/storage/storage_interface.py @@ -0,0 +1,27 @@ +from abc import ABC, abstractmethod + + +class Storage(ABC): + @abstractmethod + def save_current_state(self, chat_id, state_key): + pass + + @abstractmethod + def get_current_state(self, chat_id): + pass + + @abstractmethod + def save_playlist_to_dynamodb(self, chat_id, playlist_id): + pass + + @abstractmethod + def get_playlist_from_dynamodb(self, chat_id): + pass + + @abstractmethod + def get_user_id_from_chat_id(self, chat_id): + pass + + @abstractmethod + def get_user_id_from_channel_credentials(self, chat_id): + pass From 8b3651f2bb8013060f2aacfa678df4c9ab2d58f1 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Mon, 20 May 2024 19:31:12 -0700 Subject: [PATCH 2/8] ADD: Custom CallBack Context and Storage Changes --- bot.py | 81 ++++++++++++++++++++----------------- custom_callback_contex.py | 22 +++++----- polling_main.py | 8 ++-- storage/dynamodb_init.py | 26 ------------ storage/dynamodb_storage.py | 9 +++-- 5 files changed, 65 insertions(+), 81 deletions(-) delete mode 100644 storage/dynamodb_init.py diff --git a/bot.py b/bot.py index 3210bd9..ead552f 100644 --- a/bot.py +++ b/bot.py @@ -22,7 +22,9 @@ from spotipy.exceptions import SpotifyException import urllib.parse from custom_callback_contex import CustomCallbackContext -from storage.dynamodb_init import bot_table, credentials_table +from storage.dynamodb_storage import DynamoDBStorage + +# from storage.dynamodb_init import bot_table, credentials_table if logging.getLogger().hasHandlers(): @@ -90,11 +92,11 @@ class DynamoCredentialsCache(CacheHandler): 'ChannelCredentials' that has a primary key of chat_id. """ - def __init__(self, chat_id, user_id, credentials_table, bot_table): + def __init__(self, chat_id, user_id, storage): self.chat_id = chat_id self.user_id = user_id - self.credentials_table = credentials_table - self.bot_table = bot_table + self.credentials_table = storage.credentials_table + self.bot_table = storage.bot_table def get_cached_token(self): try: @@ -131,13 +133,12 @@ def save_token_to_cache(self, token_info): # Spotify Utility Functions # ----------------------------------------------- def get_sp_oauth(chat_id, user_id): + # TODO: Pass Storage Param return SpotifyOAuth( SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, SPOTIFY_REDIRECT_URI, - cache_handler=DynamoCredentialsCache( - chat_id, user_id, credentials_table, bot_table - ), + cache_handler=DynamoCredentialsCache(chat_id, user_id, DynamoDBStorage()), scope="playlist-modify-public ugc-image-upload", ) @@ -180,20 +181,20 @@ def create_spotify_playlist(playlist_name, sp_oauth): async def handle_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_table = context.storage.get_user_id_from_chat_id(chat_id) + user_id_table = context.storage().get_user_id_from_chat_id(chat_id) if user_id_table != str(user_id): await update.message.reply_text( "You are not authorized for this playlist process." ) return - current_state = context.storage.get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if ( current_state == BotState.AWAITING_PLAYLIST_IMAGE or current_state == BotState.CHANGING_PLAYLIST_IMAGE ): photo = update.message.photo[-1] await update.message.reply_text("Processing your image, please wait...") - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if not playlist_id: await update.message.reply_text("No playlist found for this chat.") return @@ -225,7 +226,7 @@ async def upload_image(): ), ) - context.storage.save_current_state(chat_id, BotState.NO_STATE) + context.storage().save_current_state(chat_id, BotState.NO_STATE) except asyncio.TimeoutError: logging.exception("Error Timeout: TimeoutError") await update.message.reply_text("Image upload timed out. Please try again.") @@ -237,13 +238,15 @@ async def upload_image(): async def handle_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id user_id = update.effective_user.id - user_id_bot_table = context.storage.get_user_id_from_chat_id(chat_id) - user_id_credentials_table = context.storage.get_user_id_from_channel_credentials( + user_id_bot_table = context.storage().get_user_id_from_chat_id(chat_id) + user_id_credentials_table = context.storage().get_user_id_from_channel_credentials( chat_id ) - current_state = context.storage.get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) + logging.info(f"the state here-pre is: {current_state}, {type(current_state)}") sp_oauth = get_sp_oauth(chat_id, user_id) + logging.info(f"here is sp oauth: {sp_oauth}") if current_state == BotState.CHANGING_PLAYLIST_NAME: if user_id_bot_table != str(user_id): @@ -251,7 +254,7 @@ async def handle_playlist_name(update: Update, context: CustomCallbackContext) - "Please click on authorize link before entering playlist name." ) return - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: new_name = update.message.text.strip() if change_spotify_playlist_name(playlist_id, new_name, sp_oauth): @@ -260,12 +263,13 @@ async def handle_playlist_name(update: Update, context: CustomCallbackContext) - await update.message.reply_text( "Failed to change the playlist name. Please try again later." ) - context.storage.save_current_state(chat_id, BotState.NO_STATE) + context.storage().save_current_state(chat_id, BotState.NO_STATE) else: await update.message.reply_text( "Make sure you have a playlist created before changing the name." ) - elif current_state == BotState.CREATING_PLAYLIST: + elif current_state is BotState.CREATING_PLAYLIST: + logging.info(f"the state here is: {current_state}") if str(user_id) != user_id_credentials_table: await update.message.reply_text( "You are not authorized for this playlist process." @@ -275,8 +279,8 @@ async def handle_playlist_name(update: Update, context: CustomCallbackContext) - playlist_name = update.message.text.strip() try: playlist_id = create_spotify_playlist(playlist_name, sp_oauth) - context.storage.save_playlist_to_dynamodb(chat_id, playlist_id) - context.storage.save_current_state( + context.storage().save_playlist_to_dynamodb(chat_id, playlist_id) + context.storage().save_current_state( chat_id, BotState.AWAITING_PLAYLIST_IMAGE ) await update.message.reply_text( @@ -298,14 +302,14 @@ async def handle_spotify_links(update: Update, context: CustomCallbackContext) - user_id = update.effective_user.id message_text = update.message.text match = re.search(spotify_link_pattern, message_text) - current_state = context.storage.get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if current_state == BotState.CREATING_PLAYLIST: await update.message.reply_text( "You are in the process of creating a playlist. " "Please wait until it's done before sending links." ) return - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if match and playlist_id: track_id = match.group(1) sp_oauth = get_sp_oauth(chat_id, user_id) @@ -327,10 +331,10 @@ async def handle_spotify_links(update: Update, context: CustomCallbackContext) - # ----------------------------------------------- async def change_playlist_image(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please send the new image for your playlist:") - context.storage.save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) + context.storage().save_current_state(chat_id, BotState.CHANGING_PLAYLIST_IMAGE) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." @@ -339,10 +343,10 @@ async def change_playlist_image(update: Update, context: CustomCallbackContext) async def change_playlist_name(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: await update.message.reply_text("Please enter the new name for your playlist:") - context.storage.save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) + context.storage().save_current_state(chat_id, BotState.CHANGING_PLAYLIST_NAME) else: await update.message.reply_text( "No playlist found for this chat. Create one with /createplaylist." @@ -357,9 +361,9 @@ async def create_playlist(update: Update, context: CustomCallbackContext) -> boo state_encoded = json.dumps(state_info) state_url_safe = urllib.parse.quote(state_encoded) - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) - current_state = context.storage.get_current_state(chat_id) + current_state = context.storage().get_current_state(chat_id) if playlist_id: await update.message.reply_text( @@ -372,7 +376,7 @@ async def create_playlist(update: Update, context: CustomCallbackContext) -> boo "You are already in the process of creating a playlist." ) return False - context.storage.save_current_state(chat_id, BotState.CREATING_PLAYLIST) + context.storage().save_current_state(chat_id, BotState.CREATING_PLAYLIST) sp_oauth = get_sp_oauth(chat_id, user_id) token_info = sp_oauth.cache_handler.get_cached_token() if sp_oauth.validate_token(token_info) is None: @@ -405,7 +409,7 @@ async def reset_playlist(update: Update, context: CustomCallbackContext) -> None chat_id = update.effective_chat.id # Delete the playlist entry from DynamoDB try: - bot_table.delete_item(Key={"chat_id": str(chat_id)}) + context.storage().bot_table.delete_item(Key={"chat_id": str(chat_id)}) except Exception as e: logging.error(f"Error deleting from DynamoDB: {e}") await update.message.reply_text("Failed to reset the playlist in the database.") @@ -418,7 +422,7 @@ async def reset_playlist(update: Update, context: CustomCallbackContext) -> None async def send_playlist_link(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - playlist_id = context.storage.get_playlist_from_dynamodb(chat_id) + playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if playlist_id: playlist_url = f"https://open.spotify.com/playlist/{playlist_id}" await update.message.reply_text( @@ -442,12 +446,12 @@ async def start(update: Update, context: CustomCallbackContext) -> None: async def unlink_credentials(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id # Log the storage object and tables - # logging.info(f"Storage object: {context.storage}") - # logging.info(f"Credentials table: {context.storage.credentials_table}") - # logging.info(f"Bot table: {context.storage.bot_table}") + # logging.info(f"Storage object: {context.storage()}") + # logging.info(f"Credentials table: {context.storage().credentials_table}") + # logging.info(f"Bot table: {context.storage().bot_table}") try: # Check if the item exists in the credentials table - credentials_response = context.storage.credentials_table.get_item( + credentials_response = context.storage().credentials_table.get_item( Key={"chat_id": str(chat_id)} ) logging.info(f"Credentials get_item response: {credentials_response}") @@ -459,7 +463,9 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> return # Check if the item exists in the bot table - bot_response = context.storage.bot_table.get_item(Key={"chat_id": str(chat_id)}) + bot_response = context.storage().bot_table.get_item( + Key={"chat_id": str(chat_id)} + ) logging.info(f"Bot get_item response: {bot_response}") if "Item" not in bot_response: logging.error(f"Bot data not found for chat_id {chat_id}") @@ -472,13 +478,13 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> logging.info( f"Attempting to delete item from credentials table for chat_id {chat_id}" ) - credentials_delete_response = context.storage.credentials_table.delete_item( + credentials_delete_response = context.storage().credentials_table.delete_item( Key={"chat_id": str(chat_id)} ) logging.info(f"Credentials delete_item response: {credentials_delete_response}") logging.info(f"Attempting to delete item from bot table for chat_id {chat_id}") - bot_delete_response = context.storage.bot_table.delete_item( + bot_delete_response = context.storage().bot_table.delete_item( Key={"chat_id": str(chat_id)} ) logging.info(f"Bot delete_item response: {bot_delete_response}") @@ -495,7 +501,6 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> def build_application(token, storage): logger.info(f"token: {token}") - application = Application.builder().token(token).defaults(defaults).build() # Set the custom context for each handler context_types = ContextTypes(context=CustomCallbackContext) diff --git a/custom_callback_contex.py b/custom_callback_contex.py index 65d8f86..d265c4b 100644 --- a/custom_callback_contex.py +++ b/custom_callback_contex.py @@ -1,14 +1,18 @@ 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, storage=None): + def __init__(self, application, chat_id=None, user_id=None): super().__init__(application=application, chat_id=chat_id, user_id=user_id) - self.storage = storage - - @classmethod - def from_update(cls, update, application): - """Override from_update to set storage.""" - context = super().from_update(update, application) - context.storage = application.bot_data.get("storage") - return context + + def storage(self) -> DynamoDBStorage: + return self.bot_data["storage"] + + # @classmethod + # def from_update(cls, update, application): + # """Override from_update to set storage.""" + # context = super().from_update(update, application) + # context.storage = application.bot_data.get("storage") + # return context diff --git a/polling_main.py b/polling_main.py index b835ff2..14183a6 100644 --- a/polling_main.py +++ b/polling_main.py @@ -7,7 +7,7 @@ from flask import Response from storage.dynamodb_storage import DynamoDBStorage from storage.in_memory_storage import InMemoryStorage -from storage.dynamodb_init import bot_table, credentials_table +# from storage.dynamodb_init import bot_table, credentials_table import boto3 webserver = Flask(__name__) @@ -40,11 +40,11 @@ def run_flask_app(): TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") # Initialize storage - dynamodb = boto3.resource("dynamodb", region_name="us-east-1") + # dynamodb = boto3.resource("dynamodb", region_name="us-east-1") # bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) # credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) - dynamodb_storage = DynamoDBStorage(bot_table, credentials_table) - in_memory_storage = InMemoryStorage() + dynamodb_storage = DynamoDBStorage() + # in_memory_storage = InMemoryStorage() # Use DynamoDB storage for production, in-memory storage for testing storage = dynamodb_storage diff --git a/storage/dynamodb_init.py b/storage/dynamodb_init.py deleted file mode 100644 index 76ce7e2..0000000 --- a/storage/dynamodb_init.py +++ /dev/null @@ -1,26 +0,0 @@ -import os -import boto3 -import logging - -# Initialize logging -if logging.getLogger().hasHandlers(): - logging.getLogger().setLevel(logging.INFO) -else: - logging.basicConfig(level=logging.INFO) -logger = logging.getLogger() - -# Initialize the DynamoDB resource and tables -dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - -bot_table_name = os.getenv("BOT_TABLE") -credentials_table_name = os.getenv("CREDENTIALS_TABLE") - -logger.info( - f"Initializing DynamoDB tables: BOT_TABLE={bot_table_name}, CREDENTIALS_TABLE={credentials_table_name}" -) - -bot_table = dynamodb.Table(bot_table_name) -credentials_table = dynamodb.Table(credentials_table_name) - - -__all__ = ["bot_table", "credentials_table"] diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py index 57dbb5f..6f7c126 100644 --- a/storage/dynamodb_storage.py +++ b/storage/dynamodb_storage.py @@ -2,7 +2,7 @@ import logging from .storage_interface import Storage from enum import Enum - +import os class BotState(Enum): AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" @@ -13,9 +13,10 @@ class BotState(Enum): class DynamoDBStorage(Storage): - def __init__(self, bot_table, credentials_table): - self.bot_table = bot_table - self.credentials_table = credentials_table + 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: From 645441f1d6c185d525b866b6838f5098a51be80d Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Mon, 20 May 2024 23:02:59 -0700 Subject: [PATCH 3/8] ADD: Custom CallBack Context and Storage Changes --- bot.py | 54 ++++++++++++++++-------------------- custom_callback_contex.py | 7 ----- polling_main.py | 11 ++------ requirements.txt | 2 +- storage/dynamodb_storage.py | 11 ++------ storage/storage_interface.py | 9 ++++++ 6 files changed, 38 insertions(+), 56 deletions(-) diff --git a/bot.py b/bot.py index ead552f..a0034fe 100644 --- a/bot.py +++ b/bot.py @@ -1,20 +1,17 @@ import asyncio import base64 -import boto3 import os import spotipy import logging import re import json import telegram -from enum import Enum from telegram import Update, LinkPreviewOptions from telegram.ext import ( Application, CommandHandler, MessageHandler, filters, - CallbackContext, Defaults, ContextTypes, ) @@ -22,10 +19,7 @@ from spotipy.exceptions import SpotifyException import urllib.parse from custom_callback_contex import CustomCallbackContext -from storage.dynamodb_storage import DynamoDBStorage - -# from storage.dynamodb_init import bot_table, credentials_table - +from storage.storage_interface import BotState if logging.getLogger().hasHandlers(): logging.getLogger().setLevel(logging.INFO) @@ -34,7 +28,6 @@ logger = logging.getLogger() # Constants and configurations - defaults = Defaults( link_preview_options=LinkPreviewOptions(show_above_text=False, is_disabled=True), disable_notification=True, @@ -43,20 +36,6 @@ SPOTIFY_CLIENT_SECRET = os.getenv("SPOTIFY_CLIENT_SECRET") SPOTIFY_REDIRECT_URI = os.getenv("SPOTIFY_REDIRECT_URI") spotify_link_pattern = r"https://open\.spotify\.com/track/([a-zA-Z0-9]+)" -# Dynamodb -# dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - -# bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) -# credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) - - -# Enums for bot states -class BotState(Enum): - AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" - CHANGING_PLAYLIST_NAME = "changing_playlist_name" - CREATING_PLAYLIST = "creating_playlist" - CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" - NO_STATE = None def load_html_file(file_name): @@ -64,13 +43,13 @@ def load_html_file(file_name): return file.read() -def handle_spotify_auth(state, code): +def handle_spotify_auth(state, code, storage): state_decoded = urllib.parse.unquote(state) state_info = json.loads(state_decoded) chat_id = state_info.get("chat_id") user_id = state_info.get("user_id") - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, storage) token_info = sp_oauth.get_access_token(code) if token_info: @@ -132,13 +111,13 @@ def save_token_to_cache(self, token_info): # ----------------------------------------------- # Spotify Utility Functions # ----------------------------------------------- -def get_sp_oauth(chat_id, user_id): +def get_sp_oauth(chat_id, user_id, storage): # TODO: Pass Storage Param return SpotifyOAuth( SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, SPOTIFY_REDIRECT_URI, - cache_handler=DynamoCredentialsCache(chat_id, user_id, DynamoDBStorage()), + cache_handler=DynamoCredentialsCache(chat_id, user_id, storage), scope="playlist-modify-public ugc-image-upload", ) @@ -210,7 +189,7 @@ async def handle_playlist_image(update: Update, context: CustomCallbackContext) return async def upload_image(): - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) sp = spotipy.Spotify(auth_manager=sp_oauth) sp.playlist_upload_cover_image(playlist_id, base64_image) @@ -245,7 +224,7 @@ async def handle_playlist_name(update: Update, context: CustomCallbackContext) - current_state = context.storage().get_current_state(chat_id) logging.info(f"the state here-pre is: {current_state}, {type(current_state)}") - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) logging.info(f"here is sp oauth: {sp_oauth}") if current_state == BotState.CHANGING_PLAYLIST_NAME: @@ -312,7 +291,7 @@ async def handle_spotify_links(update: Update, context: CustomCallbackContext) - playlist_id = context.storage().get_playlist_from_dynamodb(chat_id) if match and playlist_id: track_id = match.group(1) - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) if add_track_to_spotify_playlist(playlist_id, track_id, sp_oauth): await update.message.set_reaction("👍") else: @@ -377,7 +356,7 @@ async def create_playlist(update: Update, context: CustomCallbackContext) -> boo ) return False context.storage().save_current_state(chat_id, BotState.CREATING_PLAYLIST) - sp_oauth = get_sp_oauth(chat_id, user_id) + sp_oauth = get_sp_oauth(chat_id, user_id, context.storage()) token_info = sp_oauth.cache_handler.get_cached_token() if sp_oauth.validate_token(token_info) is None: auth_url = sp_oauth.get_authorize_url(state=state_url_safe) @@ -499,6 +478,20 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> ) +# TODO: Fix so it Only Works for You and in Group Chats +async def debug_stuff(update: Update, context: CustomCallbackContext): + chat_id = update.effective_chat.id + debug = "\n".join( + [ + f"Owning User ID: `{context.storage().get_user_id_from_chat_id(chat_id)}`", + f"Chat ID: `{context.storage().get_current_state(chat_id)}`", + f"Playlist: `{context.storage().get_playlist_from_dynamodb(chat_id)}`", + f"User ID from Channel Credentials: `{context.storage().get_user_id_from_channel_credentials(chat_id)}`", + ] + ) + await update.message.reply_text(debug) + + def build_application(token, storage): logger.info(f"token: {token}") @@ -527,6 +520,7 @@ def register_handlers(application: Application): CommandHandler("changeplaylistimage", change_playlist_image), CommandHandler("playlistlink", send_playlist_link), CommandHandler("unlink", unlink_credentials), + CommandHandler("debug", debug_stuff), MessageHandler( filters.TEXT & filters.Regex(spotify_link_pattern), handle_spotify_links ), diff --git a/custom_callback_contex.py b/custom_callback_contex.py index d265c4b..7412342 100644 --- a/custom_callback_contex.py +++ b/custom_callback_contex.py @@ -9,10 +9,3 @@ def __init__(self, application, chat_id=None, user_id=None): def storage(self) -> DynamoDBStorage: return self.bot_data["storage"] - - # @classmethod - # def from_update(cls, update, application): - # """Override from_update to set storage.""" - # context = super().from_update(update, application) - # context.storage = application.bot_data.get("storage") - # return context diff --git a/polling_main.py b/polling_main.py index 14183a6..2b3f183 100644 --- a/polling_main.py +++ b/polling_main.py @@ -6,9 +6,6 @@ import logging from flask import Response from storage.dynamodb_storage import DynamoDBStorage -from storage.in_memory_storage import InMemoryStorage -# from storage.dynamodb_init import bot_table, credentials_table -import boto3 webserver = Flask(__name__) @@ -28,7 +25,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") @@ -40,13 +37,10 @@ def run_flask_app(): TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") # Initialize storage - # dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - # bot_table = dynamodb.Table(os.getenv("BOT_TABLE")) - # credentials_table = dynamodb.Table(os.getenv("CREDENTIALS_TABLE")) dynamodb_storage = DynamoDBStorage() # in_memory_storage = InMemoryStorage() - # Use DynamoDB storage for production, in-memory storage for testing + # TODO: Use DynamoDB storage for production, in-memory storage for testing storage = dynamodb_storage application = build_application(TOKEN, storage) @@ -56,4 +50,3 @@ def run_flask_app(): # Start polling application.run_polling() - diff --git a/requirements.txt b/requirements.txt index 564564f..1d69188 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,7 +16,7 @@ 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==21.1.1 PyYAML==6.0.1 redis==5.0.1 requests==2.31.0 diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py index 6f7c126..80dc6b8 100644 --- a/storage/dynamodb_storage.py +++ b/storage/dynamodb_storage.py @@ -1,21 +1,14 @@ import boto3 import logging -from .storage_interface import Storage +from .storage_interface import BotState, Storage from enum import Enum import os -class BotState(Enum): - AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" - CHANGING_PLAYLIST_NAME = "changing_playlist_name" - CREATING_PLAYLIST = "creating_playlist" - CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" - NO_STATE = None - 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.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): diff --git a/storage/storage_interface.py b/storage/storage_interface.py index 4296fa7..9f998f9 100644 --- a/storage/storage_interface.py +++ b/storage/storage_interface.py @@ -1,4 +1,13 @@ from abc import ABC, abstractmethod +from enum import Enum + + +class BotState(Enum): + AWAITING_PLAYLIST_IMAGE = "awaiting_playlist_image" + CHANGING_PLAYLIST_NAME = "changing_playlist_name" + CREATING_PLAYLIST = "creating_playlist" + CHANGING_PLAYLIST_IMAGE = "changing_playlist_image" + NO_STATE = None class Storage(ABC): From b794816c87d647a22ce9b8bba6a742c109ed184c Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Wed, 22 May 2024 15:06:29 -0700 Subject: [PATCH 4/8] ADD: Custom CallBack Context and Storage Changes --- bot.py | 62 +++++++++++------------------------- polling_main.py | 5 +-- storage/dynamodb_storage.py | 45 +++++++++++++++++++++++++- storage/in_memory_storage.py | 60 +++++++++++++++++++++++++++++++--- 4 files changed, 121 insertions(+), 51 deletions(-) diff --git a/bot.py b/bot.py index a0034fe..214ed7c 100644 --- a/bot.py +++ b/bot.py @@ -74,37 +74,20 @@ class DynamoCredentialsCache(CacheHandler): def __init__(self, chat_id, user_id, storage): self.chat_id = chat_id self.user_id = user_id - self.credentials_table = storage.credentials_table - self.bot_table = storage.bot_table + self.storage = storage def get_cached_token(self): try: - response = self.credentials_table.get_item( - Key={"chat_id": str(self.chat_id)} - ) - if "Item" in response: - return response["Item"] - return None + return self.storage.get_cached_token(self.chat_id) except Exception as e: - logging.error(f"Error retrieving from DynamoDB: {e}") + logging.error(f"Error retrieving cached token from DynamoDB: {e}") raise def save_token_to_cache(self, token_info): try: - self.credentials_table.put_item( - Item={ - "chat_id": str(self.chat_id), - "user_id": self.user_id, - **token_info, - } - ) - self.bot_table.update_item( - Key={"chat_id": str(self.chat_id)}, - UpdateExpression="SET user_id = :uid", - ExpressionAttributeValues={":uid": str(self.user_id)}, - ) + self.storage.save_token_to_cache(self.chat_id, self.user_id, token_info) except Exception as e: - logging.error(f"Error saving to DynamoDB: {e}") + logging.error(f"Error saving token to DynamoDB: {e}") raise @@ -221,6 +204,7 @@ async def handle_playlist_name(update: Update, context: CustomCallbackContext) - user_id_credentials_table = context.storage().get_user_id_from_channel_credentials( chat_id ) + logger.info(f"user id table: {user_id_credentials_table}, user id: {user_id} ") current_state = context.storage().get_current_state(chat_id) logging.info(f"the state here-pre is: {current_state}, {type(current_state)}") @@ -388,9 +372,8 @@ async def reset_playlist(update: Update, context: CustomCallbackContext) -> None chat_id = update.effective_chat.id # Delete the playlist entry from DynamoDB try: - context.storage().bot_table.delete_item(Key={"chat_id": str(chat_id)}) - except Exception as e: - logging.error(f"Error deleting from DynamoDB: {e}") + context.storage().delete_item(context.storage().bot_table, chat_id) + except Exception: await update.message.reply_text("Failed to reset the playlist in the database.") return await update.message.reply_text( @@ -424,17 +407,11 @@ async def start(update: Update, context: CustomCallbackContext) -> None: async def unlink_credentials(update: Update, context: CustomCallbackContext) -> None: chat_id = update.effective_chat.id - # Log the storage object and tables - # logging.info(f"Storage object: {context.storage()}") - # logging.info(f"Credentials table: {context.storage().credentials_table}") - # logging.info(f"Bot table: {context.storage().bot_table}") try: # Check if the item exists in the credentials table - credentials_response = context.storage().credentials_table.get_item( - Key={"chat_id": str(chat_id)} - ) - logging.info(f"Credentials get_item response: {credentials_response}") - if "Item" not in credentials_response: + if not context.storage().check_item_exists( + "credentials_table", chat_id + ): logging.error(f"Credentials not found for chat_id {chat_id}") await update.message.reply_text( "Failed to unlink your Spotify credentials: credentials not found." @@ -442,11 +419,10 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> return # Check if the item exists in the bot table - bot_response = context.storage().bot_table.get_item( - Key={"chat_id": str(chat_id)} - ) - logging.info(f"Bot get_item response: {bot_response}") - if "Item" not in bot_response: + + if not context.storage().check_item_exists( + "bot_table", chat_id + ): logging.error(f"Bot data not found for chat_id {chat_id}") await update.message.reply_text( "Failed to unlink your Spotify credentials: bot data not found." @@ -457,14 +433,14 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> logging.info( f"Attempting to delete item from credentials table for chat_id {chat_id}" ) - credentials_delete_response = context.storage().credentials_table.delete_item( - Key={"chat_id": str(chat_id)} + credentials_delete_response = context.storage().delete_item( + "credentials_table", chat_id ) logging.info(f"Credentials delete_item response: {credentials_delete_response}") logging.info(f"Attempting to delete item from bot table for chat_id {chat_id}") - bot_delete_response = context.storage().bot_table.delete_item( - Key={"chat_id": str(chat_id)} + bot_delete_response = context.storage().delete_item( + "bot_table", chat_id ) logging.info(f"Bot delete_item response: {bot_delete_response}") diff --git a/polling_main.py b/polling_main.py index 2b3f183..8140ab7 100644 --- a/polling_main.py +++ b/polling_main.py @@ -6,6 +6,7 @@ import logging from flask import Response from storage.dynamodb_storage import DynamoDBStorage +from storage.in_memory_storage import InMemoryStorage webserver = Flask(__name__) @@ -38,10 +39,10 @@ def run_flask_app(): # Initialize storage dynamodb_storage = DynamoDBStorage() - # in_memory_storage = InMemoryStorage() + in_memory_storage = InMemoryStorage() # TODO: Use DynamoDB storage for production, in-memory storage for testing - storage = dynamodb_storage + storage = in_memory_storage application = build_application(TOKEN, storage) diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py index 80dc6b8..8ed5832 100644 --- a/storage/dynamodb_storage.py +++ b/storage/dynamodb_storage.py @@ -1,7 +1,6 @@ import boto3 import logging from .storage_interface import BotState, Storage -from enum import Enum import os @@ -103,3 +102,47 @@ def get_user_id_from_channel_credentials(self, chat_id): 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: + table.delete_item(Key={"chat_id": str(chat_id)}) + except Exception as e: + logging.error(f"Error deleting item from DynamoDB: {e}") + raise + + def check_item_exists(self, table, chat_id): + try: + response = table.get_item(Key={"chat_id": str(chat_id)}) + logging.info(f"check_item_exists: {response}") + return "Item" in response + except Exception as e: + logging.error(f"Error checking item in DynamoDB: {e}") + raise diff --git a/storage/in_memory_storage.py b/storage/in_memory_storage.py index 80ea66b..fead0be 100644 --- a/storage/in_memory_storage.py +++ b/storage/in_memory_storage.py @@ -1,26 +1,76 @@ from .storage_interface import Storage - +import logging class InMemoryStorage(Storage): def __init__(self): self.state = {} self.playlists = {} self.users = {} + self.tokens = {} def save_current_state(self, chat_id, state_key): self.state[chat_id] = state_key + logging.info(f"Saved state for chat_id {chat_id}: {state_key}") def get_current_state(self, chat_id): - return self.state.get(chat_id, None) + state = self.state.get(chat_id, None) + logging.info(f"Retrieved state for chat_id {chat_id}: {state}") + return state def save_playlist_to_dynamodb(self, chat_id, playlist_id): self.playlists[chat_id] = playlist_id + logging.info(f"Saved playlist for chat_id {chat_id}: {playlist_id}") def get_playlist_from_dynamodb(self, chat_id): - return self.playlists.get(chat_id, None) + playlist = self.playlists.get(chat_id, None) + logging.info(f"Retrieved playlist for chat_id {chat_id}: {playlist}") + return playlist def get_user_id_from_chat_id(self, chat_id): - return self.users.get(chat_id, None) + user_id = self.users.get(chat_id, None) + logging.info(f"Retrieved user_id from chat_id for chat_id {chat_id}: {user_id}") + return user_id def get_user_id_from_channel_credentials(self, chat_id): - return self.users.get(chat_id, None) + token_info = self.tokens.get(chat_id, None) + user_id = token_info['user_id'] if token_info else None + logging.info(f"get_user_id_from_channel_credentials: Retrieved user_id {user_id} for chat_id: {chat_id}") + return user_id + + def get_cached_token(self, chat_id): + token = self.tokens.get(chat_id, None) + logging.info(f"Retrieved token for chat_id {chat_id}: {token}") + return token + + def save_token_to_cache(self, chat_id, user_id, token_info): + self.tokens[chat_id] = {'user_id': user_id, **token_info} + self.users[chat_id] = user_id # Save user ID to users dictionary + logging.info(f"Saved token for chat_id {chat_id}: {self.tokens[chat_id]}") + logging.info(f"Saved user_id for chat_id {chat_id}: {user_id}") + + def delete_item(self, table, chat_id): + if table == "credentials_table": + if chat_id in self.tokens: + del self.tokens[chat_id] + logging.info(f"Deleted token for chat_id {chat_id} from credentials_table") + elif table == "bot_table": + if chat_id in self.state: + del self.state[chat_id] + logging.info(f"Deleted state for chat_id {chat_id} from bot_table") + if chat_id in self.playlists: + del self.playlists[chat_id] + logging.info(f"Deleted playlist for chat_id {chat_id} from bot_table") + if chat_id in self.users: + del self.users[chat_id] + logging.info(f"Deleted user for chat_id {chat_id} from bot_table") + + def check_item_exists(self, table, chat_id): + if table == "credentials_table": + exists = chat_id in self.tokens + logging.info(f"Checked if token exists for chat_id {chat_id} in credentials_table: {exists}") + return exists + elif table == "bot_table": + exists = chat_id in self.state or chat_id in self.playlists or chat_id in self.users + logging.info(f"Checked if item exists for chat_id {chat_id} in bot_table: {exists}") + return exists + return False \ No newline at end of file From 9449c24ee1c89a2c32f0ed579205f6ef2548b679 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Wed, 22 May 2024 16:49:55 -0700 Subject: [PATCH 5/8] ADD: Changes to In Memory Storage --- storage/in_memory_storage.py | 144 ++++++++++++++++++++++------------- 1 file changed, 92 insertions(+), 52 deletions(-) diff --git a/storage/in_memory_storage.py b/storage/in_memory_storage.py index fead0be..22d5091 100644 --- a/storage/in_memory_storage.py +++ b/storage/in_memory_storage.py @@ -1,76 +1,116 @@ -from .storage_interface import Storage import logging +from .storage_interface import BotState, Storage + class InMemoryStorage(Storage): def __init__(self): - self.state = {} - self.playlists = {} - self.users = {} - self.tokens = {} + 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) - def save_current_state(self, chat_id, state_key): - self.state[chat_id] = state_key - logging.info(f"Saved state for chat_id {chat_id}: {state_key}") + 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): - state = self.state.get(chat_id, None) - logging.info(f"Retrieved state for chat_id {chat_id}: {state}") - return state + 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): - self.playlists[chat_id] = playlist_id - logging.info(f"Saved playlist for chat_id {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): - playlist = self.playlists.get(chat_id, None) - logging.info(f"Retrieved playlist for chat_id {chat_id}: {playlist}") - return playlist + 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): - user_id = self.users.get(chat_id, None) - logging.info(f"Retrieved user_id from chat_id for chat_id {chat_id}: {user_id}") - return user_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): - token_info = self.tokens.get(chat_id, None) - user_id = token_info['user_id'] if token_info else None - logging.info(f"get_user_id_from_channel_credentials: Retrieved user_id {user_id} for chat_id: {chat_id}") - return user_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): - token = self.tokens.get(chat_id, None) - logging.info(f"Retrieved token for chat_id {chat_id}: {token}") - return token + 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): - self.tokens[chat_id] = {'user_id': user_id, **token_info} - self.users[chat_id] = user_id # Save user ID to users dictionary - logging.info(f"Saved token for chat_id {chat_id}: {self.tokens[chat_id]}") - logging.info(f"Saved user_id for chat_id {chat_id}: {user_id}") + 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): - if table == "credentials_table": - if chat_id in self.tokens: - del self.tokens[chat_id] - logging.info(f"Deleted token for chat_id {chat_id} from credentials_table") - elif table == "bot_table": - if chat_id in self.state: - del self.state[chat_id] - logging.info(f"Deleted state for chat_id {chat_id} from bot_table") - if chat_id in self.playlists: - del self.playlists[chat_id] - logging.info(f"Deleted playlist for chat_id {chat_id} from bot_table") - if chat_id in self.users: - del self.users[chat_id] - logging.info(f"Deleted user for chat_id {chat_id} from bot_table") + 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): - if table == "credentials_table": - exists = chat_id in self.tokens - logging.info(f"Checked if token exists for chat_id {chat_id} in credentials_table: {exists}") - return exists - elif table == "bot_table": - exists = chat_id in self.state or chat_id in self.playlists or chat_id in self.users - logging.info(f"Checked if item exists for chat_id {chat_id} in bot_table: {exists}") + try: + exists = str(chat_id) in table + logging.info(f"check_item_exists: {exists}") return exists - return False \ No newline at end of file + except Exception as e: + logging.error(f"Error checking item in InMemoryStorage: {e}") + raise From 3b890b373a3de98a8d95900298884a5805088359 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Wed, 22 May 2024 22:44:48 -0700 Subject: [PATCH 6/8] ADD: Updated Bot Test --- test_telegram_bot.py | 252 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 205 insertions(+), 47 deletions(-) diff --git a/test_telegram_bot.py b/test_telegram_bot.py index 9de3bbb..2eebfcd 100644 --- a/test_telegram_bot.py +++ b/test_telegram_bot.py @@ -1,63 +1,221 @@ +import json import unittest +from datetime import datetime +from typing import Optional, Tuple from unittest.mock import AsyncMock, MagicMock -from telegram import Update, Message, User, Chat -from telegram.ext import CallbackContext -from bot import start, help_command +from telegram import Update, Message, User, Chat, MessageEntity +from telegram._utils.defaultvalue import DEFAULT_NONE +from telegram._utils.types import ODVInput +from telegram.ext import CallbackContext, Application, ContextTypes +from telegram.request import BaseRequest, RequestData -class TestStartCommand(unittest.IsolatedAsyncioTestCase): - async def test_start_command(self): - update = MagicMock(spec=Update) +import bot +from bot import start, help_command, register_handlers +from custom_callback_contex import CustomCallbackContext +from storage.in_memory_storage import InMemoryStorage - update.message = MagicMock(spec=Message) - update.message.chat_id = 12345 - update.message.reply_text = AsyncMock() - context = MagicMock(spec=CallbackContext) +def counter(initial=1): + """counter creates a Python generator yielding incrementing integers.""" + i = initial + while True: + i += 1 + yield i - await start(update, context) - update.message.reply_text.assert_awaited_once_with( - "Hiya! I'm your Spotify Skunk bot 🦨. /createplaylist to add songs to your playlist!" +class UnexpectedTelegramEndpoint(Exception): + pass + + +class UnexpectedTelegramApiCall(Exception): + pass + + +class FakeTelegramApiServer(BaseRequest): + """FakeTelegramApiServerRequest implements a fake Telegram API server by overriding the default networking + implementation with one that operates entirely in-memory. + You can extract the outbound requests by inspecting the .messages_sent list. + """ + + fake_api_endpoint = "https://localhost/" + fake_token = "faketoken" + + def __init__(self, bot_user: User): + """Initialize a FakeTelegramApiServer. + :arg bot_user Describes the identity of the bot, as would be represented by the Telegram API server. + """ + super().__init__() + self.bot_user = bot_user + # Collected outbound requests will be stored here. + self.messages_sent = [] + + async def initialize(self) -> None: + pass + + async def shutdown(self) -> None: + pass + + async def do_request( + self, + url: str, + method: str, + request_data: Optional[RequestData] = None, + read_timeout: ODVInput[float] = DEFAULT_NONE, + write_timeout: ODVInput[float] = DEFAULT_NONE, + connect_timeout: ODVInput[float] = DEFAULT_NONE, + pool_timeout: ODVInput[float] = DEFAULT_NONE, + ) -> Tuple[int, bytes]: + """Implements a fake networking request/response.""" + expected_prefix = self.fake_api_endpoint + self.fake_token + "/" + if not url.startswith(expected_prefix): + raise UnexpectedTelegramEndpoint( + f"Expected prefix of {expected_prefix} but got URL {url}" + ) + path = url[len(expected_prefix) :] + + # AFAICT the library is only sending sendMessage and getMe requests, so these are the only two backend + # methods we fake. + if path == "getMe": + return 200, json.dumps( + {"ok": True, "result": json.loads(self.bot_user.to_json())} + ).encode("utf-8") + elif path == "sendMessage": + print(f"Sending message: {request_data.parameters}") + self.messages_sent.extend([request_data.parameters]) + # Fake a response from the Telegram servers. Note: This is not the full response that it sends to a + # message. The real response would contain the full message that was sent. + return 200, json.dumps({"ok": True, "result": {}}).encode("utf-8") + + # If the bot tries to make a Telegram API call to a different endpoint, just raise an exception and fail the + # test. + raise UnexpectedTelegramApiCall( + f"_request_wrapper(url={path}, method={method}, request_data={request_data.json_payload}" ) -class TestHelpCommand(unittest.IsolatedAsyncioTestCase): - async def setUp(self): - chat = Chat(id=12345, type="public") - user = User(id=67890, is_bot=False, first_name="Test User") - message = Message(message_id=1, date=1609459200, chat=chat, from_user=user) - self.update = Update(update_id=1, message=message) - - async def test_help_command(self): - context = CallbackContext.from_update(self.update, bot=None) - context = CallbackContext.from_update(self.update, bot=None) - context.bot = MagicMock() - context.bot_data = {} - context.user_data = {} - context.chat_data = {} - context.match = None - - self.update.message.reply_text = AsyncMock() - - # Invoke the help command - await help_command(self.update, context) - - # Check the response - expected_text = ( - "Here are the commands you can use:\n" - "/start - Start interacting with the bot\n" - "/createplaylist - Create a new Spotify playlist\n" - "/changeplaylistname - Change the name of the current playlist\n" - "/changeplaylistimage - Change the image of the current playlist\n" - "/resetplaylist - Reset so you can create a new playlist\n" - "/unlink - Unlink your Spotify credentials\n" - "/playlistlink - Get the link to the current playlist\n" - "/help - Show this help message\n" - "\nJust send me a Spotify track link to add it to your playlist!" +class TestBotWithoutMocks(unittest.IsolatedAsyncioTestCase): + """TestBotWithoutMocks allows for testing the responses generated when sending commands to the bot.""" + + # Utility values to generate incrementing values for the messages we generate. + __message_id_counter = counter() + __update_id_counter = counter(10000) + + async def asyncSetUp(self): + super().setUp() + # Note: it would be nice to test group chats AND 1:1 chats but for now we are just testing group chats. + self.chat = Chat(id=-1, type="public") + self.bot_user = User( + id=999, is_bot=False, first_name="Test", username="testskunkbot" + ) + self.request_collector = FakeTelegramApiServer(self.bot_user) + + # This must match the behavior in bot.build_application /except/ for the fields that are commented. + # TODO: unify with build_application. + self.app = ( + Application.builder() + .base_url( + self.request_collector.fake_api_endpoint + ) # disables interaction with telegram servers + .context_types(ContextTypes(context=CustomCallbackContext)) + .defaults(bot.defaults) + .get_updates_request( + self.request_collector + ) # replaces network I/O with FakeTelegramApiServerRequest + .request( + self.request_collector + ) # replaces network I/O with FakeTelegramApiServerRequest + .token(self.request_collector.fake_token) + .updater(None) # disables the library's polling or webhook behaviors + .build() ) - self.update.message.reply_text.assert_awaited_once_with(expected_text) + # Use an in-memory database rather than dynamodb + self.storage = InMemoryStorage() + self.app.bot_data["storage"] = self.storage + + register_handlers(self.app) + + await self.app.initialize() + + async def asyncTearDown(self): + await self.app.shutdown() + + async def sendMessage(self, text): + """sendMessage sends a message to the bot. + If the text starts with /, we assume it is a "Command" and send a slightly more special type of message to + deal with that. Otherwise, it is just a plain text message. + Returns the last message we received from the bot, or None if no messages were generated by the bot. + """ + + # Bot commands need to be annotated with entities=. + if text.startswith("/"): + message = Message( + chat=self.chat, + date=datetime(year=2024, month=1, day=1), + entities=[ + MessageEntity(type="bot_command", offset=0, length=len(text)) + ], + from_user=self.bot_user, + message_id=next(self.__message_id_counter), + text=text, + ) + else: + # Plain text commands do not need entities=. + message = Message( + chat=self.chat, + date=datetime(year=2024, month=1, day=1), + from_user=self.bot_user, + message_id=next(self.__message_id_counter), + text=text, + ) + + # `update` is the Update message that we would expect to receive from the Telegram API servers. + update = Update(update_id=next(self.__update_id_counter), message=message) + + # This app.process_update will trigger the appropriate command based on the user's message. Once it is + # complete, we can inspect the request_collector to determine which API calls were sent to the Telegram API + # servers. + await self.app.process_update( + Update.de_json(data=update.to_dict(), bot=self.app.bot) + ) + + if not self.request_collector.messages_sent: + return None + return self.request_collector.messages_sent.pop() + + async def test_help(self): + response = await self.sendMessage("/help") + self.assertTrue(response["text"].startswith("Here are the commands")) + + async def test_irrelevant_message(self): + response = await self.sendMessage("bark bark bark") + # This message is ignored, so there is no response. + self.assertIsNone(response) + + async def test_unlink_before_link(self): + response = await self.sendMessage("/unlink") + # TODO: Fix this behavior -- /unlink on an unlinked channel should not reply with something that looks like + # an error message. + self.assertEqual( + "Failed to unlink your Spotify credentials: credentials not found.", + response["text"], + ) + # Confirm that there are no credentials for this channel. + self.assertIsNone( + self.storage.get_user_id_from_channel_credentials(self.chat.id) + ) + # Confirm that there is no state for this channel. + self.assertIsNone(self.storage.get_current_state(self.chat.id)) + + async def test_playlistlink(self): + response = await self.sendMessage("/playlistlink") + self.assertEqual("No playlist found for this chat.", response["text"]) + + async def test_link(self): + response = await self.sendMessage("/createplaylist") + # TODO: verify the response contains a spotify URL + # self.assertTrue(response["text"] contains "http...") if __name__ == "__main__": unittest.main() From 83be35605852527c4ee538a41e2da6b71b9cd095 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Wed, 22 May 2024 23:26:35 -0700 Subject: [PATCH 7/8] ADD: Updated Lambda and deploy --- bot.py | 12 +++--------- deploy_lambda.sh | 2 +- lambda_main.py | 7 +++++-- polling_main.py | 8 ++------ requirements.txt | 4 ++-- 5 files changed, 13 insertions(+), 20 deletions(-) diff --git a/bot.py b/bot.py index 214ed7c..49028e7 100644 --- a/bot.py +++ b/bot.py @@ -409,9 +409,7 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> chat_id = update.effective_chat.id try: # Check if the item exists in the credentials table - if not context.storage().check_item_exists( - "credentials_table", chat_id - ): + if not context.storage().check_item_exists("credentials_table", chat_id): logging.error(f"Credentials not found for chat_id {chat_id}") await update.message.reply_text( "Failed to unlink your Spotify credentials: credentials not found." @@ -420,9 +418,7 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> # Check if the item exists in the bot table - if not context.storage().check_item_exists( - "bot_table", chat_id - ): + if not context.storage().check_item_exists("bot_table", chat_id): logging.error(f"Bot data not found for chat_id {chat_id}") await update.message.reply_text( "Failed to unlink your Spotify credentials: bot data not found." @@ -439,9 +435,7 @@ async def unlink_credentials(update: Update, context: CustomCallbackContext) -> logging.info(f"Credentials delete_item response: {credentials_delete_response}") logging.info(f"Attempting to delete item from bot table for chat_id {chat_id}") - bot_delete_response = context.storage().delete_item( - "bot_table", chat_id - ) + bot_delete_response = context.storage().delete_item("bot_table", chat_id) logging.info(f"Bot delete_item response: {bot_delete_response}") await update.message.reply_text( diff --git a/deploy_lambda.sh b/deploy_lambda.sh index 7d09481..ad6ce44 100755 --- a/deploy_lambda.sh +++ b/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 diff --git a/lambda_main.py b/lambda_main.py index 80fc795..c08bd21 100644 --- a/lambda_main.py +++ b/lambda_main.py @@ -7,6 +7,7 @@ import json import traceback from telegram import Update +from storage.dynamodb_storage import DynamoDBStorage if logging.getLogger().hasHandlers(): logging.getLogger().setLevel(logging.INFO) @@ -14,6 +15,7 @@ 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): diff --git a/polling_main.py b/polling_main.py index 8140ab7..8bdd052 100644 --- a/polling_main.py +++ b/polling_main.py @@ -6,7 +6,7 @@ import logging from flask import Response from storage.dynamodb_storage import DynamoDBStorage -from storage.in_memory_storage import InMemoryStorage + webserver = Flask(__name__) @@ -38,11 +38,7 @@ def run_flask_app(): TOKEN = os.getenv("TELEGRAM_BOT_TOKEN") # Initialize storage - dynamodb_storage = DynamoDBStorage() - in_memory_storage = InMemoryStorage() - - # TODO: Use DynamoDB storage for production, in-memory storage for testing - storage = in_memory_storage + storage = DynamoDBStorage() application = build_application(TOKEN, storage) diff --git a/requirements.txt b/requirements.txt index 1d69188..568a6d7 100644 --- a/requirements.txt +++ b/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==21.1.1 +python-telegram-bot>=20.8 PyYAML==6.0.1 redis==5.0.1 requests==2.31.0 From cfcfda7b25d640e98519cfda8c71821fc8d45fe6 Mon Sep 17 00:00:00 2001 From: Luis Armendariz Date: Thu, 23 May 2024 10:41:50 -0700 Subject: [PATCH 8/8] ADD: Updated Lambda --- storage/dynamodb_storage.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/storage/dynamodb_storage.py b/storage/dynamodb_storage.py index 8ed5832..e710145 100644 --- a/storage/dynamodb_storage.py +++ b/storage/dynamodb_storage.py @@ -133,16 +133,26 @@ def save_token_to_cache(self, chat_id, user_id, token_info): def delete_item(self, table, chat_id): try: - table.delete_item(Key={"chat_id": str(chat_id)}) + 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}") - raise def check_item_exists(self, table, chat_id): try: - response = table.get_item(Key={"chat_id": str(chat_id)}) - logging.info(f"check_item_exists: {response}") - return "Item" in response + 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}") - raise + return False