Skip to content

support setting token programatically #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 13, 2025
Merged

support setting token programatically #13

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3356603
support setting token programatically
davidhewitt Mar 13, 2025
25ec38f
coverage, test fixes
davidhewitt Mar 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 101 additions & 28 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,23 @@
//! - The [Logfire documentation](https://logfire.pydantic.dev/docs/) for more information about Logfire in general.
//! - The [Logfire GitHub repository](https://github.com/pydantic/logfire) for the source of the documentation, the Python SDK and an issue tracker for general questions about Logfire.
//!
//! > ***Initial release - feedback wanted!***
//! >
//! > This is an initial release of the Logfire Rust SDK. We've been using it internally to build
//! > Logfire for some time, and it is serving us well. As we're using it ourselves in production,
//! > we figured it's ready for everyone else also using Logfire.
//! >
//! > We are continually iterating to make this SDK better. We'd love your feedback on all aspects
//! > of the SDK and are keen to make the design as idiomatic and performant as possible. There are
//! > also many features currently supported by the Python SDK which are not yet supported by this
//! > SDK; please open issues to help us prioritize these to close this gap. For example, we have not
//! > yet implemented scrubbing in this Rust SDK, although we are aware it is important!
//! >
//! > In particular, the current coupling to `tracing` is an open design point. By building on top
//! > of tracing we get widest compatibility and a relatively simple SDK, however to make
//! > Logfire-specific adjustments we might prefer in future to move `tracing` to be an optional
//! > integration.
//!
//! ## Getting Started
//!
//! To use Logfire in your Rust project, add the following to your `Cargo.toml`:
Expand Down Expand Up @@ -45,6 +62,15 @@
//!
//! ## Configuration
//!
//! After adding basic setup as per above, the most two important environment variables are:
//! - `LOGFIRE_TOKEN` - (required) the token to send data to the Logfire platform
//! - `RUST_LOG` - (optional) the level of verbosity to send to the Logfire platform. By default
//! logs are captured at `TRACE` level so that all data is available for you to analyze in the
//! Logfire platform.
//!
//! All environment variables supported by the Rust Opentelemetry SDK are also supported by the
//! Logfire SDK.
//!
//! ## Integrations
//!
//! ### With `tracing`
Expand Down Expand Up @@ -99,6 +125,12 @@ mod internal;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ConfigureError {
/// No token was provided to send to logfire.
#[error(
"A logfire token is required from either `logfire::configure().with_token()` or the `LOGFIRE_TOKEN` environment variable"
)]
TokenRequired,

/// Logfire has already been configured
#[error("Logfire has already been configured")]
AlreadyConfigured,
Expand Down Expand Up @@ -195,6 +227,7 @@ impl FromStr for ConsoleMode {
pub fn configure() -> LogfireConfigBuilder {
LogfireConfigBuilder {
send_to_logfire: None,
token: None,
console_mode: ConsoleMode::Fallback,
additional_span_processors: Vec::new(),
advanced: None,
Expand All @@ -208,7 +241,7 @@ pub struct LogfireConfigBuilder {
// TODO: support all options supported by the Python SDK
// local: bool,
send_to_logfire: Option<SendToLogfire>,
// token: Option<String>,
token: Option<String>,
// service_name: Option<String>,
// service_version: Option<String>,
// environment: Option<String>,
Expand Down Expand Up @@ -257,6 +290,15 @@ impl LogfireConfigBuilder {
self
}

/// The token to use for the Logfire platform.
///
/// Defaults to the value of `LOGFIRE_TOKEN` if set.
#[must_use]
pub fn with_token<T: Into<String>>(mut self, token: T) -> Self {
self.token = Some(token.into());
self
}

/// Whether to log to the console.
#[must_use]
pub fn console_mode(mut self, console_mode: ConsoleMode) -> Self {
Expand Down Expand Up @@ -307,7 +349,7 @@ impl LogfireConfigBuilder {
tracer,
subscriber,
tracer_provider,
logfire_token,
token: logfire_token,
send_to_logfire,
base_url,
} = self.build_parts(None)?;
Expand Down Expand Up @@ -354,7 +396,10 @@ impl LogfireConfigBuilder {
self,
env: Option<&HashMap<String, String>>,
) -> Result<LogfireParts, ConfigureError> {
let logfire_token = get_optional_env("LOGFIRE_TOKEN", env)?;
let mut token = self.token;
if token.is_none() {
token = get_optional_env("LOGFIRE_TOKEN", env)?;
}

let send_to_logfire = match self.send_to_logfire {
Some(send_to_logfire) => send_to_logfire,
Expand All @@ -366,7 +411,7 @@ impl LogfireConfigBuilder {

let send_to_logfire = match send_to_logfire {
SendToLogfire::Yes => true,
SendToLogfire::IfTokenPresent => logfire_token.is_some(),
SendToLogfire::IfTokenPresent => token.is_some(),
SendToLogfire::No => false,
};

Expand All @@ -382,13 +427,13 @@ impl LogfireConfigBuilder {
};

if send_to_logfire {
if token.is_none() {
return Err(ConfigureError::TokenRequired);
}
tracer_provider_builder = tracer_provider_builder.with_span_processor(
BatchSpanProcessor::builder(RemovePendingSpansExporter::new(LogfireSpanExporter {
write_console: self.console_mode == ConsoleMode::Force,
inner: Some(span_exporter(
logfire_token.as_deref(),
&advanced_options.base_url,
)?),
inner: Some(span_exporter(token.as_deref(), &advanced_options.base_url)?),
}))
.with_batch_config(
BatchConfigBuilder::default()
Expand Down Expand Up @@ -445,7 +490,7 @@ impl LogfireConfigBuilder {
},
subscriber: Arc::new(subscriber),
tracer_provider,
logfire_token,
token,
send_to_logfire,
base_url: advanced_options.base_url,
})
Expand Down Expand Up @@ -489,7 +534,7 @@ struct LogfireParts {
tracer: LogfireTracer,
subscriber: Arc<dyn Subscriber + Send + Sync>,
tracer_provider: SdkTracerProvider,
logfire_token: Option<String>,
token: Option<String>,
send_to_logfire: bool,
base_url: String,
}
Expand Down Expand Up @@ -749,6 +794,7 @@ macro_rules! feature_required {
($feature_name:literal, $functionality:expr, $if_enabled:expr) => {{
#[cfg(feature = $feature_name)]
{
let _ = $functionality;
Ok($if_enabled)
}

Expand Down Expand Up @@ -1140,7 +1186,7 @@ mod tests {
"code.lineno",
),
value: I64(
1080,
1126,
),
},
KeyValue {
Expand Down Expand Up @@ -1266,7 +1312,7 @@ mod tests {
"code.lineno",
),
value: I64(
1081,
1127,
),
},
KeyValue {
Expand Down Expand Up @@ -1402,7 +1448,7 @@ mod tests {
"code.lineno",
),
value: I64(
1081,
1127,
),
},
KeyValue {
Expand Down Expand Up @@ -1544,7 +1590,7 @@ mod tests {
"code.lineno",
),
value: I64(
1082,
1128,
),
},
KeyValue {
Expand Down Expand Up @@ -1686,7 +1732,7 @@ mod tests {
"code.lineno",
),
value: I64(
1084,
1130,
),
},
KeyValue {
Expand Down Expand Up @@ -1856,7 +1902,7 @@ mod tests {
"code.lineno",
),
value: I64(
1085,
1131,
),
},
KeyValue {
Expand Down Expand Up @@ -1935,7 +1981,7 @@ mod tests {
),
value: String(
Owned(
"src/lib.rs:1086:17",
"src/lib.rs:1132:17",
),
),
},
Expand Down Expand Up @@ -2002,7 +2048,7 @@ mod tests {
"code.lineno",
),
value: I64(
513,
558,
),
},
KeyValue {
Expand Down Expand Up @@ -2100,7 +2146,7 @@ mod tests {
"code.lineno",
),
value: I64(
1080,
1126,
),
},
KeyValue {
Expand Down Expand Up @@ -2200,26 +2246,46 @@ mod tests {
#[test]
fn test_send_to_logfire() {
for (env, setting, expected) in [
(vec![], None, true),
(vec![("LOGFIRE_SEND_TO_LOGFIRE", "no")], None, false),
(vec![("LOGFIRE_SEND_TO_LOGFIRE", "yes")], None, true),
(vec![], None, Err(ConfigureError::TokenRequired)),
(vec![("LOGFIRE_TOKEN", "a")], None, Ok(true)),
(vec![("LOGFIRE_SEND_TO_LOGFIRE", "no")], None, Ok(false)),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "yes")],
None,
Err(ConfigureError::TokenRequired),
),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "yes"), ("LOGFIRE_TOKEN", "a")],
None,
Ok(true),
),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "if-token-present")],
None,
false,
Ok(false),
),
(
vec![
("LOGFIRE_SEND_TO_LOGFIRE", "if-token-present"),
("LOGFIRE_TOKEN", "a"),
],
None,
true,
Ok(true),
),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "no")],
vec![("LOGFIRE_SEND_TO_LOGFIRE", "no"), ("LOGFIRE_TOKEN", "a")],
Some(SendToLogfire::Yes),
true,
Ok(true),
),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "no"), ("LOGFIRE_TOKEN", "a")],
Some(SendToLogfire::IfTokenPresent),
Ok(true),
),
(
vec![("LOGFIRE_SEND_TO_LOGFIRE", "no")],
Some(SendToLogfire::IfTokenPresent),
Ok(false),
),
] {
let env = env.into_iter().map(|(k, v)| (k.into(), v.into())).collect();
Expand All @@ -2229,9 +2295,16 @@ mod tests {
config = config.send_to_logfire(value);
}

let parts = config.build_parts(Some(&env)).unwrap();
let result = config
.build_parts(Some(&env))
.map(|parts| parts.send_to_logfire);

assert_eq!(parts.send_to_logfire, expected);
match (expected, result) {
(Ok(exp), Ok(actual)) => assert_eq!(exp, actual),
// compare strings because ConfigureError doesn't implement PartialEq
(Err(exp), Err(actual)) => assert_eq!(exp.to_string(), actual.to_string()),
(expected, result) => panic!("expected {expected:?}, got {result:?}"),
}
}
}
}