From c18b2aee114a3bdbaebae208326f1baa7da8bce9 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 6 Feb 2024 22:32:16 +0400 Subject: [PATCH] playrooms and games recorded to db --- src/callbacks/receive_poll_answer.py | 14 ++++++++++++ src/data_models/Game.py | 22 +++++++++++++++---- src/data_models/Player.py | 15 +++++++++---- src/handlers/save.py | 12 ++++++++-- src/services/db_service.py | 33 ++++++++++++++++++++++++++++ src/sql/init.sql | 8 ++++--- 6 files changed, 91 insertions(+), 13 deletions(-) diff --git a/src/callbacks/receive_poll_answer.py b/src/callbacks/receive_poll_answer.py index dd2fc88..37d1668 100644 --- a/src/callbacks/receive_poll_answer.py +++ b/src/callbacks/receive_poll_answer.py @@ -3,6 +3,9 @@ from telegram import Update from telegram.ext import ContextTypes +from src.data_models.Player import Player +from src.services.db_service import save_player + async def receive_poll_answer( update: Update, context: ContextTypes.DEFAULT_TYPE @@ -18,6 +21,17 @@ async def receive_poll_answer( user_id = update.effective_user.id result = answered_poll["questions"][answer.option_ids[0]] context.bot_data[answer.poll_id]["results"][user_id] = result + await save_player( + Player( + telegram_user_id=user_id, + username=update.effective_user.username, + first_name=update.effective_user.first_name, + full_name=update.effective_user.full_name, + last_name=update.effective_user.last_name, + is_bot=update.effective_user.is_bot, + language_code=update.effective_user.language_code, + ) + ) else: logging.error( "Failed to save poll answer. Usually happens for polls that are sent to the bot before it " diff --git a/src/data_models/Game.py b/src/data_models/Game.py index af4cfff..7540e4b 100644 --- a/src/data_models/Game.py +++ b/src/data_models/Game.py @@ -1,10 +1,24 @@ from datetime import datetime from typing import Literal -from pydantic import BaseModel +from pydantic import BaseModel, field_validator class Game(BaseModel): - playroom_id: int - end_time: datetime - result: Literal["Hitler Canceler", "Fascist Law", "Hitler Death", "Liberal Law"] + poll_id: int + chat_id: int + results: dict # Literal["Hitler Canceler", "Fascist Law", "Hitler Death", "Liberal Law"] + creator_id: int + + @field_validator("results", mode="after") + @classmethod + def validate_results(cls, v: dict) -> Literal["CH", "DH", "FW", "LW"]: + outcomes = set(v.values()) + if "I'm Canceler Hitler" in outcomes: + return "CH" + if "I'm Dead Hitler" in outcomes: + return "DH" + if "I'm Liberal Winner" in outcomes: + return "LW" + if "I'm Fascistic Winner" in outcomes: + return "FW" diff --git a/src/data_models/Player.py b/src/data_models/Player.py index 332c9a2..0a6cd65 100644 --- a/src/data_models/Player.py +++ b/src/data_models/Player.py @@ -1,11 +1,18 @@ from typing import Optional -from pydantic import BaseModel +from pydantic import BaseModel, field_validator class Player(BaseModel): telegram_user_id: int username: str - first_name: Optional[str] = None - last_name: Optional[str] = None - is_bot: bool = False + first_name: Optional[str | None] + full_name: Optional[str | None] + last_name: Optional[str | None] + is_bot: Optional[bool] = False + language_code: Optional[str | None] + + @field_validator("is_bot", mode="after") + @classmethod + def validate_bot(cls, v: bool) -> str: + return "TRUE" if v else "FALSE" # sqlite3 does not support boolean type diff --git a/src/handlers/save.py b/src/handlers/save.py index 38987fc..c67db4d 100644 --- a/src/handlers/save.py +++ b/src/handlers/save.py @@ -3,10 +3,11 @@ from telegram import Update from telegram.ext import ContextTypes +from src.data_models.Game import Game from src.data_models.Record import Record from src.utils import message_is_poll, is_message_from_group_chat from src import db -from src.services.db_service import save_record +from src.services.db_service import save_record, save_game async def _pass_checks( @@ -88,7 +89,14 @@ async def save(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: for player_id, result in poll_data["results"].items() ] ) - + await save_game( + Game( + poll_id=poll_data["message_id"], + chat_id=poll_data["chat_id"], + creator_id=poll_data["creator_id"], + results=poll_data["results"].copy(), + ) + ) else: await update.effective_message.reply_text( "Something went wrong. Can't process your request." diff --git a/src/services/db_service.py b/src/services/db_service.py index 24d459c..e35fa18 100644 --- a/src/services/db_service.py +++ b/src/services/db_service.py @@ -1,6 +1,8 @@ import logging import sqlite3 +from src.data_models.Game import Game +from src.data_models.Player import Player from src.data_models.Playroom import Playroom from src.data_models.Record import Record from src.db import execute @@ -29,3 +31,34 @@ async def save_playroom(playroom: Playroom) -> None: ) except sqlite3.IntegrityError: logging.info(f"Playroom {playroom.name} already exists in the database") + + +async def save_player(player: Player) -> None: + """Add a player to the bot_data""" + try: + await execute( + "INSERT INTO players (id, username, first_name, full_name, last_name, is_bot, language_code) " + "VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + player.telegram_user_id, + player.username, + player.first_name, + player.full_name, + player.last_name, + player.is_bot, + player.language_code, + ), + ) + except sqlite3.IntegrityError: + logging.info(f"Player {player.username} already exists in the database") + + +async def save_game(game: Game) -> None: + """Add a game to the bot_data""" + try: + await execute( + "INSERT INTO games (id, playroom_id, creator_id, result) VALUES (?, ?, ?, ?)", + (game.poll_id, game.chat_id, game.creator_id, game.results), + ) + except sqlite3.IntegrityError: + logging.info(f"Game {game.poll_id} already exists in the database") diff --git a/src/sql/init.sql b/src/sql/init.sql index 37fd704..d088c52 100644 --- a/src/sql/init.sql +++ b/src/sql/init.sql @@ -3,7 +3,10 @@ CREATE TABLE IF NOT EXISTS players ( id INTEGER PRIMARY KEY NOT NULL UNIQUE, -- Telegram user id username TEXT NOT NULL, -- Telegram username first_name TEXT, -- Telegram first name - last_name TEXT -- Telegram last name + full_name TEXT, -- Telegram full name + last_name TEXT, -- Telegram last name + is_bot TEXT NOT NULL DEFAULT 'FALSE', -- Telegram is_bot + language_code TEXT -- Telegram language code ); -- Create table for playrooms @@ -27,8 +30,7 @@ CREATE TABLE IF NOT EXISTS games ( id INTEGER PRIMARY KEY NOT NULL UNIQUE, -- Telegram poll id playroom_id INTEGER, -- Telegram chat id creator_id INTEGER NOT NULL, -- Telegram user id who created the game - start_time DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, - end_time DATETIME, + time DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, result TEXT, -- ["Hitler Canceler", "Fascist Law", "Hitler Death", "Liberal Law"] FOREIGN KEY (creator_id) REFERENCES players(id), FOREIGN KEY (playroom_id) REFERENCES playrooms(id)