diff --git a/airlock_libs/src/services.rs b/airlock_libs/src/services.rs index cdca8bc..a61bf8f 100644 --- a/airlock_libs/src/services.rs +++ b/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 { 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 = 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 = serde_pyobject::to_pyobject(py, &json_data).unwrap().into(); + //py_any } -#[tokio::main] -async fn history_logging( - py: Python<'_>, - py_self: &Py, - 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) -> 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, + 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) diff --git a/models/execution.py b/models/execution.py index c5375df..13281d1 100644 --- a/models/execution.py +++ b/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", [])