hoard/src/db.rs

#[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;
        });
    }
}