Instrument s3transfer BoundExecutor submit function. (#1185)

* Instrument s3transfer BoundExecutor submit function.

Co-authored-by: Tim Pansino <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>

* Add s3transfer hooks files.

Co-authored-by: Tim Pansino <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>

* Remove newlines.

* Fixup: trailing whitespace

---------

Co-authored-by: Tim Pansino <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

newrelic/config.py

@@ -4540,6 +4540,12 @@ def _process_module_builtin_defaults():
         "instrument_botocore_client",
     )
 
+    _process_module_definition(
+        "s3transfer.futures",
+        "newrelic.hooks.external_s3transfer",
+        "instrument_s3transfer_futures",
+    )
+
     _process_module_definition(
         "tornado.httpserver",
         "newrelic.hooks.framework_tornado",

newrelic/hooks/external_s3transfer.py

@@ -0,0 +1,34 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from newrelic.api.time_trace import current_trace
+from newrelic.common.signature import bind_args
+from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.core.context import context_wrapper
+
+
+def instrument_s3transfer_futures(module):
+    if hasattr(module, "BoundedExecutor"):
+        wrap_function_wrapper(module, "BoundedExecutor.submit", wrap_BoundedExecutor_submit)
+
+
+def wrap_BoundedExecutor_submit(wrapped, instance, args, kwargs):
+    trace = current_trace()
+    if not trace:
+        return wrapped(*args, **kwargs)
+
+    bound_args = bind_args(wrapped, args, kwargs)
+    bound_args["task"] = context_wrapper(bound_args["task"], trace=trace, strict=True)
+
+    return wrapped(**bound_args)

tests/external_botocore/_test_file.txt

@@ -0,0 +1 @@
+Hello World!

tests/external_botocore/test_boto3_s3.py

@@ -27,15 +27,14 @@
 from newrelic.common.package_version_utils import get_package_version_tuple
 
 MOTO_VERSION = get_package_version_tuple("moto")
+BOTOCORE_VERSION = get_package_version_tuple("botocore")
+
 AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY"
 AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY"  # nosec
 AWS_REGION_NAME = "us-west-2"
 
 TEST_BUCKET = "python-agent-test-%s" % uuid.uuid4()
 
-BOTOCORE_VERSION = tuple(map(int, botocore.__version__.split(".")))
-
-
 if BOTOCORE_VERSION < (1, 7, 41):
     S3_URL = "s3-us-west-2.amazonaws.com"
     EXPECTED_BUCKET_URL = "https://%s/%s" % (S3_URL, TEST_BUCKET)

tests/external_botocore/test_s3transfer.py

@@ -0,0 +1,84 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import uuid
+
+import boto3
+import botocore
+from moto import mock_aws
+from testing_support.fixtures import dt_enabled
+from testing_support.validators.validate_span_events import validate_span_events
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
+from newrelic.common.package_version_utils import get_package_version_tuple
+
+MOTO_VERSION = get_package_version_tuple("moto")
+BOTOCORE_VERSION = get_package_version_tuple("botocore")
+
+AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY"
+AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY"  # nosec
+AWS_REGION_NAME = "us-west-2"
+
+TEST_BUCKET = "python-agent-test-%s" % uuid.uuid4()
+
+if BOTOCORE_VERSION < (1, 7, 41):
+    S3_URL = "s3-us-west-2.amazonaws.com"
+    EXPECTED_BUCKET_URL = "https://%s/%s" % (S3_URL, TEST_BUCKET)
+    EXPECTED_KEY_URL = EXPECTED_BUCKET_URL + "/hello_world"
+elif BOTOCORE_VERSION < (1, 28):
+    S3_URL = "s3.us-west-2.amazonaws.com"
+    EXPECTED_BUCKET_URL = "https://%s/%s" % (S3_URL, TEST_BUCKET)
+    EXPECTED_KEY_URL = EXPECTED_BUCKET_URL + "/hello_world"
+else:
+    S3_URL = "%s.s3.us-west-2.amazonaws.com" % TEST_BUCKET
+    EXPECTED_BUCKET_URL = "https://%s/" % S3_URL
+    EXPECTED_KEY_URL = EXPECTED_BUCKET_URL + "hello_world"
+
+
+@dt_enabled
+@validate_span_events(exact_agents={"aws.operation": "CreateBucket"}, count=1)
+@validate_span_events(exact_agents={"aws.operation": "PutObject"}, count=1)
+@validate_transaction_metrics(
+    "test_s3transfer:test_s3_context_propagation",
+    scoped_metrics=[
+        ("External/%s/botocore/PUT" % S3_URL, 2),
+    ],
+    rollup_metrics=[
+        ("External/all", 2),
+        ("External/allOther", 2),
+        ("External/%s/all" % S3_URL, 2),
+        ("External/%s/botocore/PUT" % S3_URL, 2),
+    ],
+    background_task=True,
+)
+@background_task()
+@mock_aws
+def test_s3_context_propagation():
+    client = boto3.client(
+        "s3",
+        aws_access_key_id=AWS_ACCESS_KEY_ID,
+        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
+        region_name=AWS_REGION_NAME,
+    )
+
+    # Create bucket
+    resp = client.create_bucket(Bucket=TEST_BUCKET, CreateBucketConfiguration={"LocationConstraint": AWS_REGION_NAME})
+    assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # Upload file
+    client.upload_file(Filename="_test_file.txt", Bucket=TEST_BUCKET, Key="_test_file.txt")
+    # No return value to check for this function currently