scouter_tracing/exporter/
grpc.rs

1use crate::error::TraceError;
2use crate::exporter::traits::SpanExporterBuilder;
3use crate::exporter::ExporterType;
4use crate::utils::{ExportConfig, GrpcConfig, OtelProtocol};
5use opentelemetry_otlp::ExportConfig as OtlpExportConfig;
6use opentelemetry_otlp::SpanExporter as OtlpSpanExporter;
7use opentelemetry_otlp::WithExportConfig;
8use opentelemetry_otlp::WithTonicConfig;
9use pyo3::prelude::*;
10use scouter_types::{CompressionType, PyHelperFuncs};
11use serde::Serialize;
12use std::time::Duration;
13
14#[derive(Debug, Clone, Serialize)]
15#[pyclass]
16pub struct GrpcSpanExporter {
17    #[pyo3(get)]
18    pub sample_ratio: Option<f64>,
19
20    #[pyo3(get)]
21    pub batch_export: bool,
22
23    #[pyo3(get)]
24    endpoint: Option<String>,
25
26    #[pyo3(get)]
27    protocol: OtelProtocol,
28
29    #[pyo3(get)]
30    timeout: Option<u64>,
31
32    #[pyo3(get)]
33    compression: Option<CompressionType>,
34}
35
36#[pymethods]
37impl GrpcSpanExporter {
38    #[new]
39    #[pyo3(signature = (batch_export=true, export_config=None, grpc_config=None, sample_ratio=None))]
40    pub fn new(
41        batch_export: bool,
42        export_config: Option<&ExportConfig>,
43        grpc_config: Option<&GrpcConfig>,
44        sample_ratio: Option<f64>,
45    ) -> Result<Self, TraceError> {
46        let (endpoint, protocol, timeout) = if let Some(config) = export_config {
47            (
48                config.endpoint.clone(),
49                config.protocol.clone(),
50                config.timeout,
51            )
52        } else {
53            (None, OtelProtocol::default(), None)
54        };
55
56        let compression = if let Some(grpc_config) = grpc_config {
57            grpc_config.compression.clone()
58        } else {
59            None
60        };
61
62        Ok(Self {
63            batch_export,
64            sample_ratio,
65            endpoint,
66            protocol,
67            timeout,
68            compression,
69        })
70    }
71
72    pub fn __str__(&self) -> String {
73        PyHelperFuncs::__str__(self)
74    }
75}
76
77impl SpanExporterBuilder for GrpcSpanExporter {
78    type Exporter = OtlpSpanExporter;
79
80    fn export_type(&self) -> ExporterType {
81        ExporterType::Grpc
82    }
83
84    fn sample_ratio(&self) -> Option<f64> {
85        self.sample_ratio
86    }
87
88    fn batch_export(&self) -> bool {
89        self.batch_export
90    }
91
92    fn build_exporter(&self) -> Result<Self::Exporter, TraceError> {
93        // Reconstruct the OtlpExportConfig each time
94        let timeout = self.timeout.map(Duration::from_secs);
95        let export_config = OtlpExportConfig {
96            endpoint: self.endpoint.clone(),
97            protocol: self.protocol.to_otel_protocol(),
98            timeout,
99        };
100
101        let mut exporter = opentelemetry_otlp::SpanExporter::builder().with_tonic();
102
103        if let Some(compression) = &self.compression {
104            exporter = exporter.with_compression(compression.to_otel_compression()?);
105        }
106
107        exporter = exporter.with_export_config(export_config);
108
109        Ok(exporter.build()?)
110    }
111}