Allow to start_as_current_span with end_on_exit=False (open-telemetry#1519)

anton-ryzhov · web-flow · commit b32365bf71a8 · 2021-01-21T10:01:46.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v0.17b0...HEAD)
 
+### Added
+- Added `end_on_exit` argument to `start_as_current_span`
+  ([#1519](https://github.com/open-telemetry/opentelemetry-python/pull/1519)])
+
 ## [0.17b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v0.17b0) - 2021-01-20
 
 ### Added
diff --git a/opentelemetry-api/src/opentelemetry/trace/__init__.py b/opentelemetry-api/src/opentelemetry/trace/__init__.py
@@ -290,6 +290,7 @@ def start_as_current_span(
         start_time: typing.Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
+        end_on_exit: bool = True,
     ) -> typing.Iterator["Span"]:
         """Context manager for creating a new span and set it
         as the current span in this tracer's context.
@@ -396,6 +397,7 @@ def start_as_current_span(
         start_time: typing.Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
+        end_on_exit: bool = True,
     ) -> typing.Iterator["Span"]:
         # pylint: disable=unused-argument,no-self-use
         yield INVALID_SPAN
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
@@ -759,6 +759,7 @@ def start_as_current_span(
         start_time: Optional[int] = None,
         record_exception: bool = True,
         set_status_on_exception: bool = True,
+        end_on_exit: bool = True,
     ) -> Iterator[trace_api.Span]:
         span = self.start_span(
             name=name,
@@ -770,7 +771,7 @@ def start_as_current_span(
             record_exception=record_exception,
             set_status_on_exception=set_status_on_exception,
         )
-        with self.use_span(span, end_on_exit=True) as span_context:
+        with self.use_span(span, end_on_exit=end_on_exit) as span_context:
             yield span_context
 
     def start_span(  # pylint: disable=too-many-locals
diff --git a/opentelemetry-sdk/tests/trace/test_trace.py b/opentelemetry-sdk/tests/trace/test_trace.py
@@ -485,6 +485,14 @@ def func():
         self.assertIsNotNone(root2.end_time)
         self.assertIsNot(root1, root2)
 
+    def test_start_as_current_span_no_end_on_exit(self):
+        tracer = new_tracer()
+
+        with tracer.start_as_current_span("root", end_on_exit=False) as root:
+            self.assertIsNone(root.end_time)
+
+        self.assertIsNone(root.end_time)
+
     def test_explicit_span_resource(self):
         resource = resources.Resource.create({})
         tracer_provider = trace.TracerProvider(resource=resource)
