datafusion_tracing/
options.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17//
18// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2025 Datadog, Inc.
19
20use crate::preview::PreviewFn;
21use std::collections::HashMap;
22use std::sync::Arc;
23
24/// Configuration options for instrumented execution plans.
25#[derive(Clone, Default)]
26pub struct InstrumentationOptions {
27    /// Whether to record metrics during execution.
28    pub record_metrics: bool,
29
30    /// Maximum number of rows to preview per span.
31    ///
32    /// If set to `0`, batch preview recording will be disabled.
33    pub preview_limit: usize,
34
35    /// Optional callback function for formatting previewed record batches.
36    ///
37    /// The provided function will be invoked for each previewed batch of at most `preview_limit` rows.
38    pub preview_fn: Option<Arc<PreviewFn>>,
39
40    /// User-defined custom fields for extensible configuration.
41    ///
42    /// This can be used to store arbitrary key-value pairs relevant to instrumentation or metadata.
43    pub custom_fields: HashMap<String, String>,
44}
45
46impl std::fmt::Debug for InstrumentationOptions {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("InstrumentationOptions")
49            .field("record_metrics", &self.record_metrics)
50            .field("preview_limit", &self.preview_limit)
51            .field(
52                "preview_fn",
53                &if self.preview_fn.is_some() {
54                    "Some(PreviewFn)"
55                } else {
56                    "None"
57                },
58            )
59            .field("custom_fields", &self.custom_fields)
60            .finish()
61    }
62}
63
64impl InstrumentationOptions {
65    /// Creates a new builder for `InstrumentationOptions`.
66    pub fn builder() -> InstrumentationOptionsBuilder {
67        InstrumentationOptionsBuilder::default()
68    }
69}
70
71/// The builder for `InstrumentationOptions`.
72#[derive(Default)]
73pub struct InstrumentationOptionsBuilder {
74    record_metrics: bool,
75    preview_limit: usize,
76    preview_fn: Option<Arc<PreviewFn>>,
77    custom_fields: HashMap<String, String>,
78}
79
80impl InstrumentationOptionsBuilder {
81    /// Sets whether to record metrics during execution.
82    pub fn record_metrics(mut self, record: bool) -> Self {
83        self.record_metrics = record;
84        self
85    }
86
87    /// Sets the maximum number of rows to preview per span.
88    /// Setting this to `0` disables the batch preview.
89    pub fn preview_limit(mut self, limit: usize) -> Self {
90        self.preview_limit = limit;
91        self
92    }
93
94    /// Sets the optional callback function for formatting previewed record batches.
95    pub fn preview_fn(mut self, func: Arc<PreviewFn>) -> Self {
96        self.preview_fn = Some(func);
97        self
98    }
99
100    /// Adds a single custom field.
101    pub fn add_custom_field<K: Into<String>, V: Into<String>>(
102        mut self,
103        key: K,
104        value: V,
105    ) -> Self {
106        self.custom_fields.insert(key.into(), value.into());
107        self
108    }
109
110    /// Replaces all custom fields with the provided `HashMap`.
111    pub fn custom_fields(mut self, fields: HashMap<String, String>) -> Self {
112        self.custom_fields = fields;
113        self
114    }
115
116    /// Consumes the builder and creates an `InstrumentationOptions` instance.
117    pub fn build(self) -> InstrumentationOptions {
118        InstrumentationOptions {
119            record_metrics: self.record_metrics,
120            preview_limit: self.preview_limit,
121            preview_fn: self.preview_fn,
122            custom_fields: self.custom_fields,
123        }
124    }
125}
datafusion_tracing/options.rs

datafusion_tracing/
options.rs
