Repository: DataDog/dd-trace-py

Issue:
`\n` stripped from kinesis JSON records when distributed tracing is enabled
### Which version of dd-trace-py are you using?

1.4.1 through `datadog-lambda-python` (https://github.com/DataDog/datadog-lambda-python/blob/main/poetry.lock#L111)

### Which version of pip are you using?

What's used to build `arn:aws:lambda:us-east-1:464622532012:layer:Datadog-Python38:62` layer.

### Which version of the libraries are you using?

What's referenced on `arn:aws:lambda:us-east-1:464622532012:layer:Datadog-Python38:62` layer.

See [datadog-lambda-python](https://github.com/DataDog/datadog-lambda-python).

### How can we reproduce your problem?

1. Create a lambda that can write to a kinesis stream using `boto3`
2. Use the [serverless](https://github.com/DataDog/serverless-plugin-datadog) plugin with default configuration.
3. Deploy the lambda
4. Run the lambda and write the following on the target kinesis stream: 

```python
kinesis_client = boto3.client('kinesis')
data = { 'foo': 1 }
kinesis_client.put_records(
  StreamName='my_stream', 
  Records=[{ 'Data': bytearray(f'{json.dumps(data)}\n', 'utf-8') }]
)
```

### What is the result that you get?

The record actually written to the kinesis stream does not include the `\n` at the end. 

### What is the result that you expected?

The record actually written to the kinesis stream should include the `\n` at the end. 

### Why is this an issue?

Our kinesis streams are ingested through Firehose and written to S3. From S3, they are read from databricks. Databricks (and other like athena), expect JSON records to be separated by a new line (see: https://spark.apache.org/docs/latest/sql-data-sources-json.html and https://jsonlines.org/). That's why we are currently writing them with a `\n` at the end. 

After adding the DataDog lambda layer, our pipeline broke because the `\n` was no longer honored, as it's discarded when read here: https://github.com/DataDog/dd-trace-py/blob/b4449028209854654291cb762d02baf4bd1cf8cd/ddtrace/contrib/botocore/patch.py#L168 

The `\n` to separate JSONs is a common enough use case that has been recently supported by AWS natively when using [Firehose dynamic partitioning](https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html#dynamic-partitioning-new-line-delimiter), which regrettably we can't use at this time.

We could add a separate transformation step from firehose => s3 and append the newline, but the behavior was unexpected, as we didn't assume the tracer would be changing our kinesis streams data payloads, and regrettably the issue went undetected for a while.

Opening the ticket to check if this use case is something you would want to consider and properly support, and discuss alternatives (from adding a configuration option to append the new line, to check for it during decoding and attempt to preserve, or providing hooks to decode/encode the message + tracing data).




Discussion:

Resolving patch:
diff --git a/ddtrace/contrib/botocore/patch.py b/ddtrace/contrib/botocore/patch.py
--- a/ddtrace/contrib/botocore/patch.py
+++ b/ddtrace/contrib/botocore/patch.py
@@ -10,6 +10,7 @@
 from typing import Dict
 from typing import List
 from typing import Optional
+from typing import Tuple
 
 import botocore.client
 import botocore.exceptions
@@ -45,6 +46,8 @@
 MAX_KINESIS_DATA_SIZE = 1 << 20  # 1MB
 MAX_EVENTBRIDGE_DETAIL_SIZE = 1 << 18  # 256KB
 
+LINE_BREAK = "\n"
+
 log = get_logger(__name__)
 
 
@@ -176,8 +179,17 @@ def inject_trace_to_eventbridge_detail(params, span):
         entry["Detail"] = detail_json
 
 
+def get_json_from_str(data_str):
+    # type: (str) -> Tuple[str, Optional[Dict[str, Any]]]
+    data_obj = json.loads(data_str)
+
+    if data_str.endswith(LINE_BREAK):
+        return LINE_BREAK, data_obj
+    return "", data_obj
+
+
 def get_kinesis_data_object(data):
-    # type: (str) -> Optional[Dict[str, Any]]
+    # type: (str) -> Tuple[str, Optional[Dict[str, Any]]]
     """
     :data: the data from a kinesis stream
 
@@ -190,13 +202,14 @@ def get_kinesis_data_object(data):
 
     # check if data is a json string
     try:
-        return json.loads(data)
+        return get_json_from_str(data)
     except ValueError:
         pass
 
     # check if data is a base64 encoded json string
     try:
-        return json.loads(base64.b64decode(data).decode("ascii"))
+        data_str = base64.b64decode(data).decode("ascii")
+        return get_json_from_str(data_str)
     except ValueError:
         raise TraceInjectionDecodingError("Unable to parse kinesis streams data string")
 
@@ -216,11 +229,15 @@ def inject_trace_to_kinesis_stream_data(record, span):
         return
 
     data = record["Data"]
-    data_obj = get_kinesis_data_object(data)
+    line_break, data_obj = get_kinesis_data_object(data)
     data_obj["_datadog"] = {}
     HTTPPropagator.inject(span.context, data_obj["_datadog"])
     data_json = json.dumps(data_obj)
 
+    # if original string had a line break, add it back
+    if line_break:
+        data_json += line_break
+
     # check if data size will exceed max size with headers
     data_size = len(data_json)
     if data_size >= MAX_KINESIS_DATA_SIZE:


Test patch:
diff --git a/tests/contrib/botocore/test.py b/tests/contrib/botocore/test.py
--- a/tests/contrib/botocore/test.py
+++ b/tests/contrib/botocore/test.py
@@ -1795,6 +1795,124 @@ def test_kinesis_put_records_base64_trace_injection(self):
 
         client.delete_stream(StreamName=stream_name)
 
+    @mock_kinesis
+    def test_kinesis_put_records_newline_json_trace_injection(self):
+        client = self.session.create_client("kinesis", region_name="us-east-1")
+
+        stream_name = "test"
+        client.create_stream(StreamName=stream_name, ShardCount=1)
+        stream = client.describe_stream(StreamName=stream_name)["StreamDescription"]
+        shard_id = stream["Shards"][0]["ShardId"]
+
+        partition_key = "1234"
+        data = [
+            {"Data": json.dumps({"Hello": "World"}) + "\n", "PartitionKey": partition_key},
+            {"Data": json.dumps({"foo": "bar"}) + "\n", "PartitionKey": partition_key},
+        ]
+
+        Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(client)
+        client.put_records(StreamName=stream_name, Records=data)
+
+        # check if the appropriate span was generated
+        spans = self.get_spans()
+        assert spans
+        span = spans[0]
+        assert len(spans) == 1
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.operation") == "PutRecords"
+        assert span.get_tag("params.MessageBody") is None
+        assert_is_measured(span)
+        assert_span_http_status_code(span, 200)
+        assert span.service == "test-botocore-tracing.kinesis"
+        assert span.resource == "kinesis.putrecords"
+        records = span.get_tag("params.Records")
+        assert records is None
+
+        resp = client.get_shard_iterator(StreamName=stream_name, ShardId=shard_id, ShardIteratorType="TRIM_HORIZON")
+        shard_iterator = resp["ShardIterator"]
+
+        # ensure headers are present in received message
+        resp = client.get_records(ShardIterator=shard_iterator)
+        assert len(resp["Records"]) == 2
+        records = resp["Records"]
+        record = records[0]
+        data_str = record["Data"].decode("ascii")
+        assert data_str.endswith("\n")
+        data = json.loads(data_str)
+        headers = data["_datadog"]
+        assert headers is not None
+        assert headers[HTTP_HEADER_TRACE_ID] == str(span.trace_id)
+        assert headers[HTTP_HEADER_PARENT_ID] == str(span.span_id)
+
+        record = records[1]
+        data_str = record["Data"].decode("ascii")
+        assert data_str.endswith("\n")
+        data = json.loads(data_str)
+        assert "_datadog" not in data
+
+        client.delete_stream(StreamName=stream_name)
+
+    @mock_kinesis
+    def test_kinesis_put_records_newline_base64_trace_injection(self):
+        client = self.session.create_client("kinesis", region_name="us-east-1")
+
+        stream_name = "test"
+        client.create_stream(StreamName=stream_name, ShardCount=1)
+        stream = client.describe_stream(StreamName=stream_name)["StreamDescription"]
+        shard_id = stream["Shards"][0]["ShardId"]
+
+        partition_key = "1234"
+        sample_string = json.dumps({"Hello": "World"}) + "\n"
+        sample_string_bytes = sample_string.encode("ascii")
+        base64_bytes = base64.b64encode(sample_string_bytes)
+        data_str = base64_bytes.decode("ascii")
+        data = [
+            {"Data": data_str, "PartitionKey": partition_key},
+            {"Data": data_str, "PartitionKey": partition_key},
+        ]
+
+        Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(client)
+        client.put_records(StreamName=stream_name, Records=data)
+
+        # check if the appropriate span was generated
+        spans = self.get_spans()
+        assert spans
+        span = spans[0]
+        assert len(spans) == 1
+        assert span.get_tag("aws.region") == "us-east-1"
+        assert span.get_tag("aws.operation") == "PutRecords"
+        assert span.get_tag("params.MessageBody") is None
+        assert_is_measured(span)
+        assert_span_http_status_code(span, 200)
+        assert span.service == "test-botocore-tracing.kinesis"
+        assert span.resource == "kinesis.putrecords"
+        records = span.get_tag("params.Records")
+        assert records is None
+
+        resp = client.get_shard_iterator(StreamName=stream_name, ShardId=shard_id, ShardIteratorType="TRIM_HORIZON")
+        shard_iterator = resp["ShardIterator"]
+
+        # ensure headers are present in received message
+        resp = client.get_records(ShardIterator=shard_iterator)
+        assert len(resp["Records"]) == 2
+        records = resp["Records"]
+        record = records[0]
+        data_str = record["Data"].decode("ascii")
+        assert data_str.endswith("\n")
+        data = json.loads(data_str)
+        headers = data["_datadog"]
+        assert headers is not None
+        assert headers[HTTP_HEADER_TRACE_ID] == str(span.trace_id)
+        assert headers[HTTP_HEADER_PARENT_ID] == str(span.span_id)
+
+        record = records[1]
+        data_str = base64.b64decode(record["Data"]).decode("ascii")
+        assert data_str.endswith("\n")
+        data = json.loads(data_str)
+        assert "_datadog" not in data
+
+        client.delete_stream(StreamName=stream_name)
+
     @unittest.skipIf(PY2, "Skipping for Python 2.7 since older moto doesn't support secretsmanager")
     def test_secretsmanager(self):
         from moto import mock_secretsmanager