Merge pull request #131 from infraweave-io/fix/add-job-id-verification-check

InfiniteTabs · web-flow · commit f98032735adb · 2025-10-19T12:38:39.000+02:00
fix: add job-id verification-check
diff --git a/defs/src/deployment.rs b/defs/src/deployment.rs
@@ -90,6 +90,13 @@ pub struct DeploymentResp {
     pub reference: String,
 }
 
+#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
+#[derive(Deserialize, Clone, Debug, Serialize)]
+pub struct JobStatus {
+    pub job_id: String,
+    pub is_running: bool,
+}
+
 #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Dependency {
diff --git a/defs/src/lib.rs b/defs/src/lib.rs
@@ -19,8 +19,8 @@ mod stack;
 pub use api::GenericFunctionResponse;
 pub use deployment::{
     get_deployment_identifier, Dependency, Dependent, DeploymentManifest, DeploymentResp,
-    DeploymentSpec, DriftDetection, Metadata as DeploymentMetadata, ProjectData, Webhook,
-    DEFAULT_DRIFT_DETECTION_INTERVAL,
+    DeploymentSpec, DriftDetection, JobStatus, Metadata as DeploymentMetadata, ProjectData,
+    Webhook, DEFAULT_DRIFT_DETECTION_INTERVAL,
 };
 pub use environment::EnvironmentResp;
 pub use errors::CloudHandlerError;
diff --git a/defs/src/provider.rs b/defs/src/provider.rs
@@ -1,8 +1,8 @@
 use std::{future::Future, pin::Pin};
 
 use crate::{
-    Dependent, DeploymentResp, EventData, GenericFunctionResponse, InfraChangeRecord, LogData,
-    ModuleResp, NotificationData, PolicyResp, ProjectData,
+    deployment::JobStatus, Dependent, DeploymentResp, EventData, GenericFunctionResponse,
+    InfraChangeRecord, LogData, ModuleResp, NotificationData, PolicyResp, ProjectData,
 };
 
 use async_trait::async_trait;
@@ -116,6 +116,7 @@ pub trait CloudProvider: Send + Sync {
         environment: &str,
         include_deleted: bool,
     ) -> Result<Option<DeploymentResp>, anyhow::Error>;
+    async fn get_job_status(&self, job_id: &str) -> Result<Option<JobStatus>, anyhow::Error>;
     async fn get_deployments_using_module(
         &self,
         module: &str,
diff --git a/env_aws/src/api.rs b/env_aws/src/api.rs
@@ -304,6 +304,15 @@ pub fn get_generate_presigned_url_query(key: &str, bucket: &str) -> Value {
     })
 }
 
+pub fn get_job_status_query(job_id: &str) -> Value {
+    json!({
+        "event": "get_job_status",
+        "data": {
+            "job_id": job_id
+        }
+    })
+}
+
 pub fn get_all_deployments_query(project_id: &str, region: &str, environment: &str) -> Value {
     json!({
         "IndexName": "DeletedIndex",
diff --git a/env_aws/src/job_id.rs b/env_aws/src/job_id.rs
@@ -11,7 +11,7 @@ use std::env;
 
 pub async fn get_current_job_id() -> Result<String, anyhow::Error> {
     if std::env::var("TEST_MODE").is_ok() {
-        return Ok("test-job-id".to_string());
+        return Ok("running-test-job-id".to_string());
     };
 
     let metadata_uri = env::var("ECS_CONTAINER_METADATA_URI_V4")
diff --git a/env_aws/src/lib.rs b/env_aws/src/lib.rs
@@ -25,6 +25,7 @@ pub use api::{
     get_deployments_using_module_query,
     get_events_query,
     get_generate_presigned_url_query,
+    get_job_status_query,
     get_latest_module_version_query,
     get_latest_stack_version_query,
     get_module_version_query,
diff --git a/env_aws/src/provider.rs b/env_aws/src/provider.rs
@@ -1,7 +1,7 @@
 use async_trait::async_trait;
 use env_defs::{
     CloudHandlerError, CloudProvider, Dependent, DeploymentResp, EventData,
-    GenericFunctionResponse, InfraChangeRecord, ModuleResp, PolicyResp, ProjectData,
+    GenericFunctionResponse, InfraChangeRecord, JobStatus, ModuleResp, PolicyResp, ProjectData,
 };
 use env_utils::{
     _get_change_records, _get_dependents, _get_deployment, _get_deployment_and_dependents,
@@ -147,6 +147,22 @@ impl CloudProvider for AwsCloudProvider {
     ) -> Result<Option<ModuleResp>, anyhow::Error> {
         _get_module_optional(self, crate::get_latest_stack_version_query(stack, track)).await
     }
+    async fn get_job_status(&self, job_id: &str) -> Result<Option<JobStatus>, anyhow::Error> {
+        match crate::run_function(
+            &self.function_endpoint,
+            &crate::get_job_status_query(job_id),
+            &self.project_id,
+            &self.region,
+        )
+        .await
+        {
+            Ok(response) => {
+                let job_status: JobStatus = serde_json::from_value(response.payload)?;
+                Ok(Some(job_status))
+            }
+            Err(e) => Err(e.into()),
+        }
+    }
     async fn generate_presigned_url(
         &self,
         key: &str,
diff --git a/env_azure/src/api.rs b/env_azure/src/api.rs
@@ -176,6 +176,15 @@ pub fn get_generate_presigned_url_query(key: &str, bucket: &str) -> Value {
     })
 }
 
+pub fn get_job_status_query(job_id: &str) -> Value {
+    json!({
+        "event": "get_job_status",
+        "data": {
+            "job_id": job_id
+        }
+    })
+}
+
 pub fn get_all_latest_modules_query(track: &str) -> Value {
     _get_all_latest_modules_query("LATEST_MODULE", track)
 }
diff --git a/env_azure/src/lib.rs b/env_azure/src/lib.rs
@@ -25,6 +25,7 @@ pub use api::{
     get_deployments_using_module_query,
     get_events_query,
     get_generate_presigned_url_query,
+    get_job_status_query,
     get_latest_module_version_query,
     get_latest_stack_version_query,
     get_module_version_query,
diff --git a/env_azure/src/provider.rs b/env_azure/src/provider.rs
@@ -1,7 +1,7 @@
 use async_trait::async_trait;
 use env_defs::{
     CloudProvider, Dependent, DeploymentResp, EventData, GenericFunctionResponse,
-    InfraChangeRecord, ModuleResp, PolicyResp, ProjectData,
+    InfraChangeRecord, JobStatus, ModuleResp, PolicyResp, ProjectData,
 };
 use env_utils::{
     _get_change_records, _get_dependents, _get_deployment, _get_deployment_and_dependents,
@@ -123,6 +123,22 @@ impl CloudProvider for AzureCloudProvider {
     ) -> Result<Option<ModuleResp>, anyhow::Error> {
         _get_module_optional(self, crate::get_latest_stack_version_query(stack, track)).await
     }
+    async fn get_job_status(&self, job_id: &str) -> Result<Option<JobStatus>, anyhow::Error> {
+        match crate::run_function(
+            &self.function_endpoint,
+            &crate::get_job_status_query(job_id),
+            &self.project_id,
+            &self.region,
+        )
+        .await
+        {
+            Ok(response) => {
+                let job_status: JobStatus = serde_json::from_value(response.payload)?;
+                Ok(Some(job_status))
+            }
+            Err(e) => Err(e),
+        }
+    }
     async fn generate_presigned_url(
         &self,
         key: &str,
diff --git a/env_common/src/interface/cloud_handlers.rs b/env_common/src/interface/cloud_handlers.rs
@@ -6,8 +6,8 @@ use env_aws::AwsCloudProvider;
 use env_azure::AzureCloudProvider;
 use env_defs::{
     CloudProvider, CloudProviderCommon, Dependent, DeploymentResp, EventData,
-    GenericFunctionResponse, InfraChangeRecord, LogData, ModuleResp, NotificationData, PolicyResp,
-    ProjectData,
+    GenericFunctionResponse, InfraChangeRecord, JobStatus, LogData, ModuleResp, NotificationData,
+    PolicyResp, ProjectData,
 };
 use serde_json::Value;
 
@@ -313,6 +313,9 @@ impl CloudProvider for GenericCloudHandler {
             .get_deployment(deployment_id, environment, include_deleted)
             .await
     }
+    async fn get_job_status(&self, job_id: &str) -> Result<Option<JobStatus>, anyhow::Error> {
+        self.provider.get_job_status(job_id).await
+    }
     async fn get_deployments_using_module(
         &self,
         module: &str,
diff --git a/env_common/src/interface/no_cloud_provider.rs b/env_common/src/interface/no_cloud_provider.rs
@@ -1,7 +1,7 @@
 use env_defs::{
     CloudProvider, CloudProviderCommon, Dependent, DeploymentResp, EventData,
-    GenericFunctionResponse, InfraChangeRecord, LogData, ModuleResp, NotificationData, PolicyResp,
-    ProjectData,
+    GenericFunctionResponse, InfraChangeRecord, JobStatus, LogData, ModuleResp, NotificationData,
+    PolicyResp, ProjectData,
 };
 use serde_json::Value;
 use std::{future::Future, pin::Pin};
@@ -140,6 +140,10 @@ impl CloudProvider for NoCloudProvider {
         Ok(None)
     }
 
+    async fn get_job_status(&self, _job_id: &str) -> Result<Option<JobStatus>, anyhow::Error> {
+        Ok(None)
+    }
+
     async fn generate_presigned_url(
         &self,
         _key: &str,
diff --git a/env_common/src/logic/api_infra.rs b/env_common/src/logic/api_infra.rs
@@ -8,7 +8,7 @@ use env_utils::{
     get_version_track, verify_required_variables_are_set, verify_variable_claim_casing,
     verify_variable_existence_and_type,
 };
-use log::{debug, error, info};
+use log::{debug, error, info, warn};
 
 use crate::{interface::GenericCloudHandler, DeploymentStatusHandler};
 
@@ -572,7 +572,38 @@ pub async fn is_deployment_in_progress(
     };
 
     if busy_statuses.contains(&deployment.status.as_str()) {
-        info!("Deployment is currently in process: {}", deployment.status);
+        warn!(
+            "Deployment is currently in process according to deployment: {}",
+            deployment.status
+        );
+        warn!(
+            "Trying to verify that a VM is running for deployment job: {}",
+            deployment.job_id
+        );
+        match handler.get_job_status(&deployment.job_id).await {
+            Ok(Some(job_status)) => {
+                if job_status.is_running {
+                    warn!("Job {} is indeed running", deployment.job_id);
+                } else {
+                    warn!("Job {} is not running, proceeding. (This may have been caused by an error in the previous run)", deployment.job_id);
+                    return (
+                        false,
+                        "".to_string(),
+                        deployment.status.to_string(),
+                        Some(deployment.clone()),
+                    );
+                }
+            }
+            Ok(None) => {
+                error!(
+                    "No job status found for {}, please talk to your administrator.",
+                    deployment.job_id
+                );
+            }
+            Err(e) => {
+                error!("Failed to get job status for {}: {}", deployment.job_id, e);
+            }
+        };
         return (
             true,
             deployment.job_id.clone(),
diff --git a/integration-tests/azure-function-code/function_app.py b/integration-tests/azure-function-code/function_app.py
@@ -64,6 +64,8 @@ def handler(req: func.HttpRequest) -> func.HttpResponse:
             return read_logs(req)
         elif event == 'generate_presigned_url':
             return generate_presigned_url(req)
+        elif event == 'get_job_status':
+            return get_job_status(req)
         elif event == 'transact_write':
             return transact_write(req)
         elif event == 'publish_notification':
@@ -160,9 +162,19 @@ def generate_presigned_url(req: func.HttpRequest) -> func.HttpResponse:
         mimetype="application/json"
     )
 
+def get_job_status(req: func.HttpRequest) -> func.HttpResponse:
+    req_body = req.get_json()
+    payload = req_body.get('data')
+    job_id = payload.get('job_id')
+    # For testing purposes, we can simulate different scenarios:
+    # - If job_id starts with 'running-', return is_running=True
+    # - Otherwise, return is_running=False
+    is_running = job_id.startswith('running-') if job_id else False
+    return func.HttpResponse(json.dumps({'job_id': job_id, 'is_running': is_running}), status_code=200, mimetype="application/json")
+
 def start_runner(req: func.HttpRequest) -> func.HttpResponse:
     # TODO: Implement the ACI task start logic as another docker container
-    return func.HttpResponse(json.dumps({"result":"Would have been started", "job_id": "test-job-id"}), status_code=200)
+    return func.HttpResponse(json.dumps({"result":"Would have been started", "job_id": "running-test-job-id"}), status_code=200)
     
 def get_id(body):
     raw = f"{body['PK']}~{body['SK']}".lower()
diff --git a/integration-tests/lambda-code/test-api.py b/integration-tests/lambda-code/test-api.py
@@ -199,11 +199,20 @@ def generate_presigned_url(event):
     url = 'http://127.0.0.1:' + ':'.join(url.split(':')[2:]) # MinIO is running in a separate network and interfaces with lambda via 172.18.x.x and local rust app using 127.0.0.1
     return {'url': url}
 
+def get_job_status(event):
+    payload = event.get('data')
+    job_id = payload.get('job_id')
+    # For testing purposes, we can simulate different scenarios:
+    # - If job_id starts with 'running-', return is_running=True
+    # - Otherwise, return is_running=False
+    is_running = job_id.startswith('running-') if job_id else False
+    return {'job_id': job_id, 'is_running': is_running}
+
 def start_runner(event):
     # ecs = boto3.client('ecs')
     payload = event.get('data')
     
-    return {'job_id': 'test-job-id'}
+    return {'job_id': 'running-test-job-id'}
 
 processes = {
     'insert_db': insert_db,
@@ -214,6 +223,7 @@ def start_runner(event):
     'start_runner': start_runner,
     'read_logs': read_logs,
     'generate_presigned_url': generate_presigned_url,
+    'get_job_status': get_job_status,
     'publish_notification': lambda event: {'status': 'success', 'message': 'Notification published successfully'},
 }
 
diff --git a/integration-tests/tests/infra.rs b/integration-tests/tests/infra.rs
diff --git a/integration-tests/tests/runner.rs b/integration-tests/tests/runner.rs
diff --git a/integration-tests/tests/runner_oci.rs b/integration-tests/tests/runner_oci.rs
