1# Copyright The OpenTelemetry Authors 
    2# 
    3# Licensed under the Apache License, Version 2.0 (the "License"); 
    4# you may not use this file except in compliance with the License. 
    5# You may obtain a copy of the License at 
    6# 
    7#     http://www.apache.org/licenses/LICENSE-2.0 
    8# 
    9# Unless required by applicable law or agreed to in writing, software 
    10# distributed under the License is distributed on an "AS IS" BASIS, 
    11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
    12# See the License for the specific language governing permissions and 
    13# limitations under the License. 
    14from typing import Optional 
    15 
    16from opentelemetry.context import create_key, get_value, set_value 
    17from opentelemetry.context.context import Context 
    18from opentelemetry.trace.span import INVALID_SPAN, Span 
    19 
    20SPAN_KEY = "current-span" 
    21_SPAN_KEY = create_key("current-span") 
    22 
    23 
    24def set_span_in_context( 
    25    span: Span, context: Optional[Context] = None 
    26) -> Context: 
    27    """Set the span in the given context. 
    28 
    29    Args: 
    30        span: The Span to set. 
    31        context: a Context object. if one is not passed, the 
    32            default current context is used instead. 
    33    """ 
    34    ctx = set_value(_SPAN_KEY, span, context=context) 
    35    return ctx 
    36 
    37 
    38def get_current_span(context: Optional[Context] = None) -> Span: 
    39    """Retrieve the current span. 
    40 
    41    Args: 
    42        context: A Context object. If one is not passed, the 
    43            default current context is used instead. 
    44 
    45    Returns: 
    46        The Span set in the context if it exists. INVALID_SPAN otherwise. 
    47    """ 
    48    span = get_value(_SPAN_KEY, context=context) 
    49    if span is None or not isinstance(span, Span): 
    50        return INVALID_SPAN 
    51    return span