airlock_libs/src/services.rs

@@ -1,7 +1,7 @@
 use chrono::{Duration, Local, NaiveDate};
 use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
 use mongodb::bson::oid::ObjectId;
-use pyo3::prelude::*;
+use pyo3::{prelude::*, types::PyString};
 use reqwest::{
     Client,
     header::{HeaderMap, HeaderName, HeaderValue},
@@ -61,7 +61,7 @@ pub fn pull_policy_exec_histories(
     exec_types: String,
     days: i64,
     //output_json: bool,
-) {
+) -> Py<PyString> {
     let file_path: PathBuf = format!(
         "{}\\cache\\chunkinator.json",
         get_base_directory().display()
@@ -87,13 +87,6 @@ pub fn pull_policy_exec_histories(
     let mut checkpoint_number: String = skipback(days).to_string();
     let multi_progress = MultiProgress::new();
     multi_progress.set_draw_target(ProgressDrawTarget::stdout());
-    let data_bar = multi_progress.add(ProgressBar::new(10_000));
-    data_bar.set_style(
-        ProgressStyle::default_bar()
-            .template("Checkpoint Progress: [{bar:40.cyan/blue}] {pos}/{len} {msg}")
-            .unwrap(),
-    );
-    data_bar.set_message("Starting");
     let progress_bar = multi_progress.add(ProgressBar::new(100));
     progress_bar.set_style(
         ProgressStyle::default_bar()
@@ -101,15 +94,21 @@ pub fn pull_policy_exec_histories(
             .unwrap(),
     );
     progress_bar.enable_steady_tick(std::time::Duration::from_millis(100));
+    let client = build_client(py, &py_self);
     let api: Py<PyAny> = py_self;
     loop {
-        let execution_histories =
-            history_logging(py, &api, &exec_types, &checkpoint_number, &policy_names);
+        let execution_histories = history_logging(
+            py,
+            &api,
+            &exec_types,
+            &checkpoint_number,
+            &policy_names,
+            &client,
+        );
         let parsed_responses = execution_histories.response.exechistories;
         if parsed_responses.is_empty() {
             break;
         }
-        data_bar.set_length(parsed_responses.len() as u64);
         let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
             let mut f = File::open(&writeable_filepath).unwrap();
             let mut contents = String::new();
@@ -145,7 +144,6 @@ pub fn pull_policy_exec_histories(
             }
             if index == parsed_responses.len() - 1 {
                 checkpoint_number = executions.checkpoint.clone();
-                data_bar.set_message(checkpoint_number.clone());
                 break;
             }
             let history_date = match NaiveDate::parse_from_str(
@@ -164,7 +162,6 @@ pub fn pull_policy_exec_histories(
                 );
                 seen.entry(key).or_insert(executions.clone());
             }
-            data_bar.inc(1);
         }
         let final_response = ApiResponse {
             error: "Success".to_string(),
@@ -186,21 +183,16 @@ pub fn pull_policy_exec_histories(
             progress_bar.set_position(percentage_diff.round() as u64);
             progress_bar.set_message("Total Percent Complete");
         }
-        data_bar.set_position(0);
     }
-    data_bar.finish_with_message("Finished Checkpoints");
     progress_bar.finish_with_message("All Checkpoints Complete");
+    let return_data = fs::read_to_string(&writeable_filepath).unwrap();
+    //let json_data = serde_json::from_str(&return_data).unwrap();
+    PyString::new(py, &return_data).into()
+    //let py_any: Py<String> = serde_pyobject::to_pyobject(py, &json_data).unwrap().into();
+    //py_any
 }
 
-#[tokio::main]
-async fn history_logging(
-    py: Python<'_>,
-    py_self: &Py<PyAny>,
-    exec_types: &String,
-    checkpoint_number: &String,
-    policy_names: &String,
-) -> ApiResponse {
-    let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
+fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
     let headers = py_self.getattr(py, "headers").unwrap().to_string();
     let headers_replace = headers.replace('\'', "\"");
     let parsed: Value = serde_json::from_str(headers_replace.as_str()).unwrap();
@@ -213,6 +205,25 @@ async fn history_logging(
             }
         }
     }
+
+    Client::builder()
+        .danger_accept_invalid_certs(true)
+        .default_headers(header_map)
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+        .unwrap()
+}
+
+#[tokio::main]
+async fn history_logging(
+    py: Python<'_>,
+    py_self: &Py<PyAny>,
+    exec_types: &String,
+    checkpoint_number: &String,
+    policy_names: &String,
+    client: &Client,
+) -> ApiResponse {
+    let base_url = py_self.getattr(py, "base_url").unwrap().to_string();
     let payload = format!(
         r#"{{
         "type": {},
@@ -221,12 +232,6 @@ async fn history_logging(
     }}"#,
         exec_types, checkpoint_number, policy_names
     );
-    let client = Client::builder()
-        .danger_accept_invalid_certs(true)
-        .default_headers(header_map)
-        .timeout(std::time::Duration::from_secs(30))
-        .build()
-        .unwrap();
     let res = client
         .post(format!("{}/v1/logging/exechistories", base_url))
         .body(payload)

models/execution.py

@@ -26,6 +26,7 @@ from typing import List, Optional, Tuple
 import dotenv
 import pandas as pd
 
+import airlock_libs
 from services.API import AirlockAPIWrapper
 from services.policyhandler import pullPolicyExechistories
 from utils.configmanager import get_protected_value, load_env_json
@@ -263,9 +264,7 @@ class ExecutionHistoryRecord:
     ) -> List["ExecutionHistoryRecord"]:
         executions = []
         for policy in selected_policies:
-            execs = pullPolicyExechistories(
-                api, policy, type_, history_days, True
-            )
+            execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days)
             if execs:
                 data = json.loads(execs)
                 exechistories = data.get("response", {}).get("exechistories", [])