#[derive(Clone)] pub struct DatabaseBackend { pub db_url: String, pub sqlite: Option<sqlx::Pool<sqlx::Sqlite>>, pub postgres: Option<sqlx::Pool<sqlx::Postgres>>, } pub fn ensure_file_exists(path: &str) { // Check if the file exists if !std::path::Path::new(path).exists() { // If the file does not exist, create an empty one match std::fs::File::create(path) { Ok(_) => log::info!("Created {path}"), Err(e) => log::error!("Failed to create file: {}", e), } } } impl DatabaseBackend { pub async fn new(db_url: &str) -> Self { let mut sqlite = None; let mut postgres = None; if db_url.starts_with("postgres") { postgres = Some( sqlx::postgres::PgPoolOptions::new() .max_connections(5) .connect(&std::env::var("DATABASE_URL").unwrap()) .await .unwrap(), ); sqlx::migrate!("./migrations") .run(postgres.as_ref().unwrap()) .await .unwrap(); } else { ensure_file_exists(db_url); sqlite = Some(sqlx::sqlite::SqlitePool::connect(db_url).await.unwrap()); sqlx::migrate!("./migrations") .run(sqlite.as_ref().unwrap()) .await .unwrap(); } Self { db_url: db_url.to_string(), sqlite, postgres, } } pub fn take_db(&self) -> Database { Database::new(self.clone()) } pub async fn query(&self, param: Query) -> Out { match param { Query::InsertUrl(ref url) => { if let Some(postgres) = self.postgres.as_ref() { sqlx::query("INSERT INTO urls (url, timestamp) VALUES ($1, CURRENT_TIMESTAMP)") .bind(url) .execute(postgres) .await .unwrap(); } else { if let Some(sqlite) = self.sqlite.as_ref() { sqlx::query( "INSERT INTO urls (url, timestamp) VALUES ($1, CURRENT_TIMESTAMP)", ) .bind(url) .execute(sqlite) .await .unwrap(); } } return Out::Ok; } Query::CheckForUrl(ref url) => { let res: (i64,) = if let Some(postgres) = self.postgres.as_ref() { sqlx::query_as("SELECT COUNT(*) FROM urls WHERE url = $1") .bind(url) .fetch_one(postgres) .await .unwrap() } else { sqlx::query_as("SELECT COUNT(*) FROM urls WHERE url = $1") .bind(url) .fetch_one(self.sqlite.as_ref().unwrap()) .await .unwrap() }; let count: i64 = res.0; return Out::Bool(count > 0); } Query::UpdateNewDownloads(ref module, ref name, ref url) => { if let Some(postgres) = self.postgres.as_ref() { sqlx::query( r#" INSERT INTO item_log (module, name, url, CURRENT_TIMESTAMP) VALUES ($1, $2, $3) ON CONFLICT (module, name, url) DO UPDATE SET timestamp = CURRENT_TIMESTAMP "#, ) .bind(module) .bind(name) .bind(url) .execute(postgres) .await .unwrap(); } else { if let Some(sqlite) = self.sqlite.as_ref() { sqlx::query( r#" INSERT INTO item_log (module, name, url, timestamp) VALUES ($1, $2, $3, CURRENT_TIMESTAMP) ON CONFLICT (module, name, url) DO UPDATE SET timestamp = CURRENT_TIMESTAMP "#, ) .bind(module) .bind(name) .bind(url) .execute(sqlite) .await .unwrap(); } } return Out::Ok; } } } } pub enum Query { InsertUrl(String), CheckForUrl(String), UpdateNewDownloads(String, String, String), } pub enum Out { Ok, Bool(bool), // Rows(Vec<String>), } #[derive(Clone)] pub struct Database { conn: DatabaseBackend, } impl Database { pub fn new(conn: DatabaseBackend) -> Self { Self { conn } } /// Insert a URL into the database as already downloaded pub fn insert_url(&self, url: &str) { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { self.conn.query(Query::InsertUrl(url.to_string())).await; }); } /// Check if a URL is already in the database /// /// # Return /// Returns `true` if already present, `false` otherwise /// /// # Example /// You could use this function like that: /// /// ```rust /// if !db.check_for_url(some_url) { /// // do download /// } /// ``` pub fn check_for_url(&self, url: &str) -> bool { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { match self.conn.query(Query::CheckForUrl(url.to_string())).await { Out::Ok => false, Out::Bool(b) => b, } }) } /// Keep a record on when download happen. /// This takes a `module`, `name` and `url` and saves a timestamp to the db. pub fn update_new_downloads(&self, module: &str, name: &str, url: &str) { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { self.conn .query(Query::UpdateNewDownloads( module.to_string(), name.to_string(), url.to_string(), )) .await; }); } }