diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/bot.rs | 45 | ||||
| -rw-r--r-- | src/config.rs | 18 | ||||
| -rw-r--r-- | src/database.rs | 15 | ||||
| -rw-r--r-- | src/main.rs | 42 | ||||
| -rw-r--r-- | src/user.rs | 10 | ||||
| -rw-r--r-- | src/web.rs | 42 |
6 files changed, 128 insertions, 44 deletions
@@ -3,6 +3,7 @@ use std::ops::Deref; use std::sync::Arc; use teloxide::{prelude2::*, types::MessageKind, utils::command::BotCommand}; +use tracing::{debug, error, info}; use crate::{database, DbPool, token}; use crate::user::User; @@ -18,36 +19,54 @@ enum Command { async fn answer(bot: AutoSend<Bot>, message: Message, command: Command, db: Arc<DbPool>) -> Result<(), Box<dyn Error + Send + Sync>> { + debug!("Answering telegram message: {:?}", message); match command { Command::Help => { - bot.send_message(message.chat.id, Command::descriptions()).await?; + debug!("Command /help."); + if let Err(why) = + bot.send_message(message.chat.id, Command::descriptions()).await { + error!("Failed to send telegram message: {:?}.", why); + } } Command::Start => { + debug!("Command /start."); if let MessageKind::Common(msg) = message.kind { if msg.from.is_none() { + debug!("Ignore message from channel."); return Ok(()); // ignore messages from channel } let from = msg.from.unwrap(); let db = db.deref(); let chat_id = message.chat.id; - match database::create_user(db, User { + let user = User { id: from.id as u64, name: from.username.unwrap_or_else(|| String::from("")), token: token::generate(), chat_id: chat_id as u64, - }).await { - Ok(_) => {} - Err(why) => println!("cannot create user: {:?}", why), + }; + if let Err(why) = database::create_user(db, &user).await { + if format!("{:?}", why).contains("UNIQUE constraint failed") { + info!("User exists: {}", user); + } else { + error!("Failed to create user {}: {:?}. Skip creating.", user, why); + } + } else { + info!("Created user {}.", user); } - bot.send_message( - chat_id, - match database::get_user_by_chat_id(db, chat_id as u64).await? { - Some(user) => - format!("Your token is `{}`. Treat it as a secret!", user.token), - _ => + let message = + match database::get_user_by_chat_id(db, chat_id as u64).await { + Ok(u) => match u { + Some(user) => format!("Your token is `{}`. Treat it as a secret!", user.token), + _ => String::from("Error: cannot fetch token.") + }, + Err(why) => { + error!("Cannot get user: {:?}.", why); String::from("Error: cannot fetch token.") - }, - ).await?; + } + }; + if let Err(why) = bot.send_message(chat_id, message).await { + error!("Failed to send telegram message: {:?}.", why); + } } } }; diff --git a/src/config.rs b/src/config.rs index c5d07a7..99e5f3d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,6 +2,7 @@ use std::fs::File; use std::io::Read; use serde_derive::Deserialize; +use tracing::error; #[derive(Deserialize)] pub struct Config { @@ -9,17 +10,26 @@ pub struct Config { pub log_file: String, pub db_file: String, pub listen: String, + pub log_level: String, } impl Config { // Read config file. Panic if any error occurs. pub fn from_file(file_path: &str) -> Config { let mut file = File::open(file_path) - .unwrap_or_else(|_| panic!("cannot open file {}", file_path)); + .unwrap_or_else(|err| { + error!("Cannot open config file {}: {:?}", file_path, err); + panic!(); + }); let mut config = String::new(); file.read_to_string(&mut config) - .unwrap_or_else(|_| panic!("cannot read config file {}", file_path)); - return serde_json::from_str(config.as_str()) - .expect("cannot decode config file in JSON"); + .unwrap_or_else(|err| { + error!("Cannot read config file {}: {:?}.", file_path, err); + panic!(); + }); + return serde_json::from_str(config.as_str()).unwrap_or_else(|err| { + error!("Cannot decode config file: {:?}.", err); + panic!(); + }); } }
\ No newline at end of file diff --git a/src/database.rs b/src/database.rs index edb2fc7..a73aafd 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,9 +1,9 @@ use std::str::FromStr; use futures::TryStreamExt; -use log::{debug, error, info, warn}; use sqlx::Row; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; +use tracing::debug; use crate::user::User; @@ -15,11 +15,13 @@ pub async fn open(file_path: &str) -> Result<DbPool, sqlx::Error> { let opt = SqliteConnectOptions::from_str(format!("sqlite://{}", file_path).as_str())? .create_if_missing(true); + debug!("Opening database pool..."); let pool = SqlitePoolOptions::new() .max_connections(SQLITE_THREAD_POOL_SIZE) .connect_with(opt).await?; // create table + debug!("Creating user table if not exist..."); sqlx::query(r#"CREATE TABLE IF NOT EXISTS "user" ( "id" INTEGER NOT NULL CHECK(id >= 0) UNIQUE, "name" TEXT, @@ -28,10 +30,12 @@ pub async fn open(file_path: &str) -> Result<DbPool, sqlx::Error> { PRIMARY KEY("id","token","chat_id") )"#).execute(&pool).await?; + debug!("Finish opening database."); Ok(pool) } pub async fn get_user(db: &DbPool, sql: &str, param: &str) -> Result<Option<User>, sqlx::Error> { + debug!(sql, param, "Database query."); let mut rows = sqlx::query(sql) .bind(param).fetch(db); @@ -52,20 +56,23 @@ pub async fn get_user(db: &DbPool, sql: &str, param: &str) -> Result<Option<User } pub async fn get_user_by_token(db: &DbPool, token: &str) -> Result<Option<User>, sqlx::Error> { + debug!(token, "Get user by token."); get_user(db, "SELECT id, name, token, chat_id FROM user WHERE token=$1 LIMIT 1", token).await } pub async fn get_user_by_chat_id(db: &DbPool, chat_id: u64) -> Result<Option<User>, sqlx::Error> { + debug!(chat_id, "Get user by chat_id."); get_user(db, "SELECT id, name, token, chat_id FROM user WHERE chat_id=$1 LIMIT 1", chat_id.to_string().as_str()).await } -pub async fn create_user(db: &DbPool, user: User) -> sqlx::Result<()> { +pub async fn create_user(db: &DbPool, user: &User) -> sqlx::Result<()> { + debug!("Create user: {}", user); sqlx::query("INSERT INTO user (id, name, token, chat_id) VALUES ($1,$2,$3,$4)") .bind(user.id.to_string()) - .bind(user.name) - .bind(user.token) + .bind(user.name.clone()) + .bind(user.token.clone()) .bind(user.chat_id.to_string()) .execute(db).await?; Ok(()) diff --git a/src/main.rs b/src/main.rs index ad5fdf3..f67a6bc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,13 @@ use std::convert::Infallible; use std::net::SocketAddr; +use std::ops::Deref; +use std::str::FromStr; use std::sync::Arc; -use log::{debug, error, info, warn}; use teloxide::prelude2::*; +use tracing::{debug, info, Level}; +use tracing::instrument; +use tracing_subscriber::FmtSubscriber; use warp::Filter; use config::Config; @@ -19,6 +23,7 @@ mod config; const CONFIG_FILE_NAME: &str = "kimikuri.json"; const MAX_BODY_LENGTH: u64 = 1024 * 16; +const DEFAULT_LOG_LEVEL: Level = Level::DEBUG; fn with_db(db_pool: DbPool) -> impl Filter<Extract=(DbPool, ), Error=Infallible> + Clone { warp::any().map(move || db_pool.clone()) @@ -29,13 +34,28 @@ fn with_bot(bot: Bot) -> impl Filter<Extract=(Bot, ), Error=Infallible> + Clone warp::any().map(move || bot.clone()) } +#[instrument] #[tokio::main] async fn main() { - pretty_env_logger::init(); - debug!("Loading bot config."); + eprintln!("Loading configuration file {}...", CONFIG_FILE_NAME); + // TODO make some fields optional let config = Config::from_file(CONFIG_FILE_NAME); - info!("Starting bot."); - let bot = Bot::new(config.bot_token); + + // configure logger + let log_level = match Level::from_str(&*config.log_level) { + Ok(l) => l, + Err(_) => { + eprintln!("Invalid log level: {}. Use {:?} instead.", + config.log_level, DEFAULT_LOG_LEVEL); + DEFAULT_LOG_LEVEL + } + }; + eprintln!("Configuration is loaded. Set log level to {:?}.", log_level); + let subscriber = FmtSubscriber::builder() + .with_max_level(log_level) + .finish(); + tracing::subscriber::set_global_default(subscriber) + .expect("Failed to set default subscriber"); let db = config.db_file.as_str(); info!(db, "Opening database..."); @@ -51,18 +71,18 @@ async fn main() { .and(with_db(db.deref().clone())) .and(with_bot(bot.clone())) .and_then(web::handler); + tokio::spawn(bot::repl(bot, db.clone())); - tokio::spawn(bot::repl(bot, Arc::new(db))); - + info!("Starting HTTP server..."); let endpoint: SocketAddr = config.listen.parse() .expect("Cannot parse `listen` as endpoint."); - - println!("Listen on {}", endpoint); - + info!("Start listening on {}", endpoint); tokio::spawn(warp::serve(send_message).run(endpoint)); + debug!("Waiting for Ctrl-C in main coroutine..."); tokio::signal::ctrl_c().await.unwrap(); - + // gracefully shutdown the database connection + info!("Closing database..."); db.deref().close().await; } diff --git a/src/user.rs b/src/user.rs index aef43fa..fb4adc7 100644 --- a/src/user.rs +++ b/src/user.rs @@ -1,6 +1,16 @@ +use std::fmt; +use std::fmt::Formatter; + pub struct User { pub id: u64, pub name: String, pub token: String, pub chat_id: u64, +} + +impl fmt::Display for User { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "User {{ id={}, name={}, token={}, chat_id={} }}", + self.id, self.name, self.token, self.chat_id) + } }
\ No newline at end of file @@ -1,5 +1,9 @@ +use std::fmt; +use std::fmt::Formatter; + use serde_derive::{Deserialize, Serialize}; use teloxide::{prelude2::*}; +use tracing::{debug, error, info, warn}; use warp::{Rejection, Reply}; use crate::{Bot, database, DbPool}; @@ -10,21 +14,33 @@ pub struct SendMessage { message: String, } +impl fmt::Display for SendMessage { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "SendMessage {{ token={}, message={} }}", self.token, self.message) + } +} + #[derive(Deserialize, Serialize)] pub struct SendMessageResponse { success: bool, message: String, } -pub async fn handler(req: SendMessage, db: DbPool, bot: Bot) -> std::result::Result<impl Reply, Rejection> { - println!("Token: {}, Message: {}", req.token, req.message); - let user = database::get_user_by_token(&db, req.token.as_str()).await; - Ok(warp::reply::json(&match user { +impl fmt::Display for SendMessageResponse { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "SendMessageResponse {{ success={}, message={} }}", self.success, self.message) + } +} + +pub async fn handler(req: SendMessage, db: DbPool, bot: Bot) + -> std::result::Result<impl Reply, Rejection> { + info!("Income API request: {}", req); + let user = + database::get_user_by_token(&db, req.token.as_str()).await; + let response = match user { Ok(u) => match u { Some(user) => { - log::info!("User: {} (id={}), message: {}", - user.name, user.id, req.message); - // TODO send message to Telegram + info!("Send message to user {}.", user); let bot = bot.auto_send(); match bot.send_message(user.chat_id as i64, req.message).await { Ok(_) => SendMessageResponse { @@ -32,7 +48,7 @@ pub async fn handler(req: SendMessage, db: DbPool, bot: Bot) -> std::result::Res message: String::new(), }, Err(why) => { - println!("Failed to send message to telegram: {:?}", why); + error!("Failed to send message to telegram: {:?}", why); SendMessageResponse { success: false, message: String::from("Failed to send message to telegram."), @@ -41,19 +57,21 @@ pub async fn handler(req: SendMessage, db: DbPool, bot: Bot) -> std::result::Res } } None => { - log::warn!("Invalid token {}, message: {}", req.token, req.message); + warn!("Invalid token: {}.", req); SendMessageResponse { success: false, message: String::from("Invalid token."), } } }, - Err(_) => { - log::error!("Error when querying the database."); + Err(err) => { + error!("Error when querying the database: {:?}.", err); SendMessageResponse { success: false, message: String::from("Invalid parameter."), } } - })) + }; + debug!("Response: {}", response); + Ok(warp::reply::json(&response)) } |