Annotation For src/main.rs
Logged in as anonymous

Lines of src/main.rs from check-in a68576cc1b that are changed by the sequence of edits moving toward check-in b2b8005309:

                         1: use std::collections::{BTreeMap, HashSet};
                         2: use std::sync::{Arc, Mutex};
                         3: 
                         4: use config;
                         5: 
                         6: use tokio;
                         7: 
                         8: use rss;
                         9: 
                        10: use chrono::DateTime;
                        11: 
                        12: use regex::Regex;
                        13: 
                        14: use telegram_bot::*;
                        15: use tokio::stream::StreamExt;
                        16: 
                        17: use sqlx::postgres::PgPoolOptions;
                        18: use sqlx::Row;
                        19: use sqlx::Done; // .rows_affected()
                        20: 
                        21: #[macro_use]
                        22: extern crate lazy_static;
                        23: 
                        24: use anyhow::{anyhow, bail, Context, Result};
                        25: 
                        26: #[derive(Clone)]
                        27: struct Core {
                        28: 	owner: i64,
                        29: 	api_key: String,
                        30: 	owner_chat: UserId,
                        31: 	tg: telegram_bot::Api,
                        32: 	my: User,
                        33: 	pool: sqlx::Pool<sqlx::Postgres>,
                        34: 	sources: Arc<Mutex<HashSet<Arc<i32>>>>,
                        35: }
                        36: 
                        37: impl Core {
                        38: 	async fn new(settings: config::Config) -> Result<Core> {
                        39: 		let owner = settings.get_int("owner")?;
                        40: 		let api_key = settings.get_str("api_key")?;
                        41: 		let tg = Api::new(&api_key);
                        42: 		let core = Core {
                        43: 			owner: owner,
                        44: 			api_key: api_key.clone(),
                        45: 			my: tg.send(telegram_bot::GetMe).await?,
                        46: 			tg: tg,
                        47: 			owner_chat: UserId::new(owner),
                        48: 			pool: PgPoolOptions::new()
                        49: 				.max_connections(5)
                        50: 				.connect_timeout(std::time::Duration::new(300, 0))
                        51: 				.idle_timeout(std::time::Duration::new(60, 0))
                        52: 				.connect_lazy(&settings.get_str("pg")?)?,
                        53: 			sources: Arc::new(Mutex::new(HashSet::new())),
                        54: 		};
                        55: 		let clone = core.clone();
                        56: 		tokio::spawn(async move {
                        57: 			if let Err(err) = &clone.autofetch().await {
a68576cc1b 2020-12-15   58: 				if let Err(err) = clone.debug(&format!("šŸ›‘ {:?}", err)) {
                        59: 					eprintln!("Autofetch error: {}", err);
                        60: 				};
                        61: 			}
                        62: 		});
                        63: 		Ok(core)
                        64: 	}
                        65: 
                        66: 	fn stream(&self) -> telegram_bot::UpdatesStream {
                        67: 		self.tg.stream()
                        68: 	}
                        69: 
a68576cc1b 2020-12-15   70: 	fn debug(&self, msg: &str) -> Result<()> {
a68576cc1b 2020-12-15   71: 		self.tg.spawn(SendMessage::new(self.owner_chat, msg));
                        72: 		Ok(())
                        73: 	}
                        74: 
                        75: 	async fn check<S>(&self, id: i32, owner: S, real: bool) -> Result<()>
                        76: 	where S: Into<i64> {
                        77: 		let owner: i64 = owner.into();
                        78: 		let id = {
                        79: 			let mut set = self.sources.lock().unwrap();
                        80: 			match set.get(&id) {
                        81: 				Some(id) => id.clone(),
                        82: 				None => {
                        83: 					let id = Arc::new(id);
                        84: 					set.insert(id.clone());
                        85: 					id.clone()
                        86: 				},
                        87: 			}
                        88: 		};
                        89: 		let count = Arc::strong_count(&id);
                        90: 		if count == 2 {
                        91: 			let mut conn = self.pool.acquire().await
                        92: 				.with_context(|| format!("Query queue fetch conn:\n{:?}", &self.pool))?;
                        93: 			let row = sqlx::query("select source_id, channel_id, url, iv_hash, owner from rsstg_source where source_id = $1 and owner = $2")
                        94: 				.bind(*id)
                        95: 				.bind(owner)
                        96: 				.fetch_one(&mut conn).await
                        97: 				.with_context(|| format!("Query source:\n{:?}", &self.pool))?;
                        98: 			drop(conn);
                        99: 			let channel_id: i64 = row.try_get("channel_id")?;
                       100: 			let destination = match real {
                       101: 				true => UserId::new(channel_id),
                       102: 				false => UserId::new(row.try_get("owner")?),
                       103: 			};
                       104: 			let url: &str = row.try_get("url")?;
                       105: 			let mut this_fetch: Option<DateTime<chrono::FixedOffset>> = None;
                       106: 			let iv_hash: Option<&str> = row.try_get("iv_hash")?;
                       107: 			let mut posts: BTreeMap<DateTime<chrono::FixedOffset>, String> = BTreeMap::new();
                       108: 			let feed = rss::Channel::from_url(url)
                       109: 				.with_context(|| format!("Problem opening feed url:\n{}", &url))?;
                       110: 			for item in feed.items() {
                       111: 				let date = match item.pub_date() {
                       112: 					Some(feed_date) => DateTime::parse_from_rfc2822(feed_date),
                       113: 					None => DateTime::parse_from_rfc3339(&item.dublin_core_ext().unwrap().dates()[0]),
                       114: 				}?;
                       115: 				let url = item.link().unwrap().to_string();
                       116: 				posts.insert(date.clone(), url.clone());
                       117: 			};
                       118: 			for (date, url) in posts.iter() {
                       119: 				let mut conn = self.pool.acquire().await
                       120: 					.with_context(|| format!("Check post fetch conn:\n{:?}", &self.pool))?;
                       121: 				let row = sqlx::query("select exists(select true from rsstg_post where url = $1 and source_id = $2) as exists;")
                       122: 					.bind(&url)
                       123: 					.bind(*id)
                       124: 					.fetch_one(&mut conn).await
                       125: 					.with_context(|| format!("Check post:\n{:?}", &conn))?;
                       126: 				let exists: bool = row.try_get("exists")?;
                       127: 				if ! exists {
                       128: 					if this_fetch == None || *date > this_fetch.unwrap() {
                       129: 						this_fetch = Some(*date);
                       130: 					};
                       131: 					self.tg.send( match iv_hash {
                       132: 							Some(x) => SendMessage::new(destination, format!("<a href=\"https://t.me/iv?url={}&rhash={}\"> </a>{0}", url, x)),
                       133: 							None => SendMessage::new(destination, format!("{}", url)),
                       134: 						}.parse_mode(types::ParseMode::Html)).await
                       135: 						.context("Can't post message:")?;
                       136: 					sqlx::query("insert into rsstg_post (source_id, posted, url) values ($1, $2, $3);")
                       137: 						.bind(*id)
                       138: 						.bind(date)
                       139: 						.bind(url)
                       140: 						.execute(&mut conn).await
                       141: 						.with_context(|| format!("Record post:\n{:?}", &conn))?;
                       142: 					drop(conn);
                       143: 					tokio::time::delay_for(std::time::Duration::new(4, 0)).await;
                       144: 				};
                       145: 			};
                       146: 			posts.clear();
                       147: 		};
                       148: 		let mut conn = self.pool.acquire().await
                       149: 			.with_context(|| format!("Update scrape fetch conn:\n{:?}", &self.pool))?;
                       150: 		sqlx::query("update rsstg_source set last_scrape = now() where source_id = $1;")
                       151: 			.bind(*id)
                       152: 			.execute(&mut conn).await
                       153: 			.with_context(|| format!("Update scrape:\n{:?}", &conn))?;
                       154: 		Ok(())
                       155: 	}
                       156: 
                       157: 	async fn delete<S>(&self, source_id: &i32, owner: S) -> Result<String>
                       158: 	where S: Into<i64> {
                       159: 		let owner: i64 = owner.into();
                       160: 		let mut conn = self.pool.acquire().await
                       161: 			.with_context(|| format!("Delete fetch conn:\n{:?}", &self.pool))?;
                       162: 		match sqlx::query("delete from rsstg_source where source_id = $1 and owner = $2;")
                       163: 			.bind(source_id)
                       164: 			.bind(owner)
                       165: 			.execute(&mut conn).await
                       166: 			.with_context(|| format!("Delete source rule:\n{:?}", &self.pool))?
                       167: 			.rows_affected() {
                       168: 			0 => { Ok("No data found found\\.".to_string()) },
                       169: 			x => { Ok(format!("{} sources removed\\.", x)) },
                       170: 		}
                       171: 	}
                       172: 
                       173: 	async fn clean<S>(&self, source_id: &i32, owner: S) -> Result<String>
                       174: 	where S: Into<i64> {
                       175: 		let owner: i64 = owner.into();
                       176: 		let mut conn = self.pool.acquire().await
                       177: 			.with_context(|| format!("Clean fetch conn:\n{:?}", &self.pool))?;
                       178: 		match sqlx::query("delete from rsstg_post p using rsstg_source s where p.source_id = $1 and owner = $2 and p.source_id = s.source_id;")
                       179: 			.bind(source_id)
                       180: 			.bind(owner)
                       181: 			.execute(&mut conn).await
                       182: 			.with_context(|| format!("Clean seen posts:\n{:?}", &self.pool))?
                       183: 			.rows_affected() {
                       184: 			0 => { Ok("No data found found\\.".to_string()) },
                       185: 			x => { Ok(format!("{} posts purged\\.", x)) },
                       186: 		}
                       187: 	}
                       188: 
                       189: 	async fn enable<S>(&self, source_id: &i32, owner: S) -> Result<&str>
                       190: 	where S: Into<i64> {
                       191: 		let owner: i64 = owner.into();
                       192: 		let mut conn = self.pool.acquire().await
                       193: 			.with_context(|| format!("Enable fetch conn:\n{:?}", &self.pool))?;
                       194: 		match sqlx::query("update rsstg_source set enabled = true where source_id = $1 and owner = $2")
                       195: 			.bind(source_id)
                       196: 			.bind(owner)
                       197: 			.execute(&mut conn).await
                       198: 			.with_context(|| format!("Enable source:\n{:?}", &self.pool))?
                       199: 			.rows_affected() {
                       200: 			1 => { Ok("Source enabled\\.") },
                       201: 			0 => { Ok("Source not found\\.") },
                       202: 			_ => { Err(anyhow!("Database error.")) },
                       203: 		}
                       204: 	}
                       205: 
                       206: 	async fn disable<S>(&self, source_id: &i32, owner: S) -> Result<&str>
                       207: 	where S: Into<i64> {
                       208: 		let owner: i64 = owner.into();
                       209: 		let mut conn = self.pool.acquire().await
                       210: 			.with_context(|| format!("Disable fetch conn:\n{:?}", &self.pool))?;
                       211: 		match sqlx::query("update rsstg_source set enabled = false where source_id = $1 and owner = $2")
                       212: 			.bind(source_id)
                       213: 			.bind(owner)
                       214: 			.execute(&mut conn).await
                       215: 			.with_context(|| format!("Disable source:\n{:?}", &self.pool))?
                       216: 			.rows_affected() {
                       217: 			1 => { Ok("Source disabled\\.") },
                       218: 			0 => { Ok("Source not found\\.") },
                       219: 			_ => { Err(anyhow!("Database error.")) },
                       220: 		}
                       221: 	}
                       222: 
                       223: 	async fn update<S>(&self, update: Option<i32>, channel: &str, channel_id: i64, url: &str, iv_hash: Option<&str>, owner: S) -> Result<String>
                       224: 	where S: Into<i64> {
                       225: 		let owner: i64 = owner.into();
                       226: 		let mut conn = self.pool.acquire().await
                       227: 			.with_context(|| format!("Update fetch conn:\n{:?}", &self.pool))?;
                       228: 
                       229: 		match match update {
                       230: 				Some(id) => {
                       231: 					sqlx::query("update rsstg_source set channel_id = $2, url = $3, iv_hash = $4, owner = $5, channel = $6 where source_id = $1").bind(id)
                       232: 				},
                       233: 				None => {
                       234: 					sqlx::query("insert into rsstg_source (channel_id, url, iv_hash, owner, channel) values ($1, $2, $3, $4, $5)")
                       235: 				},
                       236: 			}
                       237: 			.bind(channel_id)
                       238: 			.bind(url)
                       239: 			.bind(iv_hash)
                       240: 			.bind(owner)
                       241: 			.bind(channel)
                       242: 			.execute(&mut conn).await {
                       243: 			Ok(_) => return Ok(String::from("Channel added\\.")),
                       244: 			Err(sqlx::Error::Database(err)) => {
                       245: 				match err.downcast::<sqlx::postgres::PgDatabaseError>().routine() {
                       246: 					Some("_bt_check_unique", ) => {
                       247: 						return Ok("Duplicate key\\.".to_string())
                       248: 					},
                       249: 					Some(_) => {
                       250: 						return Ok("Database error\\.".to_string())
                       251: 					},
                       252: 					None => {
                       253: 						return Ok("No database error extracted\\.".to_string())
                       254: 					},
                       255: 				};
                       256: 			},
                       257: 			Err(err) => {
                       258: 				bail!("Sorry, unknown error:\n{:#?}\n", err);
                       259: 			},
                       260: 		};
                       261: 	}
                       262: 
                       263: 	async fn autofetch(&self) -> Result<()> {
                       264: 		let mut delay = chrono::Duration::minutes(5);
                       265: 		let mut now;
                       266: 		loop {
                       267: 			let mut conn = self.pool.acquire().await
                       268: 				.with_context(|| format!("Autofetch fetch conn:\n{:?}", &self.pool))?;
                       269: 			now = chrono::Local::now();
                       270: 			let mut queue = sqlx::query("select source_id, next_fetch, owner from rsstg_order natural left join rsstg_source where next_fetch < now() + interval '5 minutes';")
                       271: 				.fetch_all(&mut conn).await?;
                       272: 			for row in queue.iter() {
                       273: 				let source_id: i32 = row.try_get("source_id")?;
                       274: 				let owner: i64 = row.try_get("owner")?;
                       275: 				let next_fetch: DateTime<chrono::Local> = row.try_get("next_fetch")?;
                       276: 				if next_fetch < now {
                       277: 					//let clone = self.clone();
                       278: 					//clone.owner_chat(UserId::new(owner));
                       279: 					let clone = Core {
                       280: 						owner_chat: UserId::new(owner),
                       281: 						..self.clone()
                       282: 					};
                       283: 					tokio::spawn(async move {
                       284: 						if let Err(err) = clone.check(source_id, owner, true).await {
a68576cc1b 2020-12-15  285: 							if let Err(err) = clone.debug(&format!("šŸ›‘ {:?}", err)) {
                       286: 								eprintln!("Check error: {}", err);
                       287: 							};
                       288: 						};
                       289: 					});
                       290: 				} else {
                       291: 					if next_fetch - now < delay {
                       292: 						delay = next_fetch - now;
                       293: 					}
                       294: 				}
                       295: 			};
                       296: 			queue.clear();
                       297: 			tokio::time::delay_for(delay.to_std()?).await;
                       298: 			delay = chrono::Duration::minutes(5);
                       299: 		}
                       300: 	}
                       301: 
                       302: 	async fn list<S>(&self, owner: S) -> Result<Vec<String>>
                       303: 	where S: Into<i64> {
                       304: 		let owner = owner.into();
                       305: 		let mut reply = vec![];
                       306: 		let mut conn = self.pool.acquire().await
                       307: 			.with_context(|| format!("List fetch conn:\n{:?}", &self.pool))?;
                       308: 		reply.push("Channels:".to_string());
                       309: 		let rows = sqlx::query("select source_id, channel, enabled, url, iv_hash from rsstg_source where owner = $1 order by source_id")
                       310: 			.bind(owner)
                       311: 			.fetch_all(&mut conn).await?;
                       312: 		for row in rows.iter() {
                       313: 			let source_id: i32 = row.try_get("source_id")?;
                       314: 			let username: &str = row.try_get("channel")?;
                       315: 			let enabled: bool = row.try_get("enabled")?;
                       316: 			let url: &str = row.try_get("url")?;
                       317: 			let iv_hash: Option<&str> = row.try_get("iv_hash")?;
                       318: 			reply.push(format!("\n\\#ļøāƒ£ {} \\*ļøāƒ£ `{}` {}\nšŸ”— `{}`", source_id, username,  
                       319: 				match enabled {
                       320: 					true  => "šŸ”„ enabled",
                       321: 					false => "ā›” disabled",
                       322: 				}, url));
                       323: 			if let Some(hash) = iv_hash {
                       324: 				reply.push(format!("IV `{}`", hash));
                       325: 			}
                       326: 		};
                       327: 		Ok(reply)
                       328: 	}
                       329: }
                       330: 
                       331: #[tokio::main]
                       332: async fn main() -> Result<()> {
                       333: 	let mut settings = config::Config::default();
                       334: 	settings.merge(config::File::with_name("rsstg"))?;
                       335: 
                       336: 	let core = Core::new(settings).await?;
                       337: 
                       338: 	let mut stream = core.stream();
                       339: 
                       340: 	loop {
                       341: 		match stream.next().await {
                       342: 			Some(update) => {
a68576cc1b 2020-12-15  343: 				if let Err(err) = handle(update?, &core).await {
a68576cc1b 2020-12-15  344: 					core.debug(&format!("šŸ›‘ {:?}", err))?;
                       345: 				};
                       346: 			},
                       347: 			None => {
a68576cc1b 2020-12-15  348: 				core.debug(&format!("šŸ›‘ None error."))?;
                       349: 			}
                       350: 		};
                       351: 	}
                       352: 
                       353: 	//Ok(())
                       354: }
                       355: 
a68576cc1b 2020-12-15  356: async fn handle(update: telegram_bot::Update, core: &Core) -> Result<()> {
                       357: 	lazy_static! {
                       358: 		static ref RE_USERNAME: Regex = Regex::new(r"^@[a-zA-Z][a-zA-Z0-9_]+$").unwrap();
                       359: 		static ref RE_LINK: Regex = Regex::new(r"^https?://[a-zA-Z.0-9-]+/[-_a-zA-Z.0-9/?=]+$").unwrap();
                       360: 		static ref RE_IV_HASH: Regex = Regex::new(r"^[a-f0-9]{14}$").unwrap();
                       361: 	}
                       362: 
                       363: 	match update.kind {
                       364: 		UpdateKind::Message(message) => {
                       365: 			let mut reply: Vec<String> = vec![];
                       366: 			match message.kind {
                       367: 				MessageKind::Text { ref data, .. } => {
                       368: 					let mut words = data.split_whitespace();
                       369: 					let cmd = words.next().unwrap();
                       370: 					match cmd {
                       371: 
                       372: // start
                       373: 
                       374: 						"/start" => {
                       375: 							reply.push("We are open\\. Probably\\. Visit [channel](https://t.me/rsstg_bot_help/3) for details\\.".to_string());
                       376: 						},
                       377: 
                       378: // list
                       379: 
                       380: 						"/list" => {
                       381: 							reply.append(&mut core.list(message.from.id).await?);
                       382: 						},
                       383: 
                       384: // add
                       385: 
                       386: 						"/add" | "/update" => {
                       387: 							let mut source_id: Option<i32> = None;
                       388: 							let at_least = "Requires at least 3 parameters.";
                       389: 							if cmd == "/update" {
                       390: 								let first_word = words.next()
                       391: 									.context(at_least)?;
                       392: 								source_id = Some(first_word.parse::<i32>()
                       393: 									.with_context(|| format!("I need a number, but got {}.", first_word))?);
                       394: 							}
                       395: 							let (channel, url, iv_hash) = (
                       396: 								words.next().context(at_least)?,
                       397: 								words.next().context(at_least)?,
                       398: 								words.next());
                       399: 							if ! RE_USERNAME.is_match(&channel) {
                       400: 								reply.push("Usernames should be something like \"@\\[a\\-zA\\-Z]\\[a\\-zA\\-Z0\\-9\\_]+\", aren't they?".to_string());
                       401: 								bail!("Wrong username {:?}.", &channel);
                       402: 							}
                       403: 							if ! RE_LINK.is_match(&url) {
                       404: 								reply.push("Link should be link to atom/rss feed, something like \"https://domain/path\"\\.".to_string());
                       405: 								bail!("Url: {:?}", &url);
                       406: 							}
                       407: 							if let Some(hash) = iv_hash {
                       408: 								if ! RE_IV_HASH.is_match(&hash) {
                       409: 									reply.push("IV hash should be 14 hex digits.".to_string());
                       410: 									bail!("IV: {:?}", &iv_hash);
                       411: 								};
                       412: 							};
                       413: 							let channel_id = i64::from(core.tg.send(telegram_bot::GetChat::new(telegram_bot::types::ChatRef::ChannelUsername(channel.to_string()))).await?.id());
                       414: 							let chan_adm = core.tg.send(telegram_bot::GetChatAdministrators::new(telegram_bot::types::ChatRef::ChannelUsername(channel.to_string()))).await
                       415: 								.context("Sorry, I have no access to that chat\\.")?;
                       416: 							let (mut me, mut user) = (false, false);
                       417: 							for admin in chan_adm {
                       418: 								if admin.user.id == core.my.id {
                       419: 									me = true;
                       420: 								};
                       421: 								if admin.user.id == message.from.id {
                       422: 									user = true;
                       423: 								};
                       424: 							};
                       425: 							if ! me   { bail!("I need to be admin on that channel\\."); };
                       426: 							if ! user { bail!("You should be admin on that channel\\."); };
                       427: 							reply.push(core.update(source_id, channel, channel_id, url, iv_hash, message.from.id).await?);
                       428: 						},
                       429: 
                       430: // check
                       431: 
                       432: 						"/check" => {
                       433: 							match &words.next().unwrap().parse::<i32>() {
                       434: 								Err(err) => {
                       435: 									reply.push(format!("I need a number\\.\n{}", &err));
                       436: 								},
                       437: 								Ok(number) => {
                       438: 									core.check(*number, message.from.id, false).await
                       439: 										.context("Channel check failed.")?;
                       440: 								},
                       441: 							};
                       442: 						},
                       443: 
                       444: // clean
                       445: 
                       446: 						"/clean" => {
                       447: 							match &words.next().unwrap().parse::<i32>() {
                       448: 								Err(err) => {
                       449: 									reply.push(format!("I need a number\\.\n{}", &err));
                       450: 								},
                       451: 								Ok(number) => {
                       452: 									let result = core.clean(&number, message.from.id).await?;
                       453: 									reply.push(result.to_string());
                       454: 								},
                       455: 							};
                       456: 						},
                       457: 
                       458: // enable
                       459: 
                       460: 						"/enable" => {
                       461: 							match &words.next().unwrap().parse::<i32>() {
                       462: 								Err(err) => {
                       463: 									reply.push(format!("I need a number\\.\n{}", &err));
                       464: 								},
                       465: 								Ok(number) => {
                       466: 									let result = core.enable(&number, message.from.id).await?;
                       467: 									reply.push(result.to_string());
                       468: 								},
                       469: 							};
                       470: 						},
                       471: 
                       472: // delete
                       473: 
                       474: 						"/delete" => {
                       475: 							match &words.next().unwrap().parse::<i32>() {
                       476: 								Err(err) => {
                       477: 									reply.push(format!("I need a number\\.\n{}", &err));
                       478: 								},
                       479: 								Ok(number) => {
                       480: 									let result = core.delete(&number, message.from.id).await?;
                       481: 									reply.push(result.to_string());
                       482: 								},
                       483: 							};
                       484: 						},
                       485: 
                       486: // disable
                       487: 
                       488: 						"/disable" => {
                       489: 							match &words.next().unwrap().parse::<i32>() {
                       490: 								Err(err) => {
                       491: 									reply.push(format!("I need a number\\.\n{}", &err));
                       492: 								},
                       493: 								Ok(number) => {
                       494: 									let result = core.disable(&number, message.from.id).await?;
                       495: 									reply.push(result.to_string());
                       496: 								},
                       497: 							};
                       498: 						},
                       499: 
                       500: 						_ => {
                       501: 						},
                       502: 					};
                       503: 				},
                       504: 				_ => {
                       505: 				},
                       506: 			};
                       507: 
                       508: 			if reply.len() > 0 {
                       509: 				if let Err(err) = core.tg.send(message.text_reply(reply.join("\n")).parse_mode(types::ParseMode::MarkdownV2)).await {
                       510: 					dbg!(reply.join("\n"));
                       511: 					println!("{}", err);
                       512: 				};
                       513: 			};
                       514: 		},
                       515: 		_ => {},
                       516: 	};
                       517: 
                       518: 	Ok(())
                       519: }