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 InstrumentationOptions {
47    /// Creates a new builder for `InstrumentationOptions`.
48    pub fn builder() -> InstrumentationOptionsBuilder {
49        InstrumentationOptionsBuilder::default()
50    }
51}
52
53/// The builder for `InstrumentationOptions`.
54#[derive(Default)]
55pub struct InstrumentationOptionsBuilder {
56    record_metrics: bool,
57    preview_limit: usize,
58    preview_fn: Option<Arc<PreviewFn>>,
59    custom_fields: HashMap<String, String>,
60}
61
62impl InstrumentationOptionsBuilder {
63    /// Sets whether to record metrics during execution.
64    pub fn record_metrics(mut self, record: bool) -> Self {
65        self.record_metrics = record;
66        self
67    }
68
69    /// Sets the maximum number of rows to preview per span.
70    /// Setting this to `0` disables the batch preview.
71    pub fn preview_limit(mut self, limit: usize) -> Self {
72        self.preview_limit = limit;
73        self
74    }
75
76    /// Sets the optional callback function for formatting previewed record batches.
77    pub fn preview_fn(mut self, func: Arc<PreviewFn>) -> Self {
78        self.preview_fn = Some(func);
79        self
80    }
81
82    /// Adds a single custom field.
83    pub fn add_custom_field<K: Into<String>, V: Into<String>>(
84        mut self,
85        key: K,
86        value: V,
87    ) -> Self {
88        self.custom_fields.insert(key.into(), value.into());
89        self
90    }
91
92    /// Replaces all custom fields with the provided `HashMap`.
93    pub fn custom_fields(mut self, fields: HashMap<String, String>) -> Self {
94        self.custom_fields = fields;
95        self
96    }
97
98    /// Consumes the builder and creates an `InstrumentationOptions` instance.
99    pub fn build(self) -> InstrumentationOptions {
100        InstrumentationOptions {
101            record_metrics: self.record_metrics,
102            preview_limit: self.preview_limit,
103            preview_fn: self.preview_fn,
104            custom_fields: self.custom_fields,
105        }
106    }
107}
datafusion_tracing/options.rs

datafusion_tracing/
options.rs
