1 // Copyright (C) 2024 Huawei Device Co., Ltd. 2 // Licensed under the Apache License, Version 2.0 (the "License"); 3 // you may not use this file except in compliance with the License. 4 // You may obtain a copy of the License at 5 // 6 // http://www.apache.org/licenses/LICENSE-2.0 7 // 8 // Unless required by applicable law or agreed to in writing, software 9 // distributed under the License is distributed on an "AS IS" BASIS, 10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 14 use crate::database::REQUEST_DB; 15 16 const CREATE_TASK_CONFIG_TABLE: &str = 17 "CREATE TABLE IF NOT EXISTS task_config (task_id INTEGER PRIMARY KEY, display BOOLEAN)"; 18 19 const CREATE_GROUP_TABLE: &str = 20 "CREATE TABLE IF NOT EXISTS group_notification (task_id INTEGER PRIMARY KEY, group_id INTEGER)"; 21 22 const CREATE_GROUP_CONFIG_TABLE: &str = 23 "CREATE TABLE IF NOT EXISTS group_notification_config (group_id INTEGER PRIMARY KEY, gauge BOOLEAN, attach_able BOOLEAN, ctime INTEGER)"; 24 25 const CREATE_TASK_CONTENT_TABLE: &str = 26 "CREATE TABLE IF NOT EXISTS task_notification_content (task_id INTEGER PRIMARY KEY, title TEXT, text TEXT)"; 27 28 const CREATE_GROUP_CONTENT_TABLE: &str = 29 "CREATE TABLE IF NOT EXISTS group_notification_content (group_id INTEGER PRIMARY KEY, title TEXT, text TEXT)"; 30 31 use std::time::{SystemTime, UNIX_EPOCH}; 32 33 const MILLIS_IN_A_WEEK: u64 = 7 * 24 * 60 * 60 * 1000; 34 35 pub(crate) struct NotificationDb { 36 inner: &'static rdb::RdbStore<'static>, 37 } 38 39 #[derive(Default, Clone)] 40 pub(crate) struct CustomizedNotification { 41 pub title: Option<String>, 42 pub text: Option<String>, 43 } 44 45 impl NotificationDb { new() -> Self46 pub(crate) fn new() -> Self { 47 let me = Self { inner: &REQUEST_DB }; 48 if let Err(e) = me.create_db() { 49 error!("Failed to create notification database: {}", e); 50 } 51 me 52 } 53 create_db(&self) -> Result<(), i32>54 fn create_db(&self) -> Result<(), i32> { 55 self.inner.execute(CREATE_TASK_CONFIG_TABLE, ())?; 56 self.inner.execute(CREATE_GROUP_CONTENT_TABLE, ())?; 57 self.inner.execute(CREATE_GROUP_TABLE, ())?; 58 self.inner.execute(CREATE_TASK_CONTENT_TABLE, ())?; 59 self.inner.execute(CREATE_GROUP_CONFIG_TABLE, ())?; 60 Ok(()) 61 } 62 clear_task_info(&self, task_id: u32)63 pub(crate) fn clear_task_info(&self, task_id: u32) { 64 let sqls = [ 65 "DELETE FROM task_config WHERE task_id = ?", 66 "DELETE FROM task_notification_content WHERE task_id = ?", 67 "DELETE FROM group_notification WHERE task_id = ?", 68 ]; 69 for sql in sqls.iter() { 70 if let Err(e) = self.inner.execute(sql, task_id) { 71 error!( 72 "Failed to clear task {} notification info: {}, sql: {}", 73 task_id, e, sql 74 ); 75 } 76 } 77 } 78 clear_group_info(&self, group_id: u32)79 pub(crate) fn clear_group_info(&self, group_id: u32) { 80 let sqls = [ 81 "DELETE FROM group_notification WHERE group_id = ?", 82 "DELETE FROM group_notification_content WHERE group_id = ?", 83 "DELETE FROM group_notification_config WHERE group_id = ?", 84 ]; 85 for sql in sqls.iter() { 86 if let Err(e) = self.inner.execute(sql, group_id) { 87 error!( 88 "Failed to clear group {} notification info: {}, sql: {}", 89 group_id, e, sql 90 ); 91 } 92 } 93 } 94 clear_group_info_a_week_ago(&self)95 pub(crate) fn clear_group_info_a_week_ago(&self) { 96 let current_time = match SystemTime::now().duration_since(UNIX_EPOCH) { 97 Ok(duration) => duration, 98 Err(e) => { 99 error!("Failed to get current time: {}", e); 100 return; 101 } 102 } 103 .as_millis() as u64; 104 let group_ids = match self.inner.query::<u32>( 105 "SELECT group_id FROM group_notification_config WHERE ctime < ?", 106 current_time - MILLIS_IN_A_WEEK, 107 ) { 108 Ok(rows) => rows, 109 Err(e) => { 110 error!("Failed to clear group info: {}", e); 111 return; 112 } 113 }; 114 for group_id in group_ids { 115 let mut count = match self.inner.query::<u32>( 116 "SELECT COUNT(*) FROM group_notification WHERE group_id = ?", 117 group_id, 118 ) { 119 Ok(rows) => rows, 120 Err(e) => { 121 error!("Failed to clear group info: {}", e); 122 continue; 123 } 124 }; 125 if !count.next().is_some_and(|x| x == 0) { 126 continue; 127 } 128 129 info!( 130 "clear group {} info for have been overdue for more than a week.", 131 group_id 132 ); 133 self.clear_group_info(group_id); 134 } 135 } 136 check_task_notification_available(&self, task_id: &u32) -> bool137 pub(crate) fn check_task_notification_available(&self, task_id: &u32) -> bool { 138 let mut set = self 139 .inner 140 .query::<bool>("SELECT display FROM task_config WHERE task_id = ?", task_id) 141 .unwrap(); 142 set.next().unwrap_or(true) 143 } 144 update_task_group(&self, task_id: u32, group_id: u32)145 pub(crate) fn update_task_group(&self, task_id: u32, group_id: u32) { 146 if let Err(e) = self.inner.execute( 147 "INSERT INTO group_notification (task_id, group_id) VALUES (?, ?) ON CONFLICT(task_id) DO UPDATE SET group_id = excluded.group_id", 148 (task_id, group_id), 149 ) { 150 error!("Failed to update {} notification: {}", task_id, e); 151 } 152 } 153 query_group_tasks(&self, group_id: u32) -> Vec<u32>154 pub(crate) fn query_group_tasks(&self, group_id: u32) -> Vec<u32> { 155 let set = match self.inner.query::<u32>( 156 "SELECT task_id FROM group_notification WHERE group_id = ?", 157 group_id, 158 ) { 159 Ok(set) => set, 160 Err(e) => { 161 error!("Failed to query group tasks: {}", e); 162 return Vec::new(); 163 } 164 }; 165 set.collect() 166 } 167 query_task_gid(&self, task_id: u32) -> Option<u32>168 pub(crate) fn query_task_gid(&self, task_id: u32) -> Option<u32> { 169 let mut set = match self.inner.query::<u32>( 170 "SELECT group_id FROM group_notification WHERE task_id = ?", 171 task_id, 172 ) { 173 Ok(set) => set, 174 Err(e) => { 175 error!("Failed to query task group id: {}", e); 176 return None; 177 } 178 }; 179 set.next() 180 } 181 query_task_customized_notification( &self, task_id: u32, ) -> Option<CustomizedNotification>182 pub(crate) fn query_task_customized_notification( 183 &self, 184 task_id: u32, 185 ) -> Option<CustomizedNotification> { 186 let mut set = match self.inner.query::<(Option<String>, Option<String>)>( 187 "SELECT title, text FROM task_notification_content WHERE task_id = ?", 188 task_id, 189 ) { 190 Ok(set) => set, 191 Err(e) => { 192 error!("Failed to query task customized notification: {}", e); 193 return None; 194 } 195 }; 196 set.next() 197 .map(|(title, text)| CustomizedNotification { title, text }) 198 } 199 update_task_customized_notification( &self, task_id: u32, title: Option<String>, text: Option<String>, )200 pub(crate) fn update_task_customized_notification( 201 &self, 202 task_id: u32, 203 title: Option<String>, 204 text: Option<String>, 205 ) { 206 if let Err(e) = self.inner.execute( 207 "INSERT INTO task_notification_content (task_id, title, text) VALUES (?, ?, ?) ON CONFLICT(task_id) DO UPDATE SET title = excluded.title, text = excluded.text", 208 (task_id, title, text), 209 ) { 210 error!("Failed to insert {} notification: {}", task_id, e); 211 } 212 } 213 query_group_customized_notification( &self, group_id: u32, ) -> Option<CustomizedNotification>214 pub(crate) fn query_group_customized_notification( 215 &self, 216 group_id: u32, 217 ) -> Option<CustomizedNotification> { 218 let mut set = match self.inner.query::<(Option<String>, Option<String>)>( 219 "SELECT title, text FROM group_notification_content WHERE group_id = ?", 220 group_id, 221 ) { 222 Ok(set) => set, 223 Err(e) => { 224 error!("Failed to query task customized notification: {}", e); 225 return None; 226 } 227 }; 228 set.next() 229 .map(|(title, text)| CustomizedNotification { title, text }) 230 } 231 update_group_customized_notification( &self, group_id: u32, title: Option<String>, text: Option<String>, )232 pub(crate) fn update_group_customized_notification( 233 &self, 234 group_id: u32, 235 title: Option<String>, 236 text: Option<String>, 237 ) { 238 if let Err(e) = self.inner.execute( 239 "INSERT INTO group_notification_content (group_id, title, text) VALUES (?, ?, ?) ON CONFLICT(group_id) DO UPDATE SET title = excluded.title, text = excluded.text", 240 (group_id, title, text), 241 ) { 242 error!("Failed to insert {} notification: {}", group_id, e); 243 } 244 } 245 update_group_config(&self, group_id: u32, gauge: bool, ctime: u64)246 pub(crate) fn update_group_config(&self, group_id: u32, gauge: bool, ctime: u64) { 247 if let Err(e) = self.inner.execute( 248 "INSERT INTO group_notification_config (group_id, gauge, attach_able, ctime) VALUES (?, ?, ?, ?) ON CONFLICT(group_id) DO UPDATE SET gauge = excluded.gauge , ctime = excluded.ctime", 249 (group_id, gauge, true, ctime), 250 ) { 251 error!("Failed to update {} notification: {}", group_id, e); 252 } 253 } 254 255 #[allow(unused)] delete_task_customized(&self, task_id: u32)256 pub(crate) fn delete_task_customized(&self, task_id: u32) { 257 if let Err(e) = self.inner.execute( 258 "DELETE FROM task_notification_content WHERE task_id = ?", 259 task_id, 260 ) { 261 error!("Failed to delete {} notification: {}", task_id, e); 262 } 263 } 264 265 #[allow(unused)] delete_group_customized(&self, group_id: u32)266 pub(crate) fn delete_group_customized(&self, group_id: u32) { 267 if let Err(e) = self.inner.execute( 268 "DELETE FROM group_notification_content WHERE group_id = ?", 269 group_id, 270 ) { 271 error!("Failed to delete {} notification: {}", group_id, e); 272 } 273 } 274 contains_group(&self, group_id: u32) -> bool275 pub(crate) fn contains_group(&self, group_id: u32) -> bool { 276 let mut set = self 277 .inner 278 .query::<u32>( 279 "SELECT group_id FROM group_notification_config where group_id = ?", 280 group_id, 281 ) 282 .unwrap(); 283 set.row_count() == 1 284 } 285 attach_able(&self, group_id: u32) -> bool286 pub(crate) fn attach_able(&self, group_id: u32) -> bool { 287 let mut set = self 288 .inner 289 .query::<bool>( 290 "SELECT attach_able FROM group_notification_config where group_id = ?", 291 group_id, 292 ) 293 .unwrap(); 294 set.next().unwrap_or(false) 295 } 296 disable_attach_group(&self, group_id: u32)297 pub(crate) fn disable_attach_group(&self, group_id: u32) { 298 if let Err(e) = self.inner.execute( 299 " UPDATE group_notification_config SET attach_able = ? where group_id = ?", 300 (false, group_id), 301 ) { 302 error!("Failed to update {} notification: {}", group_id, e); 303 } 304 } 305 is_gauge(&self, group_id: u32) -> bool306 pub(crate) fn is_gauge(&self, group_id: u32) -> bool { 307 let mut set = self 308 .inner 309 .query::<bool>( 310 "SELECT gauge FROM group_notification_config where group_id = ?", 311 group_id, 312 ) 313 .unwrap(); 314 set.next().unwrap_or(false) 315 } 316 } 317 #[cfg(test)] 318 mod test { 319 use ylong_runtime::fastrand::fast_random; 320 321 use super::*; 322 const TEST_TITLE: &str = "田文镜"; 323 const TEST_TEXT: &str = "我XXX"; 324 #[test] ut_notify_database_query_tasks()325 fn ut_notify_database_query_tasks() { 326 let db = NotificationDb::new(); 327 let group_id = fast_random() as u32; 328 let mut v = vec![]; 329 for _ in 0..100 { 330 let task_id = fast_random() as u32; 331 v.push(task_id); 332 db.update_task_group(task_id, group_id); 333 } 334 v.sort(); 335 let mut ans = db.query_group_tasks(group_id); 336 ans.sort(); 337 assert_eq!(v, ans); 338 } 339 340 #[test] ut_notify_database_query_task_gid()341 fn ut_notify_database_query_task_gid() { 342 let db = NotificationDb::new(); 343 let group_id = fast_random() as u32; 344 345 for _ in 0..100 { 346 let task_id = fast_random() as u32; 347 db.update_task_group(task_id, group_id); 348 assert_eq!(db.query_task_gid(task_id).unwrap(), group_id); 349 } 350 } 351 352 #[test] ut_notify_database_query_task_customized()353 fn ut_notify_database_query_task_customized() { 354 let db = NotificationDb::new(); 355 let task_id = fast_random() as u32; 356 357 db.update_task_customized_notification( 358 task_id, 359 Some(TEST_TITLE.to_string()), 360 Some(TEST_TEXT.to_string()), 361 ); 362 let customized = db.query_task_customized_notification(task_id).unwrap(); 363 assert_eq!(customized.title.unwrap(), TEST_TITLE); 364 assert_eq!(customized.text.unwrap(), TEST_TEXT); 365 } 366 367 #[test] ut_notify_database_query_group_customized()368 fn ut_notify_database_query_group_customized() { 369 let db = NotificationDb::new(); 370 let group_id = fast_random() as u32; 371 372 db.update_group_customized_notification( 373 group_id, 374 Some(TEST_TITLE.to_string()), 375 Some(TEST_TEXT.to_string()), 376 ); 377 let customized = db.query_group_customized_notification(group_id).unwrap(); 378 assert_eq!(customized.title.unwrap(), TEST_TITLE); 379 assert_eq!(customized.text.unwrap(), TEST_TEXT); 380 } 381 382 #[test] ut_notify_database_group_config()383 fn ut_notify_database_group_config() { 384 let db = NotificationDb::new(); 385 let group_id = fast_random() as u32; 386 387 assert!(!db.contains_group(group_id)); 388 db.update_group_config(group_id, true, 0); 389 assert!(db.contains_group(group_id)); 390 assert!(db.is_gauge(group_id)); 391 assert!(db.attach_able(group_id)); 392 db.update_group_config(group_id, false, 0); 393 db.disable_attach_group(group_id); 394 assert!(!db.attach_able(group_id)); 395 assert!(!db.is_gauge(group_id)); 396 } 397 398 #[test] ut_clear_task_info()399 fn ut_clear_task_info() { 400 let db = NotificationDb::new(); 401 402 let group_id = fast_random() as u32; 403 let task_id = fast_random() as u32; 404 405 db.update_task_customized_notification(task_id, None, None); 406 db.update_task_group(task_id, group_id); 407 assert!(db.query_task_customized_notification(task_id).is_some()); 408 assert_eq!(db.query_task_gid(task_id).unwrap(), group_id); 409 410 db.clear_task_info(task_id); 411 assert!(db.query_task_customized_notification(task_id).is_none()); 412 assert!(db.query_task_gid(task_id).is_none()); 413 } 414 415 #[test] ut_clear_group_info()416 fn ut_clear_group_info() { 417 let db = NotificationDb::new(); 418 419 let group_id = fast_random() as u32; 420 let task_id = fast_random() as u32; 421 db.update_group_customized_notification(group_id, None, None); 422 db.update_group_config(group_id, true, 0); 423 db.update_task_group(task_id, group_id); 424 425 assert!(db.query_group_customized_notification(group_id).is_some()); 426 assert!(db.contains_group(group_id)); 427 assert_eq!(db.query_task_gid(task_id).unwrap(), group_id); 428 429 db.clear_group_info(group_id); 430 assert!(db.query_group_customized_notification(group_id).is_none()); 431 assert!(!db.contains_group(group_id)); 432 assert!(db.query_task_gid(task_id).is_none()); 433 } 434 435 #[test] ut_clear_group_info_a_week_ago()436 fn ut_clear_group_info_a_week_ago() { 437 let current_time = SystemTime::now() 438 .duration_since(UNIX_EPOCH) 439 .unwrap() 440 .as_millis() as u64; 441 let a_week_ago = current_time - MILLIS_IN_A_WEEK; 442 443 let db = NotificationDb::new(); 444 let group_id = fast_random() as u32; 445 let task_id = fast_random() as u32; 446 447 db.update_group_customized_notification(group_id, None, None); 448 db.update_group_config(group_id, true, current_time); 449 450 db.clear_group_info_a_week_ago(); 451 assert!(db.query_group_customized_notification(group_id).is_some()); 452 assert!(db.contains_group(group_id)); 453 454 db.update_group_config(group_id, true, a_week_ago); 455 db.update_task_group(task_id, group_id); 456 db.clear_group_info_a_week_ago(); 457 assert!(db.query_group_customized_notification(group_id).is_some()); 458 assert!(db.contains_group(group_id)); 459 460 db.clear_task_info(task_id); 461 db.clear_group_info_a_week_ago(); 462 assert!(db.query_group_customized_notification(group_id).is_none()); 463 assert!(!db.contains_group(group_id)); 464 } 465 } 466