geo-engine / geoengine / 12767614094

pub(crate) fn init_task_routes<C>(cfg: &mut web::ServiceConfig)

where

    C: ApplicationContext,

    C::Session: FromRequest,

{

    cfg.service(

        web::scope("/tasks")

            .service(web::resource("/list").route(web::get().to(list_handler::<C>)))

            .service(

                web::scope("/{task_id}")

                    .service(web::resource("/status").route(web::get().to(status_handler::<C>)))

                    .service(web::resource("").route(web::delete().to(abort_handler::<C>))),

            ),

    );

}

#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]

    pub fn new(task_id: TaskId) -> Self {

        Self { task_id }

    }

#[utoipa::path(

    tag = "Tasks",

    get,

    path = "/tasks/{id}/status",

    responses(

        (status = 200, description = "Status of the task (running)", body = TaskStatus,

            example = json!({

                "status": "running",

                "taskType": "dummy-task",

                "description": "Demo",

                "pctComplete": "0.00%",

                "timeStarted": "2023-02-16T15:25:45.390Z",

                "estimatedTimeRemaining": "? (± ?)",

                "info": null,

            })

        )

    ),

    params(

        ("id" = TaskId, description = "Task id")

    ),

    security(

        ("session_token" = [])

    )

)]

async fn status_handler<C: ApplicationContext>(

    session: C::Session, // TODO: check for session auth

    app_ctx: web::Data<C>,

    task_id: web::Path<TaskId>,

) -> Result<impl Responder> {

    let task_id = task_id.into_inner();

    let task = app_ctx

        .session_context(session)

        .tasks()

        .get_task_status(task_id)

        .await?;

    Ok(web::Json(task))

}

#[utoipa::path(

    tag = "Tasks",

    get,

    path = "/tasks/list",

    responses(

        (status = 200, description = "Status of all tasks", body = Vec<TaskStatusWithId>,

            example = json!([

                {

                    "taskId": "420b06de-0a7e-45cb-9c1c-ea901b46ab69",

                    "status": "completed",

                    "taskType": "dummy-task",

                    "description": "Demo",

                    "info": null,

                    "timeTotal": "00:00:30",

                    "timeStarted": "2023-02-16T15:25:45.390Z"

                }

            ])

        )

    ),

    params(

        TaskListOptions

    ),

    security(

        ("session_token" = [])

    )

)]

async fn list_handler<C: ApplicationContext>(

    session: C::Session,

    app_ctx: web::Data<C>,

    task_list_options: ValidatedQuery<TaskListOptions>,

) -> Result<web::Json<Vec<TaskStatusWithId>>> {

    let task_list_options = task_list_options.into_inner();

    let task = app_ctx

        .session_context(session)

        .tasks()

        .list_tasks(task_list_options)

        .await?;

    Ok(web::Json(task))

}

#[derive(Debug, Serialize, Deserialize, Clone, ToSchema, IntoParams)]

#[utoipa::path(

    tag = "Tasks",

    delete,

    path = "/tasks/{id}",

    responses(

        (status = 202, description = "Task will be aborted.")

    ),

    params(

        TaskAbortOptions,

        ("id" = TaskId, description = "Task id")

    ),

    security(

        ("session_token" = [])

    )

)]

async fn abort_handler<C: ApplicationContext>(

    session: C::Session, // TODO: check for session auth

    app_ctx: web::Data<C>,

    task_id: web::Path<TaskId>,

    options: web::Query<TaskAbortOptions>,

) -> Result<HttpResponse> {

    let task_id = task_id.into_inner();

    app_ctx

        .session_context(session)

        .tasks()

        .abort_tasks(task_id, options.force)

        .await?;

    Ok(HttpResponse::Accepted().finish())

}

        pub fn new_with_sender() -> (Self, oneshot::Sender<()>) {

            let (complete_tx, complete_rx) = oneshot::channel();

            let this = Self {

                complete_rx: Arc::new(Mutex::new(complete_rx)),

                unique_id: None,

            };

            (this, complete_tx)

        }

        pub fn new_with_sender_and_unique_id(unique_id: String) -> (Self, oneshot::Sender<()>) {

            let (complete_tx, complete_rx) = oneshot::channel();

            let this = Self {

                complete_rx: Arc::new(Mutex::new(complete_rx)),

                unique_id: Some(unique_id),

            };

            (this, complete_tx)

        }

        async fn run(&self, _ctx: C) -> Result<Box<dyn TaskStatusInfo>, Box<dyn ErrorSource>> {

            let mut complete_rx_lock = self.complete_rx.lock().await;

            let pinned_receiver: Pin<&mut oneshot::Receiver<()>> = Pin::new(&mut complete_rx_lock);

            pinned_receiver.await.unwrap();

            Ok("completed".to_string().boxed())

        }

        async fn cleanup_on_error(&self, _ctx: C) -> Result<(), Box<dyn ErrorSource>> {

            let mut complete_rx_lock = self.complete_rx.lock().await;

            let pinned_receiver: Pin<&mut oneshot::Receiver<()>> = Pin::new(&mut complete_rx_lock);

            pinned_receiver.await.unwrap();

            Ok(())

        }

        fn task_type(&self) -> &'static str {

            "nopTask"

        }

        fn task_unique_id(&self) -> Option<String> {

            self.unique_id.clone()

        }

        fn task_description(&self) -> String {

            "No operation".to_string()

        }

        pub fn new_with_sender(

            subtasks: Vec<Box<dyn Task<C>>>,

            task_manager: Arc<T>,

        ) -> (Self, oneshot::Sender<()>) {

            let task_description = subtasks

                .iter()

                .map(|subtask| subtask.task_description())

                .collect::<Vec<_>>()

                .join(", ");

            let (complete_tx, complete_rx) = oneshot::channel();

            let this = Self {

                subtasks: Arc::new(Mutex::new(subtasks)),

                subtask_ids: Arc::new(Mutex::new(vec![])),

                task_manager,

                complete_rx: Arc::new(Mutex::new(complete_rx)),

                task_description,

            };

            (this, complete_tx)

        }

        async fn run(&self, _ctx: C) -> Result<Box<dyn TaskStatusInfo>, Box<dyn ErrorSource>> {

            for subtask in self.subtasks.lock().await.drain(..) {

                let subtask_id = self

                    .task_manager

                    .schedule_task(subtask, None)

                    .await

                    .map_err(ErrorSource::boxed)?;

                self.subtask_ids.lock().await.push(subtask_id);

            let mut complete_rx_lock = self.complete_rx.lock().await;

            let pinned_receiver: Pin<&mut oneshot::Receiver<()>> = Pin::new(&mut complete_rx_lock);

            pinned_receiver.await.unwrap();

        }

        async fn cleanup_on_error(&self, _ctx: C) -> Result<(), Box<dyn ErrorSource>> {

            Ok(())

        }

        fn task_type(&self) -> &'static str {

            stringify!(TaskTree)

        }

        fn task_unique_id(&self) -> Option<String> {

            None

        }

        fn task_description(&self) -> String {

            self.task_description.clone()

        }

        async fn subtasks(&self) -> Vec<TaskId> {

            self.subtask_ids.lock().await.clone()

        }

        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

            write!(f, "FailingTaskWithFailingCleanupError")

        }

        async fn run(&self, _ctx: C) -> Result<Box<dyn TaskStatusInfo>, Box<dyn ErrorSource>> {

            Err(Box::new(FailingTaskWithFailingCleanupError))

        }

        async fn cleanup_on_error(&self, _ctx: C) -> Result<(), Box<dyn ErrorSource>> {

            Err(Box::new(FailingTaskWithFailingCleanupError))

        }

        fn task_type(&self) -> &'static str {

            "FailingTaskWithFailingCleanup"

        }

        fn task_unique_id(&self) -> Option<String> {

            None

        }

        fn task_description(&self) -> String {

            "Failing task with failing cleanup".to_string()

        }

    #[ge_context::test]

    async fn test_get_status(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let (task, complete_tx) = NopTask::new_with_sender();

        let tasks = Arc::new(ctx.tasks());

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        // 1. initially, we should get a running status

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        let res_status = res.status();

        let res_body = read_body_json(res).await;

        assert_eq!(res_status, 200, "{res_body:?}");

        assert_eq!(res_body["status"], json!("running"));

        assert_eq!(res_body["pctComplete"], json!("0.00%"));

        assert!(res_body["info"].is_null());

        assert_eq!(res_body["estimatedTimeRemaining"], json!("? (± ?)"));

        assert!(res_body["timeStarted"].is_string());

        complete_tx.send(()).unwrap();

        wait_for_task_to_finish(tasks, task_id).await;

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let app_ctx_clone = app_ctx.clone();

        let res = send_test_request(req, app_ctx_clone).await;

        let res_status = res.status();

        let res_body = read_body_json(res).await;

        assert_eq!(res_status, 200, "{res_body:?}");

        assert_eq!(res_body["status"], json!("completed"));

        assert_eq!(res_body["info"], json!("completed"));

        assert_eq!(res_body["timeTotal"], json!("00:00:00"));

        assert!(res_body["timeStarted"].is_string());

    }

    #[ge_context::test]

    async fn test_get_status_with_admin_session(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let (task, _complete_tx) = NopTask::new_with_sender();

        let task_id = ctx.tasks().schedule_task(task.boxed(), None).await.unwrap();

        // 1. initially, we should get a running status

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        let res_status = res.status();

        let res_body = read_body_json(res).await;

        assert_eq!(res_status, 200, "{res_body:?}");

        assert_eq!(res_body["status"], json!("running"));

        assert_eq!(res_body["pctComplete"], json!("0.00%"));

        assert!(res_body["info"].is_null());

        assert_eq!(res_body["estimatedTimeRemaining"], json!("? (± ?)"));

        assert!(res_body["timeStarted"].is_string());

    }

    #[ge_context::test]

    async fn test_get_list(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let tasks = Arc::new(ctx.tasks());

        let (task, complete_tx) = NopTask::new_with_sender();

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        complete_tx.send(()).unwrap();

        wait_for_task_to_finish(tasks, task_id).await;

        let req = actix_web::test::TestRequest::get()

            .uri("/tasks/list")

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        let res_status = res.status();

        let res_body = read_body_json(res).await;

        assert_eq!(res_status, 200, "{res_body:?}");

        let res_body = res_body.get(0).unwrap();

        assert_eq!(res_body["taskId"], json!(task_id));

        assert_eq!(res_body["status"], json!("completed"));

        assert_eq!(res_body["info"], json!("completed"));

        assert_eq!(res_body["timeTotal"], json!("00:00:00"));

        assert!(res_body["timeStarted"].is_string());

    }

    #[ge_context::test]

    async fn test_abort_task(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let tasks = Arc::new(ctx.tasks());

        // 1. Create task

        let (task, complete_tx) = NopTask::new_with_sender();

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        // 2. Abort task

        let req = actix_web::test::TestRequest::delete()

            .uri(&format!("/tasks/{task_id}"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 202, "{:?}", res.response().error());

        complete_tx.send(()).unwrap();

        wait_for_task_to_finish(tasks, task_id).await;

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 200, "{:?}", res.response().error());

        let status = read_body_json(res).await;

        assert_eq!(

            status,

            json!({

                "status": "aborted",

                "cleanUp": {

                    "status": "completed",

                    "info": null

                },

            })

        );

    }

    #[ge_context::test]

    async fn test_force_abort_task(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let tasks = Arc::new(ctx.tasks());

        // 1. Create task

        let (task, _complete_tx) = NopTask::new_with_sender();

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        // 2. Abort task

        let req = actix_web::test::TestRequest::delete()

            .uri(&format!("/tasks/{task_id}?force=true"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 202, "{:?}", res.response().error());

        wait_for_task_to_finish(tasks, task_id).await;

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 200, "{:?}", res.response().error());

        let status = read_body_json(res).await;

        assert_eq!(

            status,

            json!({

                "status": "aborted",

                "cleanUp": {

                    "status": "noCleanUp"

                },

            })

        );

    }

    #[ge_context::test]

    async fn test_abort_after_abort(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let session_id = session.id();

        let tasks = Arc::new(ctx.tasks());

        // 1. Create task

        let (task, _complete_tx) = NopTask::new_with_sender();

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        // 2. Abort task

        let req = actix_web::test::TestRequest::delete()

            .uri(&format!("/tasks/{task_id}"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 202, "{:?}", res.response().error());

        let req = actix_web::test::TestRequest::delete()

            .uri(&format!("/tasks/{task_id}"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 400, "{:?}", res.response().error());

        let status = read_body_json(res).await;

        assert_eq!(

            status,

            json!({

                "error": "Task",

                "message": format!("TaskError: Task was already aborted by the user: {task_id}"),

            })

        );

        let req = actix_web::test::TestRequest::delete()

            .uri(&format!("/tasks/{task_id}?force=true"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 202, "{:?}", res.response().error());

        let req = actix_web::test::TestRequest::get()

            .uri(&format!("/tasks/{task_id}/status"))

            .append_header((header::CONTENT_LENGTH, 0))

            .append_header((header::AUTHORIZATION, Bearer::new(session_id.to_string())));

        let res = send_test_request(req, app_ctx.clone()).await;

        assert_eq!(res.status(), 200, "{:?}", res.response().error());

        let status = read_body_json(res).await;

        assert_eq!(

            status,

            json!({

                "status": "aborted",

                "cleanUp": {

                    "status": "noCleanUp"

                },

            })

        );

    }

    #[ge_context::test]

    async fn test_duplicate(app_ctx: ProPostgresContext<NoTls>) {

        let unique_id = "highlander".to_string();

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let tasks = app_ctx.session_context(session).tasks();

        let (task, complete_tx) = NopTask::new_with_sender_and_unique_id(unique_id.clone());

        tasks.schedule_task(task.boxed(), None).await.unwrap();

        let (task, _) = NopTask::new_with_sender_and_unique_id(unique_id.clone());

        assert!(tasks.schedule_task(task.boxed(), None).await.is_err());

        complete_tx.send(()).unwrap();

    }

    #[ge_context::test]

    async fn test_duplicate_after_finish(app_ctx: ProPostgresContext<NoTls>) {

        let unique_id = "highlander".to_string();

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let tasks = Arc::new(ctx.tasks());

        // 1. start first task

        let (task, complete_tx) = NopTask::new_with_sender_and_unique_id(unique_id.clone());

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        // 2. wait for task to finish

        complete_tx.send(()).unwrap();

        wait_for_task_to_finish(tasks.clone(), task_id).await;

        let (task, complete_tx) = NopTask::new_with_sender_and_unique_id(unique_id.clone());

        assert!(tasks.schedule_task(task.boxed(), None).await.is_ok());

        complete_tx.send(()).unwrap();

    }

    #[ge_context::test]

    async fn test_notify(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let tasks = app_ctx.session_context(session).tasks();

        // 1. start first task

        let (schedule_complete_tx, schedule_complete_rx) = oneshot::channel();

        let (task, complete_tx) = NopTask::new_with_sender();

        tasks

            .schedule_task(task.boxed(), Some(schedule_complete_tx))

            .await

            .unwrap();

        // finish task

        complete_tx.send(()).unwrap();

        // wait for completion notification

        assert!(matches!(

            schedule_complete_rx.await.unwrap(),

    }

    #[ge_context::test]

    async fn abort_subtasks(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let tasks = Arc::new(ctx.tasks());

        // 1. start super task

        let (subtask_a, complete_tx_a) = NopTask::new_with_sender();

        let (subtask_b, complete_tx_b) = NopTask::new_with_sender();

        let (task, _complete_tx) =

            TaskTree::new_with_sender(vec![subtask_a.boxed(), subtask_b.boxed()], tasks.clone());

        let task_id = tasks.schedule_task(task.boxed(), None).await.unwrap();

        let all_task_ids: Vec<TaskId> =

            geoengine_operators::util::retry::retry(5, 100, 2., None, || {

                let task_manager = tasks.clone();

                async move {

                    let task_list = task_manager

                        .list_tasks(TaskListOptions {

                            filter: None,

                            offset: 0,

                            limit: 10,

                        })

                        .await

                        .unwrap();

                    if task_list.len() == 3 {

                        Ok(task_list.into_iter().map(|t| t.task_id).collect())

                }

            })

            .await

            .unwrap();

        // 3. abort task

        tasks.abort_tasks(task_id, false).await.unwrap();

        // allow clean-up to complete

        complete_tx_a.send(()).unwrap();

        complete_tx_b.send(()).unwrap();

        for task_id in all_task_ids {

            wait_for_task_to_finish(tasks.clone(), task_id).await;

        let list = tasks

            .list_tasks(TaskListOptions {

                filter: None,

                offset: 0,

                limit: 10,

            })

            .await

            .unwrap();

        assert_eq!(list.len(), 3);

        assert_eq!(

            serde_json::to_value(&list[0].status).unwrap(),

            json!({

                "status": "aborted",

                "cleanUp": {"status": "completed", "info": null}

            })

        );

        assert_eq!(

            serde_json::to_value(&list[1].status).unwrap(),

            json!({

                "status": "aborted",

                "cleanUp": {"status": "completed", "info": null}

            })

        );

        assert_eq!(

            serde_json::to_value(&list[2].status).unwrap(),

            json!({

                "status": "aborted",

                "cleanUp": {"status": "completed", "info": null}

            })

        );

    }

    #[ge_context::test]

    async fn test_failing_task_with_failing_cleanup(app_ctx: ProPostgresContext<NoTls>) {

        let session = app_ctx.create_anonymous_session().await.unwrap();

        let ctx = app_ctx.session_context(session.clone());

        let tasks = Arc::new(ctx.tasks());

        let task_id = tasks

            .schedule_task(FailingTaskWithFailingCleanup.boxed(), None)

            .await

            .unwrap();

        // 2. wait for completion

        wait_for_task_to_finish(tasks.clone(), task_id).await;

        assert_eq!(

            serde_json::to_value(tasks.get_task_status(task_id).await.unwrap()).unwrap(),

            json!({

                "status": "failed",

                "error": "FailingTaskWithFailingCleanupError",

                "cleanUp": {"status": "failed", "error": "FailingTaskWithFailingCleanupError"}

            })

    }