Merge pull request #11 from axiomhq/arne/add-dataset-name-param

Add dataset name param

.github/workflows/ci.yaml

@@ -69,6 +69,7 @@ jobs:
         env:
           AXIOM_TOKEN: ${{ secrets.AXIOM_TOKEN }}
           AXIOM_URL: https://cloud.dev.axiomtestlabs.co
+          AXIOM_DATASET: _traces
         with:
           command: test
   publish_on_crates_io:
@@ -86,4 +87,4 @@ jobs:
         with:
           command: publish
         env:
-          CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
+          CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tracing-axiom"
-version = "0.3.0"
+version = "0.4.0"
 authors = ["Arne Bahlo <[email protected]>"]
 edition = "2021"
 rust-version = "1.60"

README.md

@@ -30,15 +30,15 @@ Add the following to your Cargo.toml:
 
 ```toml
 [dependencies]
-tracing-axiom = "0.3"
+tracing-axiom = "0.4"
 ```
 
-If you use the [Axiom CLI](https://github.com/axiomhq/cli), run `eval $(axiom config export -f)` to configure your environment variables.
+Create a dataset in Axiom and export the name as `AXIOM_DATASET`.
+Then create an API token with ingest permission into that dataset in
+[the Axiom settings](https://cloud.axiom.co/settings/profile) and export it as
+`AXIOM_TOKEN`.
 
-Otherwise create an API token in [the Axiom settings](https://cloud.axiom.co/settings/profile) and export it as `AXIOM_TOKEN`. 
-The token needs to have ingest permission to all datasets, using it will automatically create a `_traces` dataset for you.
-
-Set up tracing in one line like this:
+Now you can set up tracing in one line like this:
 
 ```rust
 #[tokio::main]

src/builder.rs

@@ -38,6 +38,7 @@ impl Drop for Guard {
 /// environment variables.
 #[derive(Debug, Default)]
 pub struct Builder {
+    dataset_name: Option<String>,
     token: Option<String>,
     url: Option<String>,
     trace_config: Option<TraceConfig>,
@@ -51,6 +52,11 @@ impl Builder {
         Self::default()
     }
 
+    pub fn with_dataset(mut self, dataset_name: impl Into<String>) -> Self {
+        self.dataset_name = Some(dataset_name.into());
+        self
+    }
+
     /// Set the Axiom API token to use.
     pub fn with_token(mut self, token: impl Into<String>) -> Self {
         self.token = Some(token.into());
@@ -121,6 +127,15 @@ impl Builder {
             return Err(Error::InvalidToken);
         }
 
+        let mut dataset_name = self.dataset_name;
+        if !self.no_env {
+            dataset_name = dataset_name.or_else(|| env::var("AXIOM_DATASET").ok());
+        }
+        let dataset_name = dataset_name.ok_or(Error::MissingDatasetName)?;
+        if dataset_name.is_empty() {
+            return Err(Error::EmptyDatasetName);
+        }
+
         let mut url = self.url;
         if !self.no_env {
             url = url.or_else(|| env::var("AXIOM_URL").ok());
@@ -133,6 +148,7 @@ impl Builder {
 
         let mut headers = HashMap::with_capacity(2);
         headers.insert("Authorization".to_string(), format!("Bearer {}", token));
+        headers.insert("X-Axiom-Dataset".to_string(), dataset_name);
         headers.insert(
             "User-Agent".to_string(),
             format!("tracing-axiom/{}", env!("CARGO_PKG_VERSION")),
@@ -200,6 +216,7 @@ mod tests {
         match Builder::new()
             .no_env()
             .with_token("xaat-123456789")
+            .with_dataset("test")
             .with_url("<invalid>")
             .try_init()
         {
@@ -213,8 +230,10 @@ mod tests {
         // Note that we can't test the init/try_init funcs here because OTEL
         // gets confused with the global subscriber.
 
-        let result: Result<(OpenTelemetryLayer<Registry, Tracer>, Guard), Error> =
-            Builder::new().with_token("xaat-123456789").layer();
+        let result: Result<(OpenTelemetryLayer<Registry, Tracer>, Guard), Error> = Builder::new()
+            .with_dataset("test")
+            .with_token("xaat-123456789")
+            .layer();
         assert!(result.is_ok(), "{:?}", result.err());
     }
 
@@ -227,7 +246,7 @@ mod tests {
         env::set_var("AXIOM_TOKEN", "xaat-1234567890");
 
         let result: Result<(OpenTelemetryLayer<Registry, Tracer>, Guard), Error> =
-            Builder::new().layer();
+            Builder::new().with_dataset("test").layer();
 
         if let Ok(token) = env_backup {
             env::set_var("AXIOM_TOKEN", token);

src/error.rs

@@ -14,6 +14,10 @@ pub enum Error {
     EmptyToken,
     #[error("Invalid token (please provide a personal token)")]
     InvalidToken,
+    #[error("Dataset name is missing")]
+    MissingDatasetName,
+    #[error("Dataset name is empty")]
+    EmptyDatasetName,
     #[error("Invalid URL: {0}")]
     InvalidUrl(#[from] url::ParseError),
 }

src/lib.rs

@@ -20,7 +20,8 @@
 //! ```
 //!
 //! The example above gets the Axiom API token from the `AXIOM_TOKEN` env and
-//! panics if setup fails. If you want to handle the error, use [`try_init`].
+//! the dataset name from `AXIOM_DATASET` and panics if setup fails.
+//! If you want to handle the error, use [`try_init`].
 //! For more advanced configuration, see [`builder()`].
 
 mod builder;
@@ -42,23 +43,24 @@ pub struct ReadmeDoctests;
 /// # Panics
 ///
 /// Panics if the initialization was unsuccessful, likely because a global
-/// subscriber was already installed or `AXIOM_TOKEN` is not set or invalid.
+/// subscriber was already installed or `AXIOM_TOKEN` and/or `AXIOM_DATASET`
+/// is not set or invalid.
 /// If you want to handle the error instead, use [`try_init`].
 pub fn init() -> Guard {
     Builder::new().init()
 }
 
 /// Initialize a global subscriber which sends traces to Axiom.
 ///
-/// It uses the environment variables `AXIOM_TOKEN` and optionally `AXIOM_URL`
-/// to configure the endpoint.
+/// It uses the environment variables `AXIOM_TOKEN`, `AXIOM_DATASET` and
+/// optionally `AXIOM_URL` to configure the endpoint.
 /// If you want to manually set these, see [`Builder`].
 ///
 /// # Errors
 ///
 /// Returns an error if the initialization was unsuccessful, likely because a
-/// global subscriber was already installed or `AXIOM_TOKEN` is not set or
-/// invalid.
+/// global subscriber was already installed or `AXIOM_TOKEN` or `AXIOM_DATASET`
+/// is not set or invalid.
 pub fn try_init() -> Result<Guard, Error> {
     Builder::new().try_init()
 }