lambda_function.py

import asyncio
import json
import logging

from aiogram import types

from core import BOT_TOKEN, LAMBDA_URL, SECRET_KEY, bot, dp
from utils import setup_logging

# Configure logging
setup_logging(BOT_TOKEN)


async def set_webhook():
    """Sets the webhook for the bot."""
    if not LAMBDA_URL:
        logging.error("LAMBDA_URL is not set in environment variables.")
        return
    await bot.set_webhook(LAMBDA_URL)
    logging.info(f"Webhook set to {LAMBDA_URL}")


async def delete_webhook():
    """Deletes the webhook for the bot."""
    await bot.delete_webhook()
    logging.info("Webhook deleted.")


async def process_update(update: dict):
    """Process a single Telegram update."""
    telegram_update = types.Update(**update)
    await dp.feed_update(bot=bot, update=telegram_update)


async def main(event):
    """
    Asynchronous main function to handle the Lambda event.
    This allows for proper session management.
    """
    # The AWS Console 'Test' button sends the payload directly in the event.
    # API Gateway sends it in a 'body' string.
    body = json.loads(event.get('body', '{}')) if isinstance(event.get('body'), str) else event

    # Check for a management command
    if 'command' in body and body.get('secret') == SECRET_KEY:
        command = body['command']
        logging.info(f"Received management command: {command}")
        if command == "set_webhook":
            await set_webhook()
            return {'statusCode': 200, 'body': json.dumps({'status': 'Webhook set successfully'})}
        elif command == "delete_webhook":
            await delete_webhook()
            return {'statusCode': 200, 'body': json.dumps({'status': 'Webhook deleted successfully'})}
        else:
            return {'statusCode': 400, 'body': json.dumps({'error': 'Unknown command'})}

    # If not a management command, process as a Telegram update
    elif 'update_id' in body:
        await process_update(body)
        return {'statusCode': 200, 'body': json.dumps({'status': 'ok'})}

    else:
        logging.info("Received a request that was neither a command nor a Telegram update.")
        return {'statusCode': 200, 'body': json.dumps({'status': 'ok - no action taken'})}


def lambda_handler(event, context):
    """
    AWS Lambda handler that can process both incoming Telegram updates
    and special management commands.
    """
    logging.info(f"Received event: {event}")

    try:
        return asyncio.run(main(event))
    finally:
        # Ensure the bot session is closed after each invocation.
        # This is crucial to prevent 'Event loop is closed' errors.
        asyncio.run(bot.session.close())