From 5546c2391097103c74688c66b0dd57f985194d8f Mon Sep 17 00:00:00 2001 From: Zachary Groves <32471391+ZStriker19@users.noreply.github.com> Date: Mon, 23 Sep 2024 10:44:12 -0400 Subject: [PATCH] fix(botocore): account for None data_obj case in Kinesis (#10628) This accounts for the possibility of None being returned for the data_obj value from the method `get_kinesis_data_object` in `_patched_kinesis_api_call`. ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) --------- Co-authored-by: erikayasuda <153395705+erikayasuda@users.noreply.github.com> --- .../contrib/internal/botocore/services/kinesis.py | 2 +- ddtrace/contrib/internal/botocore/utils.py | 5 ++--- .../kinesis_none_type_fix-4b39f2059184359e.yaml | 4 ++++ tests/contrib/botocore/test.py | 13 +++++++++++-- 4 files changed, 18 insertions(+), 6 deletions(-) create mode 100644 releasenotes/notes/kinesis_none_type_fix-4b39f2059184359e.yaml diff --git a/ddtrace/contrib/internal/botocore/services/kinesis.py b/ddtrace/contrib/internal/botocore/services/kinesis.py index 1e8d4972bc..2d60252bdc 100644 --- a/ddtrace/contrib/internal/botocore/services/kinesis.py +++ b/ddtrace/contrib/internal/botocore/services/kinesis.py @@ -102,7 +102,7 @@ def _patched_kinesis_api_call(parent_ctx, original_func, instance, args, kwargs, parent_ctx, params, time_estimate, - data_obj.get("_datadog"), + data_obj.get("_datadog") if data_obj else None, record, result, config.botocore.propagation_enabled, diff --git a/ddtrace/contrib/internal/botocore/utils.py b/ddtrace/contrib/internal/botocore/utils.py index dba7f5f703..664bc5d774 100644 --- a/ddtrace/contrib/internal/botocore/utils.py +++ b/ddtrace/contrib/internal/botocore/utils.py @@ -29,7 +29,7 @@ def get_json_from_str(data_str: str) -> Tuple[str, Optional[Dict[str, Any]]]: return None, data_obj -def get_kinesis_data_object(data: str) -> Tuple[str, Optional[Dict[str, Any]]]: +def get_kinesis_data_object(data: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: """ :data: the data from a kinesis stream The data from a kinesis stream comes as a string (could be json, base64 encoded, etc.) @@ -37,9 +37,8 @@ def get_kinesis_data_object(data: str) -> Tuple[str, Optional[Dict[str, Any]]]: - json string - byte encoded json string - base64 encoded json string - If it's none of these, then we leave the message as it is. + If it's none of these, then we return None """ - # check if data is a json string try: return get_json_from_str(data) diff --git a/releasenotes/notes/kinesis_none_type_fix-4b39f2059184359e.yaml b/releasenotes/notes/kinesis_none_type_fix-4b39f2059184359e.yaml new file mode 100644 index 0000000000..f12ab1a0d6 --- /dev/null +++ b/releasenotes/notes/kinesis_none_type_fix-4b39f2059184359e.yaml @@ -0,0 +1,4 @@ +--- +fixes: + - | + kinesis: This fix resolves an issue where unparsable data in a Kinesis record would cause a NoneType error. diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py index 39bdd2d771..c85099fd3e 100644 --- a/tests/contrib/botocore/test.py +++ b/tests/contrib/botocore/test.py @@ -2846,7 +2846,9 @@ def _test_kinesis_put_record_trace_injection(self, test_name, data, client=None, return decoded_record_data - def _test_kinesis_put_records_trace_injection(self, test_name, data, client=None, enable_stream_arn=False): + def _test_kinesis_put_records_trace_injection( + self, test_name, data, client=None, enable_stream_arn=False, verify=True + ): if not client: client = self.session.create_client("kinesis", region_name="us-east-1") @@ -2858,7 +2860,8 @@ def _test_kinesis_put_records_trace_injection(self, test_name, data, client=None client.put_records(StreamName=stream_name, Records=data, StreamARN=stream_arn) else: client.put_records(StreamName=stream_name, Records=data) - + if not verify: + return None # assert commons for span span = self._kinesis_assert_spans() @@ -3249,6 +3252,12 @@ def test_kinesis_put_records_newline_json_trace_injection(self): assert decoded_record_data.endswith("\n") + @mock_kinesis + def test_kinesis_put_records_unparsable_data_object_avoid_nonetype_error(self): + # If the data is unparsable we should not error in tracer code + records = [{"Data": b"", "PartitionKey": "1234"}] + self._test_kinesis_put_records_trace_injection("unparsable_data_obj", records, verify=False) + @mock_kinesis def test_kinesis_put_records_newline_bytes_trace_injection(self): # (dict -> json string -> bytes + new line)[]