use std::env;
use std::io::Cursor;
use std::path::PathBuf;
use async_compression::tokio::bufread::ZstdDecoder;
use async_compression::tokio::write::ZstdEncoder;
use image::{DynamicImage, ImageReader};
use jiff::Timestamp;
use toasty::Db;
use toasty::migration::MigrationSet;
use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
use url::Url;
use crate::{ChapterInfo, Error};
#[derive(toasty::Model)]
struct Text {
#[key]
pub id: u32,
pub date_time: Option<Timestamp>,
pub content: Vec<u8>,
}
#[derive(toasty::Model)]
struct Image {
#[key]
pub url: String,
pub content: Vec<u8>,
}
#[must_use]
pub(crate) struct NovelDB {
db: Db,
}
#[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,
}
static MIGRATIONS: MigrationSet = toasty::embed_migrations!();
impl NovelDB {
const DB_NAME: &'static str = "novel.sqlite.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 db = toasty::Db::builder()
.models(toasty::models!(crate::*))
.connect(&db_url)
.await?;
let report = MIGRATIONS.apply(&db).await?;
tracing::debug!(?report);
Ok(Self { db })
}
#[cfg(test)]
pub(crate) async fn drop(&self) -> Result<(), Error> {
Ok(self.db.reset_db().await?)
}
pub(crate) async fn find_text(&self, info: &ChapterInfo) -> Result<FindTextResult, Error> {
let mut conn = self.db.connection().await?;
match Text::get_by_id(&mut conn, info.id).await {
Ok(text) => {
let saved_data_time = text.date_time;
let time = NovelDB::get_time(info);
if time.is_some()
&& saved_data_time.is_some()
&& saved_data_time.unwrap() < time.unwrap()
{
Ok(FindTextResult::Outdate)
} else {
Ok(FindTextResult::Ok(unsafe {
String::from_utf8_unchecked(zstd_decompress(&text.content).await?)
}))
}
}
Err(error) if error.is_record_not_found() => Ok(FindTextResult::None),
Err(error) => Err(error.into()),
}
}
pub(crate) async fn insert_text<T>(&self, info: &ChapterInfo, text: T) -> Result<(), Error>
where
T: AsRef<str>,
{
let mut conn = self.db.connection().await?;
toasty::create!(Text {
id: info.id,
date_time: NovelDB::get_time(info),
content: zstd_compress(text.as_ref().as_bytes()).await?,
})
.exec(&mut conn)
.await?;
Ok(())
}
pub(crate) async fn update_text<T>(&self, info: &ChapterInfo, text: T) -> Result<(), Error>
where
T: AsRef<str>,
{
let mut conn = self.db.connection().await?;
Text::update_by_id(info.id)
.date_time(NovelDB::get_time(info))
.content(zstd_compress(text.as_ref().as_bytes()).await?)
.exec(&mut conn)
.await?;
Ok(())
}
pub(crate) async fn find_image(&self, url: &Url) -> Result<FindImageResult, Error> {
let mut conn = self.db.connection().await?;
let image = Image::get_by_url(&mut conn, url.to_string()).await;
match image {
Ok(image) => {
let bytes = zstd_decompress(&image.content).await?;
let image = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()?
.decode()?;
Ok(FindImageResult::Ok(image))
}
Err(error) if error.is_record_not_found() => Ok(FindImageResult::None),
Err(error) => Err(error.into()),
}
}
pub(crate) async fn insert_image<T>(&self, url: &Url, bytes: T) -> Result<(), Error>
where
T: AsRef<[u8]>,
{
let mut conn = self.db.connection().await?;
toasty::create!(Image {
url: url.to_string(),
content: zstd_compress(bytes).await?,
})
.exec(&mut conn)
.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<Timestamp> {
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 pretty_assertions::assert_eq;
use testresult::TestResult;
use toasty_cli::{Config, ToastyCli};
use super::*;
#[tokio::test]
async fn zstd() -> TestResult {
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() -> TestResult {
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("2020-07-08T15:25:15Z".parse()?),
..Default::default()
};
let chapter_info_new = ChapterInfo {
id: 0,
update_time: Some("2020-07-08T15:25:17Z".parse()?),
..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(())
}
#[tokio::test]
async fn migration() -> TestResult {
let config = Config::load()?;
let db = NovelDB::new("migration").await?;
let cli = ToastyCli::with_config(db.db, config);
cli.parse_from(["", "migration", "generate"]).await?;
Ok(())
}
}