From 7d2fe8cb2acfd7806c550866688e382e6af3c60b Mon Sep 17 00:00:00 2001
From: vmarchaud <contact@vmarchaud.fr>
Date: Sat, 15 May 2021 16:30:15 +0200
Subject: [PATCH] feat(context): add utils method to remove keys from context
 #66

---
 src/api/propagation.ts     |  9 ++++++++-
 src/api/trace.ts           |  6 ++++++
 src/baggage/index.ts       | 13 +++++++++++++
 src/trace/context-utils.ts | 19 +++++++++++++++++++
 4 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/api/propagation.ts b/src/api/propagation.ts
index a554d2be..ede9b17a 100644
--- a/src/api/propagation.ts
+++ b/src/api/propagation.ts
@@ -28,7 +28,12 @@ import {
   registerGlobal,
   unregisterGlobal,
 } from '../internal/global-utils';
-import { getBaggage, createBaggage, setBaggage } from '../baggage/index';
+import {
+  getBaggage,
+  createBaggage,
+  setBaggage,
+  deleteBaggage,
+} from '../baggage/index';
 
 const API_NAME = 'propagation';
 
@@ -107,6 +112,8 @@ export class PropagationAPI {
 
   public setBaggage = setBaggage;
 
+  public deleteBaggage = deleteBaggage;
+
   private _getGlobalPropagator(): TextMapPropagator {
     return getGlobal(API_NAME) || NOOP_TEXT_MAP_PROPAGATOR;
   }
diff --git a/src/api/trace.ts b/src/api/trace.ts
index f44fe256..ec5b2ae1 100644
--- a/src/api/trace.ts
+++ b/src/api/trace.ts
@@ -27,6 +27,8 @@ import {
 import { Tracer } from '../trace/tracer';
 import { TracerProvider } from '../trace/tracer_provider';
 import {
+  deleteSpan,
+  deleteSpanContext,
   getSpan,
   getSpanContext,
   setSpan,
@@ -89,6 +91,10 @@ export class TraceAPI {
 
   public isSpanContextValid = isSpanContextValid;
 
+  public deleteSpan = deleteSpan;
+
+  public deleteSpanContext = deleteSpanContext;
+
   public getSpan = getSpan;
 
   public getSpanContext = getSpanContext;
diff --git a/src/baggage/index.ts b/src/baggage/index.ts
index 03eda420..ef52a471 100644
--- a/src/baggage/index.ts
+++ b/src/baggage/index.ts
@@ -41,6 +41,8 @@ export function createBaggage(
 }
 
 /**
+ * Get baggage from context if its present
+ *
  * @param {Context} Context that manage all context values
  * @returns {Baggage} Extracted baggage from the context
  */
@@ -49,6 +51,8 @@ export function getBaggage(context: Context): Baggage | undefined {
 }
 
 /**
+ * Set a baggage in a new context and returns it
+ *
  * @param {Context} Context that manage all context values
  * @param {Baggage} baggage that will be set in the actual context
  */
@@ -56,6 +60,15 @@ export function setBaggage(context: Context, baggage: Baggage): Context {
   return context.setValue(BAGGAGE_KEY, baggage);
 }
 
+/**
+ * Delete baggage from the given context
+ *
+ * @param {Context} Context context to delete baggage from
+ */
+export function deleteBaggage(context: Context): Context {
+  return context.deleteValue(BAGGAGE_KEY);
+}
+
 /**
  * Create a serializable BaggageEntryMetadata object from a string.
  *
diff --git a/src/trace/context-utils.ts b/src/trace/context-utils.ts
index 3714eb96..41517aa1 100644
--- a/src/trace/context-utils.ts
+++ b/src/trace/context-utils.ts
@@ -19,6 +19,7 @@ import { Context } from '../context/types';
 import { Span } from './span';
 import { SpanContext } from './span_context';
 import { NonRecordingSpan } from './NonRecordingSpan';
+import { INVALID_SPAN_CONTEXT } from './spancontext-utils';
 
 /**
  * span key
@@ -44,6 +45,15 @@ export function setSpan(context: Context, span: Span): Context {
   return context.setValue(SPAN_KEY, span);
 }
 
+/**
+ * Remove current span stored in the context
+ *
+ * @param context context to delete span from
+ */
+export function deleteSpan(context: Context): Context {
+  return context.deleteValue(SPAN_KEY);
+}
+
 /**
  * Wrap span context in a NoopSpan and set as span in a new
  * context
@@ -66,3 +76,12 @@ export function setSpanContext(
 export function getSpanContext(context: Context): SpanContext | undefined {
   return getSpan(context)?.spanContext();
 }
+
+/**
+ * Replace the span context from the context with an invalid span context
+ *
+ * @param context context to remove the span context from
+ */
+export function deleteSpanContext(context: Context): Context {
+  return setSpan(context, new NonRecordingSpan(INVALID_SPAN_CONTEXT));
+}