Как я написал telegram бота на Rust для отслеживания рейтинга при поступлении в вуз / Хабр

Привет, меня зовут Илья, я сейчас сдаю вступительные экзамены в магистратуру. Столкнулся при поступлении с проблемой, что результаты экзаменов в рейтинговом списке появляются не сразу, а постоянно его открывать и находить себя на странице - после раза двадцатого надоело. После исследования devtools я захотел написать приложение для отслеживания изменений рейтинга, а уведомления отправлять в телеграм. А Rust был выбран по ��ростой причине - он мне понравился, ну и есть удобные штуки всякие.

Не судите строго, мой первый опыт написания статьи (и бота). Также она не претендует на звание полноценного туториала по разработке телеграм ботов на Rust, но я старался. И тем более это не туториал по самому языку.

Реализация

Для запросов использовал reqwest с включенной фичей json для десериализации, для поддержки async/await - tokio, для работы с SQLite - rusqlite.

Пока что есть возможность смотреть рейтинг только для магистратуры ИТМО

Схема работы

Получаем список программ и сохраняем в базе, далее запускаем цикл, на каждом шаге запрашиваем и обрабатываем обновления с телеграма, и примерно каждые 10 минут обновляем рейтинг с сайта.

Получение рейтинга

Начал я с получения данных о рейтинге с abitlk.itmo.ru. Здесь получаем данные, функция find_score ищет в них нужную запись, в случае если не находит, возвращает ошибку, которая будет обработана, и будет отправлено сообщение, что ничего не найдено

get_rating_competition

pub async fn get_rating_competition(
    degree: &str,
    program_id: &str,
    case_number: &str,
) -> Result<Option<Competition>, Box<dyn std::error::Error>> {
    let rating_response: RatingResponse = reqwest::get(format!(
        "{API_PREFIX}/{API_KEY}/rating/{degree}/budget?program_id={program_id}"
    ))
    .await?
    .json()
    .await?;

    match find_score(rating_response, case_number) {
        None => Err(Box::from("no matching competition")),
        competition => Ok(competition),
    }
}

Поиск записи в данных. Если в поле ok вернулось false, сразу возвращаем None. Иначе достаем массив всех записей из response.result.general_competition, и находим в нем позицию по номеру дела из личного кабинета (у меня вида Мх-хххх-2022)

find_score

fn find_score(response: RatingResponse, case_number: &str) -> Option<Competition> {
    if !response.ok {
        return None;
    }

    response
        .result
        .general_competition
        .iter()
        .find(|c| -> bool {
            if let Some(c) = &c.case_number {
                c == case_number
            } else {
                false
            }
        })
        .cloned()
}

Структура RatingResponse выглядит так. Тут я задал только нужные поля, остальные, которые придут в ответе, будут проигнорированы при парсинге.

Модели ответа

pub struct Response<T> {
    pub ok: bool,
    pub message: String,
    pub result: T,
}

pub type RatingResponse = Response<RatingResult>;

pub struct RatingResult {
    pub general_competition: Vec<Competition>,
}

pub struct Competition {
    pub position: i32,
    pub priority: i32,
    pub total_scores: f64,
    pub case_number: Option<String>,
    pub exam_scores: Option<f64>,
}

Также каждая структура наследует serde::Deserialize для генерации нужных методов, которые будут потом преобразовывать сырые данные (в данном случае json) в структуру:

#[derive(Deserialize)]
pub struct Response<T> {}

Получение команд от пользователей

Для получения сообщений используется эндпойнт https://api.telegram.org/botTOKEN/getUpdates?offset=x для получения обновлений, где есть и новые сообщения. offset необходим для пометки обновлений как прочитанных, чтобы они больше не возвращались. Функция для их получения:

get_updates

async fn get_updates(offset: i32) -> Result<GetUpdatesResponse, Box<dyn std::error::Error>> {
    let params = [("offset", &offset.to_string())];
    let url =
        reqwest::Url::parse_with_params(&format!("{TG_API_PREFIX}{TOKEN}/getUpdates"), &params)?;

    let response = reqwest::get(url).await?;
    if !response.status().is_success() {
        let error: ErrorResponse = response.json().await?;
        let text = format!(
            "cannot get updates\nerror: `{}`",
            error.description.unwrap_or_default()
        );
        send_message(&text, LOGS_CHAT_ID).await?;
        return Err(Box::from("cannot get updates"));
    }

    Ok(response.json().await?)
}

Полученные данные обрабатываем. Здесь во-первых сохраняется максимальный update_id для последующего использования в качестве offset, далее в MessageRequest::from(text) парсится текст сообщения, если это известная команда, она далее обрабатывается, если нет, отправляется сообщение, что это не команда.

handle_updates

pub async fn handle_updates(db: &DB, offset: i32) -> Result<i32, Box<dyn std::error::Error>> {
    let data = get_updates(offset).await?;

    let mut max_update_id = 0;
    for update in data.result {
        if update.update_id > max_update_id {
            max_update_id = update.update_id;
        }

        if let Some(message) = update.message {
            if let Some(text) = message.text {
                let chat_id = message.from.id.to_string();
                match MessageRequest::from(text) {
                    Some(request) => handle_message_request(db, request, &chat_id).await?,
                    None => send_message(messages::unknown_message, &chat_id).await?,
                }
            }
        }
    }

    Ok(max_update_id + 1)
}

Парсится так (функция from). Текст разбивается по пробелам, далее в случае односложных команд просто возвращается значение перечисления, а в случае /watch из значений массива создается объект для их передачи. При неправильной команде возвращается её имя, далее оно будет использовано для отправки синтаксиса этой команды.

MessageRequest::from

impl MessageRequest {
    pub fn from(text: String) -> Option<Self> {
        let text: Vec<String> = text.split(' ').map(|w| w.to_string()).collect();
        if text.is_empty() {
            return None;
        }

        let command = text[0].as_str();
        match command {
            "/watch" | "/unwatch" => {
                let incorrect_command = Some(Self::IncorrectCommand(command.to_string()));

                if text.len() < 5 {
                    if text.len() == 2 && text[1] == "all" {
                        return Some(Self::UnwatchAll);
                    }
                    return incorrect_command;
                }

                // waiting for let-chain
                if let Some(watch) = Watch::new("itmo", &text[2], &text[3], &text[4]) {
                    if watch.degree == Degree::Master {
                        return match command {
                            "/watch" => Some(Self::Watch(watch)),
                            "/unwatch" => Some(Self::Unwatch(watch)),
                            _ => incorrect_command,
                        };
                    }
                }

                incorrect_command
            }
            "/about" => Some(Self::About),
            "/help" => Some(Self::Help),
            "/start" => Some(Self::Start),
            _ => None,
        }
    }
}

Структура MessageRequest и вспомогательные Degree и Watch:

Модели запроса пользователя

pub enum Degree {
    Bachelor,
    Master,
    Postgraduate,
}

pub struct Watch {
    pub uni: String,
    pub degree: Degree,
    pub program_id: String,
    pub case_number: String,
}

pub enum MessageRequest {
    About,
    Help,
    Start,
    Unwatch(Watch),
    UnwatchAll,
    Watch(Watch),
    IncorrectCommand(String),
}

MessageRequest обрабатывается так. Для простых команд отправляем соответствующие сообщения, команды /watch и /unwatch обрабатываем (сама функция обработки чуть ниже).

handle_message_request

async fn handle_message_request(
    db: &DB,
    request: MessageRequest,
    chat_id: &str,
) -> Result<(), CrateError> {
    match request {
        MessageRequest::Watch(args) => {
            let result = handle_competition(db, chat_id, &args.degree.to_string(), &args.case_number, &args.program_id, true).await;
            if let Err(_) = result {
                send_message(messages::rating_not_found, chat_id).await?;
            }
        }
        MessageRequest::Unwatch(args) => {
            db.delete_competition(
                &args.case_number,
                chat_id,
                &args.program_id,
                &args.degree.to_string(),
            )?;
            send_message(messages::done, chat_id).await?;
        }
        MessageRequest::UnwatchAll => {
            db.delete_competition_by_user(chat_id)?;
            send_message(messages::done, chat_id).await?;
        }
        MessageRequest::IncorrectCommand(command) => {
            send_incorrect_command_message(&command, chat_id).await?
        }
        MessageRequest::Help => send_message(messages::help, chat_id).await?,
        MessageRequest::Start => send_message(messages::start, chat_id).await?,
        MessageRequest::About => send_message(messages::about, chat_id).await?,
    };
    Ok(())
}

Код обработки довольно скучный, но довольно большой. Сначала получается позиция, достается уже ранее запрошенная позиция из базы, если они не совпадают, то отсылается сообщение пользователю, и значение в базе обновляется. Если не равны, то сообщение отсылается только в случае, если эта функция была вызвана после обработки команды от пользователя, а не на этапе регулярной проверки обновлений. За это отвечает параметр is_user_request:

handle_competition

pub async fn handle_competition(
    db: &DB,
    chat_id: &str,
    degree: &str,
    case_number: &str,
    program_id: &str,
    is_user_request: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    let competition = get_rating_competition(db, degree, program_id, case_number).await?;

    match db.select_competition(chat_id, case_number, degree, program_id) {
        Ok(old_competition) => {
            if let Some(competition) = competition {
                let program = db.select_program("itmo", program_id)?;
                let program_name = if let Some(program) = program {
                    program.title_ru
                } else {
                    "Названия нет".to_string()
                };

                let mut should_send_message = false;

                // update if competition is old (competition != old_competition)
                // insert if is new (when old == None, on first user request)
                if let Some(old_competition) = old_competition {
                    if competition != old_competition {
                        db.update_competition(&competition, chat_id, program_id, degree)?;
                        should_send_message = true;
                    }
                } else {
                    db.insert_competition(&competition, chat_id, program_id, degree)?;
                }

                // send if it's user request or record in db was updated
                if is_user_request || should_send_message {
                    send_competition_message(&competition, chat_id, &program_name).await?;
                }
            }
        }
        Err(e) => {
            eprintln!("cannot select competition: {e}")
        }
    };
    Ok(())
}

Отправка сообщения выглядит так. Экранируем некоторые символы, отправляем запрос, проверяем результат.

send_message

pub async fn send_message(text: &str, chat_id: &str) -> Result<(), Box<dyn std::error::Error>> {
    let text = &text.replace('-', "\\-").replace('.', "\\.");

    let params = [
        ("chat_id", chat_id),
        ("text", text),
        ("parse_mode", "MarkdownV2"),
    ];

    let url =
        reqwest::Url::parse_with_params(&format!("{TG_API_PREFIX}{TOKEN}/sendMessage"), &params)?;
    let response = reqwest::get(url).await?;
    if !response.status().is_success() {
        let error: ErrorResponse = response.json().await?;
        let msg = "Cannot send message request";
        if let Some(description) = error.description {
            eprintln!("{msg}: {description}");
        }
        return Err(Box::from(msg));
    }

    let data: SendMessageResponse = response.json().await?;

    if !data.ok {
        eprintln!(
            "Cannot send message: {}",
            data.description
                .unwrap_or_else(|| "error has no description".to_string())
        )
    }

    Ok(())
}

Объединяем

В функции main сначала запрашивается список всех программ (для сохранения их названий), далее запускается бесконечный цикл, который каждую секунду проверяет обновления со стороны телеграма, вызывая функцию handle_updates, а также раз в примерно 10 минут проверяет обновления рейтингов и рассылает сообщения об обновлении. #[tokio::main] нужно для старта main сразу как асинхронной функции, без необходимости ручного создания loop-а.

main

const TEN_MIN_IN_SEC: i32 = 10 * 60;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let db = init_db()?;

    load_programs(&db).await.unwrap();

    let mut offset = 0;
    let mut sec_counter = 0;
    loop {
        offset = handle_updates(&db, offset).await.unwrap();

        if sec_counter == 0 {
            check_rating_updates(&db).await?;
        }
        sec_counter = (sec_counter + 1) % TEN_MIN_IN_SEC;

        time::sleep(time::Duration::from_secs(1)).await;
    }
}

check_rating_updates

async fn check_rating_updates(db: &DB) -> Result<(), Box<dyn std::error::Error>> {
    // select registered watchers from 'results'
    for c in db.select_all_competitions()? {
        if let Some(case_number) = c.competition.case_number {
            handle_competition(db, &c.tg_chat_id, &c.degree, &case_number, &c.program_id, false)
            .await?;
        }
    }
    Ok(())
}

Итоги

Надеюсь, этот бот поможет кому-нибудь ещё кроме меня. Ссылка на бота @uni_rating_checker_bot и на исходный код

P.S.

Хотелось бы проверить, как приложение справится под нагрузкой (сервер довольно слабенький)
Из возможных фич - учитывать, что кто-то из тех, кто находится выше по рейтингу, уже проходят на направление, которое у них выше по приоритету
После написания основной части статьи я переписал часть проекта, улучшив передачу ошибок наверх и их обработку
Только магистратура потому что для бакалавриата вроде как приходит другая структура данных, а времени особо не было на это
По хорошему, ввод данных надо делать не передачей нескольких аргументов в одной строке, а чтобы пользователь мог ввести все по очереди, в идеале - чтобы надо было просто выбирать из предложенных вариантов, например, с помощью кнопок
Вдохновился этой статьей https://habr.com/p/679832
Команда потестить: /watch itmo master 15850 M1-0979-2022. Это не я :) Выбрал наугад