Merge branch 'master' into gh-actions

Author: willarmiros

.github/workflows/UnitTesting.yaml

@@ -36,7 +36,7 @@ jobs:
           python-version: ${{ env[matrix.python-version] }}
 
       - name: Install tox
-        run: pip install -U tox-factor
+        run: pip install "tox<=3.27.1" -U tox-factor
 
       - name: Cache tox environment
         # Preserves .tox directory between runs for faster installs

CHANGELOG.rst

@@ -5,6 +5,20 @@ CHANGELOG
 Unreleased
 ==========
 
+2.11.0
+==========
+* bugfix: Fix TypeError by patching register_default_jsonb from psycopg2 `PR350 https://github.com/aws/aws-xray-sdk-python/pull/350`
+* improvement: Add annotations `PR348 https://github.com/aws/aws-xray-sdk-python/pull/348`
+* bugfix: Use service parameter to match centralized sampling rules `PR 353 https://github.com/aws/aws-xray-sdk-python/pull/353`
+* bugfix: Implement PEP3134 to discover underlying problems with python3 `PR355 https://github.com/aws/aws-xray-sdk-python/pull/355`
+* improvement: Allow list TopicArn for SNS PublishBatch request `PR358 https://github.com/aws/aws-xray-sdk-python/pull/358`
+* bugfix: Version pinning flask-sqlalchemy version to 2.5.1 or less `PR360 https://github.com/aws/aws-xray-sdk-python/pull/360`
+* bugfix: Fix UnboundLocalError when aiohttp server raises a CancelledError `PR356 https://github.com/aws/aws-xray-sdk-python/pull/356`
+* improvement: Instrument httpx >= 0.20 `PR357 https://github.com/aws/aws-xray-sdk-python/pull/357`
+* improvement: [LambdaContext] persist original trace header `PR362 https://github.com/aws/aws-xray-sdk-python/pull/362`
+* bugfix: Run tests against Django 4.x `PR361 https://github.com/aws/aws-xray-sdk-python/pull/361`
+* improvement: Oversampling Mitigation `PR366 https://github.com/aws/aws-xray-sdk-python/pull/366`
+
 2.10.0
 ==========
 * bugfix: Only import future for py2. `PR343 <https://github.com/aws/aws-xray-sdk-python/pull/343>`_.

README.md

@@ -124,6 +124,46 @@ xray_recorder.end_subsegment()
 xray_recorder.end_segment()
 ```
 
+### Oversampling Mitigation
+To modify the sampling decision at the subsegment level, subsegments that inherit the decision of their direct parent (segment or subsegment) can be created using `xray_recorder.begin_subsegment()` and unsampled subsegments can be created using
+`xray_recorder.begin_subsegment_without_sampling()`.
+
+The code snippet below demonstrates creating a sampled or unsampled subsegment based on the sampling decision of each SQS message processed by Lambda.
+
+```python
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.models.subsegment import Subsegment
+from aws_xray_sdk.core.utils.sqs_message_helper import SqsMessageHelper
+
+def lambda_handler(event, context):
+
+    for message in event['Records']:
+        if SqsMessageHelper.isSampled(message):
+            subsegment = xray_recorder.begin_subsegment('sampled_subsegment')
+            print('sampled - processing SQS message')
+
+        else:
+            subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled_subsegment')
+            print('unsampled - processing SQS message')
+    
+    xray_recorder.end_subsegment()   
+```
+
+The code snippet below demonstrates wrapping a downstream AWS SDK request with an unsampled subsegment.
+```python
+from aws_xray_sdk.core import xray_recorder, patch_all
+import boto3
+
+patch_all()
+
+def lambda_handler(event, context):
+    subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled_subsegment')
+    client = boto3.client('sqs')
+    print(client.list_queues())
+    
+    xray_recorder.end_subsegment()
+```
+
 ### Capture
 
 As a decorator:

aws_xray_sdk/core/async_recorder.py

@@ -1,5 +1,4 @@
 import time
-import six
 
 from aws_xray_sdk.core.recorder import AWSXRayRecorder
 from aws_xray_sdk.core.utils import stacktrace
@@ -82,10 +81,10 @@ async def record_subsegment_async(self, wrapped, instance, args, kwargs, name,
         try:
             return_value = await wrapped(*args, **kwargs)
             return return_value
-        except Exception as exc:
-            exception = exc
+        except Exception as e:
+            exception = e
             stack = stacktrace.get_stacktrace(limit=self._max_trace_back)
-            six.raise_from(exc, exc)
+            raise
         finally:
             # No-op if subsegment is `None` due to `LOG_ERROR`.
             if subsegment is not None:

aws_xray_sdk/core/context.py

@@ -27,7 +27,7 @@ class Context(object):
 
     This data structure is thread-safe.
     """
-    def __init__(self, context_missing='RUNTIME_ERROR'):
+    def __init__(self, context_missing='LOG_ERROR'):
 
         self._local = threading.local()
         strategy = os.getenv(CXT_MISSING_STRATEGY_KEY, context_missing)

aws_xray_sdk/core/models/dummy_entities.py

@@ -11,7 +11,7 @@ class DummySegment(Segment):
     the segment based on sampling rules.
     Adding data to a dummy segment becomes a no-op except for
     subsegments. This is to reduce the memory footprint of the SDK.
-    A dummy segment will not be sent to the X-Ray daemon. Manually create
+    A dummy segment will not be sent to the X-Ray daemon. Manually creating
     dummy segments is not recommended.
     """
 

aws_xray_sdk/core/models/entity.py

@@ -81,6 +81,10 @@ def add_subsegment(self, subsegment):
         """
         self._check_ended()
         subsegment.parent_id = self.id
+
+        if not self.sampled and subsegment.sampled:
+            log.warning("This sampled subsegment is being added to an unsampled parent segment/subsegment and will be orphaned.")
+
         self.subsegments.append(subsegment)
 
     def remove_subsegment(self, subsegment):

aws_xray_sdk/core/patcher.py

@@ -6,7 +6,6 @@
 import re
 import sys
 import wrapt
-import six
 
 from aws_xray_sdk import global_sdk_config
 from .utils.compat import PY2, is_classmethod, is_instance_method
@@ -110,9 +109,9 @@ def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
 def _patch_module(module_to_patch, raise_errors=True):
     try:
         _patch(module_to_patch)
-    except Exception as exc:
+    except Exception:
         if raise_errors:
-            six.raise_from(exc, exc)
+            raise
         log.debug('failed to patch module %s', module_to_patch)
 
 

aws_xray_sdk/core/recorder.py

@@ -4,7 +4,6 @@
 import os
 import platform
 import time
-import six
 
 from aws_xray_sdk import global_sdk_config
 from aws_xray_sdk.version import VERSION
@@ -275,16 +274,10 @@ def current_segment(self):
         else:
             return entity
 
-    def begin_subsegment(self, name, namespace='local'):
-        """
-        Begin a new subsegment.
-        If there is open subsegment, the newly created subsegment will be the
-        child of latest opened subsegment.
-        If not, it will be the child of the current open segment.
-
-        :param str name: the name of the subsegment.
-        :param str namespace: currently can only be 'local', 'remote', 'aws'.
-        """
+    def _begin_subsegment_helper(self, name, namespace='local', beginWithoutSampling=False):
+        '''
+        Helper method to begin_subsegment and begin_subsegment_without_sampling
+        '''
         # Generating the parent dummy segment is necessary.
         # We don't need to store anything in context. Assumption here
         # is that we only work with recorder-level APIs.
@@ -295,16 +288,42 @@ def begin_subsegment(self, name, namespace='local'):
         if not segment:
             log.warning("No segment found, cannot begin subsegment %s." % name)
             return None
-
-        if not segment.sampled:
+        
+        current_entity = self.get_trace_entity()
+        if not current_entity.sampled or beginWithoutSampling:
             subsegment = DummySubsegment(segment, name)
         else:
             subsegment = Subsegment(name, namespace, segment)
 
         self.context.put_subsegment(subsegment)
-
         return subsegment
 
+
+
+    def begin_subsegment(self, name, namespace='local'):
+        """
+        Begin a new subsegment.
+        If there is open subsegment, the newly created subsegment will be the
+        child of latest opened subsegment.
+        If not, it will be the child of the current open segment.
+
+        :param str name: the name of the subsegment.
+        :param str namespace: currently can only be 'local', 'remote', 'aws'.
+        """
+        return self._begin_subsegment_helper(name, namespace)
+
+
+    def begin_subsegment_without_sampling(self, name):
+        """
+        Begin a new unsampled subsegment.
+        If there is open subsegment, the newly created subsegment will be the
+        child of latest opened subsegment.
+        If not, it will be the child of the current open segment.
+
+        :param str name: the name of the subsegment.
+        """
+        return self._begin_subsegment_helper(name, beginWithoutSampling=True)
+
     def current_subsegment(self):
         """
         Return the latest opened subsegment. In a multithreading environment,
@@ -436,10 +455,10 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
         try:
             return_value = wrapped(*args, **kwargs)
             return return_value
-        except Exception as exc:
-            exception = exc
+        except Exception as e:
+            exception = e
             stack = stacktrace.get_stacktrace(limit=self.max_trace_back)
-            six.raise_from(exc, exc)
+            raise
         finally:
             # No-op if subsegment is `None` due to `LOG_ERROR`.
             if subsegment is not None:
@@ -487,7 +506,8 @@ def _send_segment(self):
 
     def _stream_subsegment_out(self, subsegment):
         log.debug("streaming subsegments...")
-        self.emitter.send_entity(subsegment)
+        if subsegment.sampled:
+            self.emitter.send_entity(subsegment)
 
     def _load_sampling_rules(self, sampling_rules):
 

aws_xray_sdk/core/utils/sqs_message_helper.py

@@ -0,0 +1,11 @@
+SQS_XRAY_HEADER = "AWSTraceHeader"
+class SqsMessageHelper:
+    
+    @staticmethod 
+    def isSampled(sqs_message):
+        attributes = sqs_message['attributes']
+
+        if SQS_XRAY_HEADER not in attributes:
+            return False
+
+        return 'Sampled=1' in attributes[SQS_XRAY_HEADER]

aws_xray_sdk/ext/aiohttp/middleware.py

@@ -1,5 +1,3 @@
-import six
-
 """
 AioHttp Middleware
 """
@@ -66,14 +64,14 @@ async def middleware(request, handler):
     except HTTPException as exc:
         # Non 2XX responses are raised as HTTPExceptions
         response = exc
-        six.raise_from(exc, exc)
-    except BaseException as exc:
+        raise
+    except BaseException as err:
         # Store exception information including the stacktrace to the segment
         response = None
         segment.put_http_meta(http.STATUS, 500)
         stack = stacktrace.get_stacktrace(limit=xray_recorder.max_trace_back)
-        segment.add_exception(exc, stack)
-        six.raise_from(exc, exc)
+        segment.add_exception(err, stack)
+        raise
     finally:
         if response is not None:
             segment.put_http_meta(http.STATUS, response.status)

aws_xray_sdk/ext/django/conf.py

@@ -6,7 +6,7 @@
 DEFAULTS = {
     'AWS_XRAY_DAEMON_ADDRESS': '127.0.0.1:2000',
     'AUTO_INSTRUMENT': True,
-    'AWS_XRAY_CONTEXT_MISSING': 'RUNTIME_ERROR',
+    'AWS_XRAY_CONTEXT_MISSING': 'LOG_ERROR',
     'PLUGINS': (),
     'SAMPLING': True,
     'SAMPLING_RULES': None,

aws_xray_sdk/ext/sqlalchemy_core/patch.py

@@ -1,13 +1,12 @@
 import logging
 import sys
-import wrapt
-import six
 
 if sys.version_info >= (3, 0, 0):
     from urllib.parse import urlparse, uses_netloc
 else:
     from urlparse import urlparse, uses_netloc
 
+import wrapt
 
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.patcher import _PATCHED_MODULES
@@ -73,12 +72,12 @@ def _process_request(wrapped, engine_instance, args, kwargs):
         subsegment = None
     try:
         res = wrapped(*args, **kwargs)
-    except Exception as exc:
+    except Exception:
         if subsegment is not None:
             exception = sys.exc_info()[1]
             stack = stacktrace.get_stacktrace(limit=xray_recorder._max_trace_back)
             subsegment.add_exception(exception, stack)
-        six.raise_from(exc, exc)
+        raise
     finally:
         if subsegment is not None:
             subsegment.set_sql(sql)

aws_xray_sdk/ext/util.py

@@ -35,7 +35,6 @@ def inject_trace_header(headers, entity):
     else:
         header = entity.get_origin_trace_header()
     data = header.data if header else None
-
     to_insert = TraceHeader(
         root=entity.trace_id,
         parent=entity.id,

aws_xray_sdk/version.py

@@ -1 +1 @@
-VERSION = '2.10.0'
+VERSION = '2.11.0'

docs/conf.py

@@ -62,9 +62,9 @@
 # built documents.
 #
 # The short X.Y version.
-version = u'2.10.0'
+version = u'2.11.0'
 # The full version, including alpha/beta/rc tags.
-release = u'2.10.0'
+release = u'2.11.0'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.

docs/frameworks.rst

@@ -39,7 +39,7 @@ The default values are as follows::
     XRAY_RECORDER = {
         'AWS_XRAY_DAEMON_ADDRESS': '127.0.0.1:2000',
         'AUTO_INSTRUMENT': True,  # If turned on built-in database queries and template rendering will be recorded as subsegments
-        'AWS_XRAY_CONTEXT_MISSING': 'RUNTIME_ERROR',
+        'AWS_XRAY_CONTEXT_MISSING': 'LOG_ERROR',
         'PLUGINS': (),
         'SAMPLING': True,
         'SAMPLING_RULES': None,

sample-apps/flask/requirements.txt

@@ -1,5 +1,5 @@
 boto3
-certifi==2020.4.5.1
+certifi==2022.12.7
 chardet==3.0.4
 click==7.1.2
 Flask==1.1.2

tests/ext/django/test_db.py

@@ -10,8 +10,7 @@
 @pytest.fixture(scope='module', autouse=True)
 def setup():
     django.setup()
-    xray_recorder.configure(context=Context(),
-                            context_missing='LOG_ERROR')
+    xray_recorder.configure(context=Context())
     patch_db()
 
 

tests/ext/django/test_middleware.py

@@ -15,8 +15,7 @@ class XRayTestCase(TestCase):
 
     def setUp(self):
         django.setup()
-        xray_recorder.configure(context=Context(),
-                                context_missing='LOG_ERROR')
+        xray_recorder.configure(context=Context())
         xray_recorder.clear_trace_entities()
         global_sdk_config.set_sdk_enabled(True)
 

tests/test_facade_segment.py

@@ -55,3 +55,18 @@ def test_structure_intact():
 
     assert segment.subsegments[0] is subsegment
     assert subsegment.subsegments[0] is subsegment2
+
+def test_adding_unsampled_subsegment():
+
+    segment = FacadeSegment('name', 'id', 'id', True)
+    subsegment = Subsegment('sampled', 'local', segment)
+    subsegment2 = Subsegment('unsampled', 'local', segment)
+    subsegment2.sampled = False
+
+    segment.add_subsegment(subsegment)
+    subsegment.add_subsegment(subsegment2)
+
+
+    assert segment.subsegments[0] is subsegment
+    assert subsegment.subsegments[0] is subsegment2
+    assert subsegment2.sampled == False