feat.: Impl. auto user insertion

Signed-off-by: Dheshan Mohandass <dmohandass@mgh.harvard.edu>

Author: dheshanm

src/bin/diskusage/main.rs

@@ -4,6 +4,7 @@ mod users;
 
 use clap::Parser;
 use rayon::prelude::*;
+use crate::models::definitions::DbModel;
 
 #[derive(clap::Parser, Default, Debug)]
 #[clap(
@@ -38,6 +39,29 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
  return Err("Root directory does not exist".into());
  }
 
+ // get the database url from the environment
+ let database_url = match std::env::var("DATABASE_URL") {
+ Ok(url) => url,
+ Err(_) => {
+ log::error!("DATABASE_URL environment variable not set.");
+ return Err("DATABASE_URL environment variable not set.".into());
+ }
+ };
+
+ let pool_result = sqlx::postgres::PgPoolOptions::new().connect(&database_url).await;
+ let handle = tokio::runtime::Handle::current();
+
+ let pool = match pool_result {
+ Ok(pool) => {
+ log::info!("Connected to database: {}", database_url);
+ pool
+ }
+ Err(e) => {
+ log::error!("Failed to connect to database: {}", e);
+ return Err(e.into());
+ }
+ };
+
  walkdir::WalkDir::new(root_dir)
  .into_iter()
  .par_bridge() // Allows rayon to process entries in parallel
@@ -46,31 +70,47 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
  // Check if file or directory
  if entry.file_type().is_dir() {
  let dir_path = entry.path();
- let owner = filesystem::fetch::owner(dir_path).unwrap_or_default();
+ let owner: Option<u32> = filesystem::fetch::owner(dir_path);
+ let owner_i32: Option<i32> = owner.map(|x| x as i32);
  let parent_dir = dir_path.parent().unwrap_or(std::path::Path::new("/") );
+
+ let handle = handle.clone();
+ let select_user_where_clause = format!("WHERE user_id = {}", owner_i32.unwrap_or_default());
+ let pool_c = pool.clone();
+ handle.block_on(async move {
+ let user = models::definitions::User::select_where(&pool_c, &select_user_where_clause).await.unwrap_or_default();
+ if user.is_empty() {
+ let user = models::definitions::User {
+ user_id: owner_i32.unwrap_or_default(),
+ username: users::username::get_username(owner.unwrap_or_default())
+ };
+ user.insert(&pool_c).await.unwrap_or_default();
+ }
+ });
 
  let directory: models::definitions::Directory = models::definitions::Directory {
  directory_id: dir_path.to_string_lossy().to_string(),
- owner_id: owner as i32,
+ owner_id: owner_i32,
  parent_id: Some(parent_dir.to_string_lossy().to_string()),
  };
 
- log::info!("{:?}", directory);
+ log::info!("directory: {:?}", directory);
  } else {
  let file_path = entry.path();
- let owner = filesystem::fetch::owner(file_path).unwrap_or_default();
+ let owner: Option<u32> = filesystem::fetch::owner(file_path);
+ let owner_i32: Option<i32> = owner.map(|x| x as i32);
  let file_size = filesystem::fetch::file_size(file_path).unwrap_or_default();
 
  let parent_dir = file_path.parent().unwrap_or(std::path::Path::new("/") );
  let file: models::definitions::File = models::definitions::File {
  file_id: file_path.to_string_lossy().to_string(),
  name: file_path.file_name().unwrap().to_string_lossy().to_string(),
  size: file_size as i64,
- owner_id: owner as i32,
+ owner_id: owner_i32.unwrap_or_default(),
  directory_id: parent_dir.to_string_lossy().to_string(),
  };
 
- log::info!("{:?}", file);
+ log::info!("file: {:?}", file);
  }
  });
 

src/bin/diskusage/models/definitions.rs

@@ -11,7 +11,7 @@ pub struct User {
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 pub struct Directory {
  pub directory_id: String,
- pub owner_id: i32,
+ pub owner_id: Option<i32>,
  pub parent_id: Option<String>,
 }
 
@@ -112,3 +112,80 @@ impl DbModel for User {
  Ok(users)
  }
 }
+
+#[async_trait::async_trait]
+impl DbModel for Directory {
+ async fn insert(&self, pool: &sqlx::Pool<sqlx::postgres::Postgres>) -> Result<(), sqlx::Error> {
+ sqlx::query!(
+ "INSERT INTO directories (directory_id, owner_id, parent_id) VALUES ($1, $2, $3)",
+ self.directory_id,
+ self.owner_id,
+ self.parent_id
+ )
+ .execute(pool)
+ .await?;
+ Ok(())
+ }
+
+ async fn update(&self, pool: &sqlx::Pool<sqlx::postgres::Postgres>) -> Result<(), sqlx::Error> {
+ sqlx::query!(
+ "UPDATE directories SET owner_id = $2, parent_id = $3 WHERE directory_id = $1",
+ self.directory_id,
+ self.owner_id,
+ self.parent_id
+ )
+ .execute(pool)
+ .await?;
+ Ok(())
+ }
+
+ async fn delete(&self, pool: &sqlx::Pool<sqlx::postgres::Postgres>) -> Result<(), sqlx::Error> {
+ sqlx::query!(
+ "DELETE FROM directories WHERE directory_id = $1",
+ self.directory_id
+ )
+ .execute(pool)
+ .await?;
+ Ok(())
+ }
+
+ async fn select(&self, pool: &sqlx::Pool<sqlx::postgres::Postgres>) -> Result<Box<Self>, sqlx::Error> {
+ let directory = sqlx::query_as!(
+ Directory,
+ "SELECT directory_id, owner_id, parent_id FROM directories WHERE directory_id = $1",
+ self.directory_id
+ )
+ .fetch_one(pool)
+ .await?;
+ Ok(Box::new(directory))
+ }
+
+ async fn select_all(pool: &sqlx::Pool<sqlx::postgres::Postgres>) -> Result<Vec<Box<Self>>, sqlx::Error> {
+ let directories = sqlx::query_as!(
+ Directory,
+ "SELECT directory_id, owner_id, parent_id FROM directories"
+ )
+ .fetch_all(pool)
+ .await?;
+ Ok(directories.into_iter().map(|d| Box::new(d)).collect())
+ }
+
+ async fn select_where(pool: &sqlx::Pool<sqlx::postgres::Postgres>, where_clause: &str) -> Result<Vec<Box<Directory>>, sqlx::Error> {
+ let query_string = format!("SELECT directory_id, owner_id, parent_id FROM directories WHERE {}", where_clause);
+ let rows = sqlx::query(&query_string)
+ .fetch_all(pool)
+ .await?;
+
+ let directories: Vec<Box<Directory>> = rows.into_iter()
+ .map(|row| {
+ Box::new(Directory {
+ directory_id: row.get("directory_id"),
+ owner_id: row.get("owner_id"),
+ parent_id: row.get("parent_id"),
+ })
+ })
+ .collect();
+
+ Ok(directories)
+ }
+}

src/bin/init_db/main.rs

@@ -26,8 +26,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
  let debug = args.debug;
 
- let database_url = "postgres://pipeline:piedpiper@localhost:5432/test";
- let pool_result = sqlx::postgres::PgPoolOptions::new().connect(database_url).await;
+ // get the database url from the environment
+ let database_url = match std::env::var("DATABASE_URL") {
+ Ok(url) => url,
+ Err(_) => {
+ log::error!("DATABASE_URL environment variable not set.");
+ return Err("DATABASE_URL environment variable not set.".into());
+ }
+ };
+
+ let pool_result = sqlx::postgres::PgPoolOptions::new().connect(&database_url).await;
 
  let pool = match pool_result {
  Ok(pool) => {