Merge pull request #44 from wobcom/refactor/vickyctl

refactor for vickyctl

Author: Kek5chen

Diff for: Cargo.lock

@@ -80,16 +80,16 @@ pub struct FlakeRef {
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(tag = "result")]
 pub enum TaskResult {
-    SUCCESS,
-    ERROR,
+    Success,
+    Error,
 }
 
 #[derive(Debug, Deserialize)]
 #[serde(tag = "state")]
 pub enum TaskStatus {
-    NEW,
-    RUNNING,
-    FINISHED(TaskResult),
+    New,
+    Running,
+    Finished(TaskResult),
 }
 
 #[derive(Debug, Deserialize)]
@@ -173,9 +173,9 @@ async fn run_task(cfg: Arc<AppConfig>, task: Task) {
     let result = match try_run_task(cfg.clone(), &task).await {
         Err(e) => {
             log::info!("task failed: {} {} {:?}", task.id, task.display_name, e);
-            TaskResult::ERROR
+            TaskResult::Error
         }
-        Ok(_) => TaskResult::SUCCESS,
+        Ok(_) => TaskResult::Success,
     };
     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
     let _ = api::<_, ()>(

Diff for: fairy/src/main.rs

@@ -1,8 +1,10 @@
 use diesel::PgConnection;
-use rocket::get;
+use rocket::{get, patch};
 use rocket::serde::json::Json;
+use uuid::Uuid;
 use vickylib::database::entities::{Database, Lock};
 use vickylib::database::entities::lock::db_impl::LockDatabase;
+use vickylib::database::entities::lock::PoisonedLock;
 use crate::auth::{Machine, User};
 use crate::errors::AppError;
 
@@ -11,6 +13,11 @@ async fn locks_get_poisoned(db: &Database) -> Result<Json<Vec<Lock>>, AppError>
     Ok(Json(poisoned_locks))
 }
 
+async fn locks_get_detailed_poisoned(db: &Database) -> Result<Json<Vec<PoisonedLock>>, AppError> {
+    let poisoned_locks: Vec<PoisonedLock> = db.run(PgConnection::get_poisoned_locks_with_tasks).await?;
+    Ok(Json(poisoned_locks))
+}
+
 #[get("/poisoned")]
 pub async fn locks_get_poisoned_user(
     db: Database,
@@ -27,6 +34,23 @@ pub async fn locks_get_poisoned_machine(
     locks_get_poisoned(&db).await
 }
 
+
+#[get("/poisoned_detailed")]
+pub async fn locks_get_detailed_poisoned_user(
+    db: Database,
+    _user: User,
+) -> Result<Json<Vec<PoisonedLock>>, AppError> {
+    locks_get_detailed_poisoned(&db).await
+}
+
+#[get("/poisoned_detailed", rank = 2)]
+pub async fn locks_get_detailed_poisoned_machine(
+    db: Database,
+    _machine: Machine,
+) -> Result<Json<Vec<PoisonedLock>>, AppError> {
+    locks_get_detailed_poisoned(&db).await
+}
+
 async fn locks_get_active(db: &Database) -> Result<Json<Vec<Lock>>, AppError> {
     let locks: Vec<Lock> = db.run(PgConnection::get_active_locks).await?;
     Ok(Json(locks))
@@ -47,3 +71,15 @@ pub async fn locks_get_active_machine(
 ) -> Result<Json<Vec<Lock>>, AppError> {
     locks_get_active(&db).await
 }
+
+#[patch("/unlock/<lock_id>")]
+pub async fn locks_unlock(
+    db: Database,
+    _user: Machine, // TODO: Should actually be user-only, but we don't have that yet
+    lock_id: String,
+) -> Result<(), AppError> {
+    let lock_uuid = Uuid::try_parse(&lock_id)?;
+    
+    db.run(move |conn| conn.unlock_lock(&lock_uuid)).await?;
+    Ok(())
+}

Diff for: vicky/src/bin/vicky/locks.rs

@@ -16,7 +16,11 @@ use vickylib::logs::LogDrain;
 use vickylib::s3::client::S3Client;
 
 use crate::events::{get_global_events, GlobalEvent};
-use crate::locks::{locks_get_active_machine, locks_get_active_user, locks_get_poisoned_machine, locks_get_poisoned_user};
+use crate::locks::{
+    locks_get_active_machine, locks_get_active_user, locks_get_detailed_poisoned_machine,
+    locks_get_detailed_poisoned_user, locks_get_poisoned_machine, locks_get_poisoned_user,
+    locks_unlock,
+};
 use crate::tasks::{
     tasks_add, tasks_claim, tasks_finish, tasks_get_logs, tasks_get_machine, tasks_get_user,
     tasks_put_logs, tasks_specific_get_machine, tasks_specific_get_user,
@@ -77,19 +81,19 @@ fn run_migrations(connection: &mut impl MigrationHarness<diesel::pg::Pg>) -> Res
         Ok(_) => {
             log::info!("Migrations successfully completed");
             Ok(())
-        },
+        }
         Err(e) => {
             log::error!("Error running migrations {e}");
             Err(AppError::MigrationError(e.to_string()))
         }
     }
 }
 
-async fn run_rocket_migrations(rocket: Rocket<Build>) -> Result<Rocket<Build>,Rocket<Build>> {
+async fn run_rocket_migrations(rocket: Rocket<Build>) -> Result<Rocket<Build>, Rocket<Build>> {
     let db: Database = Database::get_one(&rocket).await.unwrap();
     match db.run(run_migrations).await {
         Ok(_) => Ok(rocket),
-        Err(_) => Err(rocket)
+        Err(_) => Err(rocket),
     }
 }
 
@@ -117,12 +121,16 @@ async fn main() -> anyhow::Result<()> {
 
     let app_config = build_rocket.figment().extract::<Config>()?;
 
-    let oidc_config_resolved: OIDCConfigResolved = reqwest::get(app_config.oidc_config.well_known_uri)
-        .await?
-        .json()
-        .await?;
+    let oidc_config_resolved: OIDCConfigResolved =
+        reqwest::get(app_config.oidc_config.well_known_uri)
+            .await?
+            .json()
+            .await?;
 
-    log::info!("Fetched OIDC configuration, found jwks_uri={}", oidc_config_resolved.jwks_uri);
+    log::info!(
+        "Fetched OIDC configuration, found jwks_uri={}",
+        oidc_config_resolved.jwks_uri
+    );
 
     let jwks_verifier = RemoteJwksVerifier::new(
         oidc_config_resolved.jwks_uri.clone(),
@@ -162,7 +170,10 @@ async fn main() -> anyhow::Result<()> {
         .manage(oidc_config_resolved)
         .attach(Database::fairing())
         .attach(AdHoc::config::<Config>())
-        .attach(AdHoc::try_on_ignite("run migrations", run_rocket_migrations))
+        .attach(AdHoc::try_on_ignite(
+            "run migrations",
+            run_rocket_migrations,
+        ))
         .mount("/api/v1/web-config", routes![get_web_config])
         .mount("/api/v1/user", routes![get_user])
         .mount("/api/v1/events", routes![get_global_events])
@@ -185,8 +196,11 @@ async fn main() -> anyhow::Result<()> {
             routes![
                 locks_get_poisoned_user,
                 locks_get_poisoned_machine,
+                locks_get_detailed_poisoned_user,
+                locks_get_detailed_poisoned_machine,
                 locks_get_active_user,
                 locks_get_active_machine,
+                locks_unlock
             ],
         )
         .launch()

Diff for: vicky/src/bin/vicky/main.rs

@@ -107,8 +107,8 @@ pub async fn tasks_get_logs<'a>(
     EventStream! {
 
         match task.status {
-            TaskStatus::NEW => {},
-            TaskStatus::RUNNING => {
+            TaskStatus::New => {},
+            TaskStatus::Running => {
                 let mut recv = log_drain.send_handle.subscribe();
                 let existing_log_messages = log_drain.get_logs(task_uuid.to_string()).await.unwrap();
 
@@ -137,7 +137,7 @@ pub async fn tasks_get_logs<'a>(
                     }
                 }
             },
-            TaskStatus::FINISHED(_) => {
+            TaskStatus::Finished(_) => {
                 let logs = s3.get_logs(&id).await.unwrap();
                 for element in logs {
                     yield Event::data(element)
@@ -166,7 +166,7 @@ pub async fn tasks_put_logs(
         .ok_or(AppError::HttpError(Status::NotFound))?;
 
     match task.status {
-        TaskStatus::RUNNING => {
+        TaskStatus::Running => {
             log_drain.push_logs(id, logs.lines.clone())?;
             Ok(Json(()))
         }
@@ -193,7 +193,7 @@ pub async fn tasks_claim(
                 .run(move |conn| conn.get_task(next_task.id))
                 .await?
                 .ok_or(AppError::HttpError(Status::NotFound))?;
-            task.status = TaskStatus::RUNNING;
+            task.status = TaskStatus::Running;
             let task2 = task.clone();
             db.run(move |conn| conn.update_task(&task2)).await?;
             global_events.send(GlobalEvent::TaskUpdate { uuid: task.id })?;
@@ -220,9 +220,9 @@ pub async fn tasks_finish(
 
     log_drain.finish_logs(&id).await?;
 
-    task.status = TaskStatus::FINISHED(finish.result.clone());
+    task.status = TaskStatus::Finished(finish.result.clone());
 
-    if finish.result == TaskResult::ERROR {
+    if finish.result == TaskResult::Error {
         task.locks.iter_mut().for_each(|lock| lock.poison(&task.id));
     }
 
@@ -253,7 +253,7 @@ pub async fn tasks_add(
 ) -> Result<Json<RoTask>, AppError> {
     let task_uuid = Uuid::new_v4();
 
-    let task_manifest = Task::builder()
+    let task = Task::builder()
         .with_id(task_uuid)
         .with_display_name(&task.display_name)
         .with_flake(&task.flake_ref.flake)
@@ -262,16 +262,16 @@ pub async fn tasks_add(
         .requires_features(task.features.clone())
         .build();
 
-    if check_lock_conflict(&task_manifest) {
+    if check_lock_conflict(&task) {
         return Err(AppError::HttpError(Status::Conflict));
     }
 
-    db.run(move |conn| conn.put_task(&task_manifest)).await?;
+    db.run(move |conn| conn.put_task(task)).await?;
     global_events.send(GlobalEvent::TaskAdd)?;
 
     let ro_task = RoTask {
         id: task_uuid,
-        status: TaskStatus::NEW,
+        status: TaskStatus::New,
     };
 
     Ok(Json(ro_task))