Add license key to trace resource (#6455)

GitOrigin-RevId: 039747618a0689fa9818cb8756fc9934684eeb79

Author: voodoo11

Cargo.lock

@@ -58,6 +58,7 @@ pyo3-asyncio = "0.20.0"
 pyo3-log = "0.9.0"
 rand = "0.8.5"
 rdkafka = { version = "0.36.2", features = ["ssl-vendored", "cmake-build", "zstd"] }
+regex = "1.10.4"
 reqwest = { version = "0.12.4", features = ["blocking", "json"] }
 rusqlite = { version = "0.31.0", features = ["bundled"] }
 rust-s3 = { version = "0.33.0", features = ["sync-native-tls-vendored", "sync-native-tls", "fail-on-err"], default-features = false }

Cargo.toml

@@ -790,6 +790,7 @@ class TelemetryConfig:
     service_namespace: str | None
     service_instance_id: str | None
     run_id: str
+    license_key: str | None
     @staticmethod
     def create(
         *,

python/pathway/engine.pyi

@@ -67,6 +67,7 @@ def _resource(self) -> Resource:
                 SERVICE_INSTANCE_ID: self.config.service_instance_id or "",
                 "run.id": self.config.run_id,
                 "python.version": sys.version,
+                "license.key": self.config.license_key or "",
             }
         )
 

python/pathway/internals/graph_runner/telemetry.py

@@ -4,6 +4,7 @@ use cached::proc_macro::cached;
 use log::info;
 use nix::sys::resource::Resource;
 use once_cell::sync::Lazy;
+use regex::Regex;
 use serde::Deserialize;
 use serde_json::{json, Value};
 
@@ -49,6 +50,23 @@ impl License {
             }
         }
     }
+
+    pub fn shortcut(&self) -> String {
+        match self {
+            License::NoLicenseKey => String::new(),
+            License::LicenseKey(key) => {
+                let pattern = r"^([^-\s]+-){4}[^-\s]+(-[^-\s]+)*$";
+                let re = Regex::new(pattern).unwrap();
+
+                if re.is_match(key) {
+                    let parts: Vec<&str> = key.split('-').collect();
+                    format!("{}-{}", parts[0], parts[1])
+                } else {
+                    String::new()
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, thiserror::Error)]

src/engine/license.rs

@@ -46,6 +46,7 @@ const OUTPUT_LATENCY: &str = "latency.output";
 
 const ROOT_TRACE_ID: Key = Key::from_static_str("root.trace.id");
 const RUN_ID: Key = Key::from_static_str("run.id");
+const LICENSE_KEY: Key = Key::from_static_str("license.key");
 
 const LOCAL_DEV_NAMESPACE: &str = "local-dev";
 
@@ -68,6 +69,7 @@ impl Telemetry {
             SERVICE_NAMESPACE.string(self.config.service_namespace.clone()),
             ROOT_TRACE_ID.string(root_trace_id.to_string()),
             RUN_ID.string(self.config.run_id.clone()),
+            LICENSE_KEY.string(self.config.license_key.clone()),
         ])
     }
 
@@ -164,6 +166,7 @@ pub struct TelemetryEnabled {
     pub service_instance_id: String,
     pub run_id: String,
     pub trace_parent: Option<String>,
+    pub license_key: String,
 }
 
 #[derive(Clone, Debug)]
@@ -198,9 +201,13 @@ impl Config {
         }
 
         match license {
-            License::LicenseKey(_) => {
-                Config::create_enabled(run_id, telemetry_server, monitoring_server, trace_parent)
-            }
+            License::LicenseKey(_) => Config::create_enabled(
+                run_id,
+                telemetry_server,
+                monitoring_server,
+                trace_parent,
+                license,
+            ),
             License::NoLicenseKey => Ok(Config::Disabled),
         }
     }
@@ -210,6 +217,7 @@ impl Config {
         telemetry_server: Option<String>,
         monitoring_server: Option<String>,
         trace_parent: Option<String>,
+        license: &License,
     ) -> Result<Self> {
         let service_instance_id: String = parse_env_var("PATHWAY_SERVICE_INSTANCE_ID")
             .map_err(DynError::from)?
@@ -233,6 +241,7 @@ impl Config {
             service_instance_id,
             run_id,
             trace_parent,
+            license_key: license.shortcut(),
         })))
     }
 }

src/engine/telemetry.rs

@@ -3495,6 +3495,7 @@ pub struct TelemetryConfig {
     service_namespace: Option<String>,
     service_instance_id: Option<String>,
     run_id: String,
+    license_key: Option<String>,
 }
 
 #[pymethods]
@@ -3529,6 +3530,7 @@ impl From<EngineTelemetryConfig> for TelemetryConfig {
                 service_namespace: Some(config.service_namespace),
                 service_instance_id: Some(config.service_instance_id),
                 run_id: config.run_id,
+                license_key: Some(config.license_key),
             },
             EngineTelemetryConfig::Disabled => Self::default(),
         }