mod entity;
mod migration;
use std::env;
use std::io::Cursor;
use std::path::PathBuf;
use async_compression::tokio::bufread::ZstdDecoder;
use async_compression::tokio::write::ZstdEncoder;
use chrono::NaiveDateTime;
use image::{DynamicImage, ImageReader};
use sea_orm::sqlx::sqlite::{SqliteJournalMode, SqliteSynchronous};
use sea_orm::{ActiveModelTrait, ConnectOptions, Database, DatabaseConnection, EntityTrait};
use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
use url::Url;
use self::entity::{Image, Text};
use self::migration::{Migrator, MigratorTrait};
use crate::{ChapterInfo, Error};
#[must_use]
pub(crate) struct NovelDB {
db: DatabaseConnection,
}
#[must_use]
#[derive(Debug, PartialEq)]
pub(crate) enum FindTextResult {
Ok(String),
None,
Outdate,
}
#[must_use]
#[derive(Debug, PartialEq)]
pub(crate) enum FindImageResult {
Ok(DynamicImage),
None,
}
impl NovelDB {
const DB_NAME: &'static str = "novel.db";
pub(crate) async fn new(app_name: &str) -> Result<Self, Error> {
let db_path = NovelDB::db_path(app_name)?;
if fs::try_exists(&db_path).await? {
tracing::info!("The database file is located at `{}`", db_path.display());
} else {
tracing::info!(
"The database file will be created at `{}`",
db_path.display()
);
fs::create_dir_all(db_path.parent().unwrap()).await?;
}
let db_url = format!("sqlite:{}?mode=rwc", db_path.display());
let mut db = NovelDB::connect_db(&db_url).await?;
if Migrator::up(&db, None).await.is_err() {
tracing::error!("The file may not be a database, try recreating it");
db.close().await?;
let backup_path = db_path.with_extension("backup");
fs::rename(&db_path, &backup_path).await?;
tracing::info!("The file has been backed up to `{}`", backup_path.display());
db = NovelDB::connect_db(&db_url).await?;
Migrator::up(&db, None).await?;
}
Ok(Self { db })
}
async fn connect_db(db_url: &str) -> Result<DatabaseConnection, Error> {
let mut options = ConnectOptions::new(db_url);
options.map_sqlx_sqlite_opts(|opt| {
opt.journal_mode(SqliteJournalMode::Wal)
.synchronous(SqliteSynchronous::Normal)
});
Ok(Database::connect(options).await?)
}
#[cfg(test)]
pub(crate) async fn drop(&self) -> Result<(), Error> {
Ok(Migrator::down(&self.db, None).await?)
}
pub(crate) async fn find_text(&self, info: &ChapterInfo) -> Result<FindTextResult, Error> {
match Text::find_by_id(info.id).one(&self.db).await? {
Some(model) => {
let saved_data_time = model.date_time;
let time = NovelDB::get_time(info);
if (time.is_some()
&& saved_data_time.is_some()
&& saved_data_time.unwrap() < time.unwrap())
|| env::var("FORCE_UPDATE_NOVEL_DB").is_ok()
{
Ok(FindTextResult::Outdate)
} else {
Ok(FindTextResult::Ok(unsafe {
String::from_utf8_unchecked(zstd_decompress(&model.content).await?)
}))
}
}
None => Ok(FindTextResult::None),
}
}
pub(crate) async fn insert_text<T>(&self, info: &ChapterInfo, text: T) -> Result<(), Error>
where
T: AsRef<str>,
{
let model = entity::text::ActiveModel {
id: sea_orm::Set(info.id),
date_time: sea_orm::Set(NovelDB::get_time(info)),
content: sea_orm::Set(zstd_compress(text.as_ref().as_bytes()).await?),
};
model.insert(&self.db).await?;
Ok(())
}
pub(crate) async fn update_text<T>(&self, info: &ChapterInfo, text: T) -> Result<(), Error>
where
T: AsRef<str>,
{
let model = entity::text::ActiveModel {
id: sea_orm::Set(info.id),
date_time: sea_orm::Set(NovelDB::get_time(info)),
content: sea_orm::Set(zstd_compress(text.as_ref().as_bytes()).await?),
};
model.update(&self.db).await?;
Ok(())
}
pub(crate) async fn find_image(&self, url: &Url) -> Result<FindImageResult, Error> {
let model = Image::find_by_id(url.to_string()).one(&self.db).await?;
match model {
Some(model) => {
let bytes = zstd_decompress(&model.content).await?;
let image = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()?
.decode()?;
Ok(FindImageResult::Ok(image))
}
None => Ok(FindImageResult::None),
}
}
pub(crate) async fn insert_image<T>(&self, url: &Url, bytes: T) -> Result<(), Error>
where
T: AsRef<[u8]>,
{
let model = entity::image::ActiveModel {
url: sea_orm::Set(url.to_string()),
content: sea_orm::Set(zstd_compress(bytes).await?),
};
model.insert(&self.db).await?;
Ok(())
}
fn db_path(app_name: &str) -> Result<PathBuf, Error> {
let mut db_path = crate::data_dir_path(app_name)?;
db_path.push(NovelDB::DB_NAME);
Ok(db_path)
}
fn get_time(info: &ChapterInfo) -> Option<NaiveDateTime> {
if info.update_time.is_some() {
info.update_time
} else {
info.create_time
}
}
}
async fn zstd_decompress<T>(data: T) -> Result<Vec<u8>, Error>
where
T: AsRef<[u8]>,
{
let mut reader = ZstdDecoder::new(BufReader::new(data.as_ref()));
let mut buf = Vec::new();
reader.read_to_end(&mut buf).await?;
Ok(buf)
}
async fn zstd_compress<T>(data: T) -> Result<Vec<u8>, Error>
where
T: AsRef<[u8]>,
{
let mut writer = ZstdEncoder::new(Vec::new());
writer.write_all(data.as_ref()).await?;
writer.shutdown().await?;
let mut res = writer.into_inner();
res.flush().await?;
Ok(res)
}
#[cfg(test)]
mod tests {
use std::str::FromStr;
use pretty_assertions::assert_eq;
use super::*;
#[tokio::test]
async fn zstd() -> Result<(), Error> {
let data = "test-data";
let compressed_data = zstd_compress(data).await?;
let decompressed_data = zstd_decompress(compressed_data).await?;
assert_eq!(data.as_bytes(), decompressed_data.as_slice());
Ok(())
}
#[tokio::test]
async fn db() -> Result<(), Error> {
let app_name = "test-app";
let contents = "test-contents";
let db = NovelDB::new(app_name).await?;
let chapter_info_old = ChapterInfo {
id: 0,
update_time: Some(NaiveDateTime::from_str("2020-07-08T15:25:15")?),
..Default::default()
};
let chapter_info_new = ChapterInfo {
id: 0,
update_time: Some(NaiveDateTime::from_str("2020-07-08T15:25:17")?),
..Default::default()
};
assert_eq!(db.find_text(&chapter_info_new).await?, FindTextResult::None);
db.insert_text(&chapter_info_old, contents).await?;
assert_eq!(
db.find_text(&chapter_info_new).await?,
FindTextResult::Outdate
);
db.update_text(&chapter_info_new, contents).await?;
if let FindTextResult::Ok(result) = db.find_text(&chapter_info_new).await? {
assert_eq!(result, contents);
} else {
panic!("Incorrect database query result");
}
db.drop().await?;
Ok(())
}
}