Added in Telemetry

Next Build will have opt-in/opt-out capabilities
2025-11-14 13:41:04 -05:00
parent 0a0f39542d
commit bf5c7d156b
4 changed files with 1219 additions and 277 deletions
@@ -1,6 +1,12 @@
 use chrono::{Duration, Local, NaiveDate};
 use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
 use mongodb::bson::oid::ObjectId;
+use opentelemetry::global::shutdown_tracer_provider;
+use opentelemetry::sdk::Resource;
+use opentelemetry::trace::{Status, TraceContextExt, TraceError};
+use opentelemetry::{Context, KeyValue, sdk::trace as sdktrace, trace::Tracer};
+use opentelemetry::{Key, global};
+use opentelemetry_otlp::WithExportConfig;
 use pyo3::{prelude::*, types::PyString};
 use reqwest::{
    Client,
@@ -17,6 +23,7 @@ use std::{
    path::PathBuf,
    str::FromStr,
 };
+use tracing_subscriber::prelude::*;

 #[derive(Debug, Deserialize, Serialize)]
 struct ApiResponse {
@@ -61,6 +68,12 @@ pub fn pull_policy_exec_histories(
    exec_types: String,
    days: i64,
 ) -> Py<PyString> {
+    let rt = tokio::runtime::Runtime::new().unwrap();
+    rt.block_on(async {
+        let _ = init_tracer();
+    });
+    let tracer = global::tracer("global_tracer");
+    let _cx = Context::new();
    let file_path: PathBuf = format!(
        "{}\\cache\\chunkinator.json",
        get_base_directory().display()
@@ -93,102 +106,142 @@ 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 client = tracer.in_span("Building HTTP Client", |cx| {
+        let client_result = build_client(py, &py_self);
+        match client_result {
+            Ok(client_result) => {
+                cx.span().add_event(
+                    "info",
+                    vec![KeyValue::new(
+                        "Client Built Successfully",
+                        format!("{:?}", client_result),
+                    )],
+                );
+                client_result
+            }
+            Err(client_result) => {
+                cx.span().add_event(
+                    "warn",
+                    vec![KeyValue::new(
+                        "Client Failed to Build",
+                        format!("{:?}", &client_result),
+                    )],
+                );
+                cx.span()
+                    .set_status(Status::error("Client Failed to Build"));
+                panic!("Failed to Build Client: {:?}", client_result);
+            }
+        }
+    });
    let api: Py<PyAny> = py_self;
    let cutoff = Local::now().naive_local() - Duration::days(days);
    let mut f = File::open(&writeable_filepath).unwrap();
-    loop {
-        f.seek(SeekFrom::Start(0)).unwrap();
-        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;
-        }
-        let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists() {
-            let mut contents = String::new();
-            f.read_to_string(&mut contents).unwrap();
-            let existing_data: ApiResponse =
-                serde_json::from_str(&contents).unwrap_or(ApiResponse {
-                    error: "Success".to_string(),
-                    response: ExecHistories {
-                        exechistories: vec![],
-                    },
-                });
-            existing_data
-                .response
-                .exechistories
-                .into_iter()
-                .map(|entry| {
-                    (
-                        (
-                            entry.sha256.clone(),
-                            entry.filename.clone(),
-                            entry.hostname.clone(),
-                        ),
-                        entry,
-                    )
-                })
-                .collect()
-        } else {
-            HashMap::new()
-        };
-        for (index, executions) in parsed_responses.iter().enumerate() {
-            if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
-                continue;
-            }
-            if index == parsed_responses.len() - 1 {
-                checkpoint_number = executions.checkpoint.clone();
+    tracer.in_span("Airlock Data Retreival", |cx| {
+        let span = cx.span();
+        span.set_attribute(Key::new("Days").string(days.to_string().to_string()));
+        loop {
+            tracing::info!("Starting Airlock Data Retrieval");
+            f.seek(SeekFrom::Start(0)).unwrap();
+            let execution_histories = tracer.in_span(checkpoint_number.to_string(), |cx| {
+                let results: ApiResponse = history_logging(
+                    py,
+                    &api,
+                    &exec_types,
+                    &checkpoint_number,
+                    &policy_names,
+                    &client,
+                );
+                cx.span().set_attribute(KeyValue::new(
+                    "Items in Response",
+                    results.response.exechistories.len().to_string(),
+                ));
+                results
+            });
+            let parsed_responses = execution_histories.response.exechistories;
+            if parsed_responses.is_empty() {
                break;
            }
-            let history_date = match NaiveDate::parse_from_str(
-                &executions.datetime.replace(" +0000 UTC", ""),
-                "%Y-%m-%dT%H:%M:%SZ",
-            ) {
-                Ok(date) => date,
-                Err(_) => continue,
+            let mut seen: HashMap<(String, String, String), Group> = if writeable_filepath.exists()
+            {
+                let mut contents = String::new();
+                f.read_to_string(&mut contents).unwrap();
+                let existing_data: ApiResponse =
+                    serde_json::from_str(&contents).unwrap_or(ApiResponse {
+                        error: "Success".to_string(),
+                        response: ExecHistories {
+                            exechistories: vec![],
+                        },
+                    });
+                existing_data
+                    .response
+                    .exechistories
+                    .into_iter()
+                    .map(|entry| {
+                        (
+                            (
+                                entry.sha256.clone(),
+                                entry.filename.clone(),
+                                entry.hostname.clone(),
+                            ),
+                            entry,
+                        )
+                    })
+                    .collect()
+            } else {
+                HashMap::new()
            };
-            if history_date >= cutoff.into() {
-                let key = (
-                    executions.sha256.clone(),
-                    executions.filename.clone(),
-                    executions.hostname.clone(),
-                );
-                seen.entry(key).or_insert(executions.clone());
+            for (index, executions) in parsed_responses.iter().enumerate() {
+                if executions.checkpoint.is_empty() || executions.datetime.is_empty() {
+                    continue;
+                }
+                if index == parsed_responses.len() - 1 {
+                    checkpoint_number = executions.checkpoint.clone();
+                    break;
+                }
+                let history_date = match NaiveDate::parse_from_str(
+                    &executions.datetime.replace(" +0000 UTC", ""),
+                    "%Y-%m-%dT%H:%M:%SZ",
+                ) {
+                    Ok(date) => date,
+                    Err(_) => continue,
+                };
+                if history_date >= cutoff.into() {
+                    let key = (
+                        executions.sha256.clone(),
+                        executions.filename.clone(),
+                        executions.hostname.clone(),
+                    );
+                    seen.entry(key).or_insert(executions.clone());
+                }
+            }
+            let final_response = ApiResponse {
+                error: "Success".to_string(),
+                response: ExecHistories {
+                    exechistories: seen.values().cloned().collect(),
+                },
+            };
+            let data_write = serde_json::to_string_pretty(&final_response).unwrap();
+            fs::write(&writeable_filepath, data_write).unwrap();
+            if let Some(last_item) = &final_response.response.exechistories.last()
+                && let Ok(last_date) = NaiveDate::parse_from_str(
+                    &last_item.datetime.replace(" +0000 UTC", ""),
+                    "%Y-%m-%dT%H:%M:%SZ",
+                )
+            {
+                let date_diff = Local::now().naive_local().date() - last_date;
+                let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
+                progress_bar.set_position(percentage_diff.round() as u64);
+                progress_bar.set_message("Total Percent Complete");
            }
        }
-        let final_response = ApiResponse {
-            error: "Success".to_string(),
-            response: ExecHistories {
-                exechistories: seen.values().cloned().collect(),
-            },
-        };
-        let data_write = serde_json::to_string_pretty(&final_response).unwrap();
-        fs::write(&writeable_filepath, data_write).unwrap();
-        if let Some(last_item) = &final_response.response.exechistories.last()
-            && let Ok(last_date) = NaiveDate::parse_from_str(
-                &last_item.datetime.replace(" +0000 UTC", ""),
-                "%Y-%m-%dT%H:%M:%SZ",
-            )
-        {
-            let date_diff = Local::now().naive_local().date() - last_date;
-            let percentage_diff = (days - date_diff.num_days()) as f64 / days as f64 * 100.0;
-            progress_bar.set_position(percentage_diff.round() as u64);
-            progress_bar.set_message("Total Percent Complete");
-        }
-    }
+    });
    progress_bar.finish_with_message("All Checkpoints Complete");
    let return_data = fs::read_to_string(&writeable_filepath).unwrap();
+    shutdown_tracer_provider();
    PyString::new(py, &return_data).into()
 }

-fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
+fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Result<reqwest::Client, reqwest::Error> {
    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();
@@ -207,7 +260,6 @@ fn build_client(py: Python<'_>, py_self: &Py<PyAny>) -> Client {
        .default_headers(header_map)
        .timeout(std::time::Duration::from_secs(300))
        .build()
-        .unwrap()
 }

 #[tokio::main]
@@ -276,3 +328,20 @@ fn skipback(days: i64) -> ObjectId {
    let objectid_hex = format!("{}0000000000000000", hex_timestamp);
    ObjectId::parse_str(&objectid_hex).expect("Invalid ObjectId hex")
 }
+
+fn init_tracer() -> Result<sdktrace::Tracer, TraceError> {
+    opentelemetry_otlp::new_pipeline()
+        .tracing()
+        .with_exporter(
+            opentelemetry_otlp::new_exporter()
+                .tonic()
+                .with_endpoint("https://signoz.racooncity.org"),
+        )
+        .with_trace_config(
+            sdktrace::config().with_resource(Resource::new(vec![KeyValue::new(
+                "service.name",
+                "LoxideLibs",
+            )])),
+        )
+        .install_simple()
+}