feat(aws): add partition tag (#14577)

<!--
* New contributors are highly encouraged to read our
  [CONTRIBUTING](/CONTRIBUTING.md) documentation.
* Commit and PR titles should be prefixed with the general area of the
pull request's change.

-->
### What does this PR do?
Adds a `aws.partition` tag onto AWS traces based on the region. The
logic to derive this partition is
[identical](https://github.com/DataDog/dd-trace-go/blob/main/contrib/aws/aws-sdk-go-v2/aws/aws.go#L206-L215)
to what we do within other tracing libraries.

### Motivation
Part of a larger effort to better link cloud resources within Datadog.
We often need the AWS partition during downstream enrichment - it is
almost always a required part of constructing an ARN, but it's currently
not captured in the library

[Jira Ticket](https://datadoghq.atlassian.net/browse/CLP-1021) and more
[context](https://datadoghq.atlassian.net/wiki/x/toGyLwE) around this
effort
<!--
* What inspired you to submit this pull request?
* Link any related GitHub issues or PRs here.
* If this resolves a GitHub issue, include "Fixes #XXXX" to link the
issue and auto-close it on merge.
-->

## 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
- [ ] 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: wantsui <[email protected]>

Author: seungleelee

ddtrace/_trace/utils_botocore/span_tags.py

@@ -92,6 +92,7 @@ def set_botocore_patched_api_call_span_tags(span: Span, instance, args, params,
     if region_name is not None:
         span.set_tag_str("aws.region", region_name)
         span.set_tag_str("region", region_name)
+        span.set_tag_str("aws.partition", aws.get_aws_partition(region_name))
 
         # Derive peer hostname only in serverless environments to avoid
         # unnecessary tag noise in traditional hosts/containers.

ddtrace/contrib/internal/aiobotocore/patch.py

@@ -160,6 +160,10 @@ async def _wrapped_api_call(original_func, instance, args, kwargs):
             "aws.region": region_name,
             "region": region_name,
         }
+
+        if region_name:
+            meta["aws.partition"] = aws.get_aws_partition(region_name)
+
         span.set_tags(meta)
 
         result = await original_func(*args, **kwargs)

ddtrace/contrib/internal/boto/patch.py

@@ -124,6 +124,7 @@ def patched_query_request(original_func, instance, args, kwargs):
         if region_name:
             meta[aws.REGION] = region_name
             meta[aws.AWSREGION] = region_name
+            meta[aws.PARTITION] = aws.get_aws_partition(region_name)
 
             if in_aws_lambda():
                 # Derive the peer hostname now that we have both service and region.
@@ -192,6 +193,7 @@ def patched_auth_request(original_func, instance, args, kwargs):
         if region_name:
             meta[aws.REGION] = region_name
             meta[aws.AWSREGION] = region_name
+            meta[aws.PARTITION] = aws.get_aws_partition(region_name)
 
             if in_aws_lambda():
                 # Derive the peer hostname

ddtrace/ext/aws.py

@@ -82,7 +82,19 @@ def _add_api_param_span_tags(span, endpoint_name, params):
         span.set_tag_str("statemachinearn", state_machine_arn)
 
 
+def get_aws_partition(region_name):
+    # type: (str) -> str
+    """Determine AWS partition from region name."""
+    if region_name.startswith("cn-"):
+        return "aws-cn"
+    elif region_name.startswith("us-gov-"):
+        return "aws-us-gov"
+
+    return "aws"
+
+
 AWSREGION = "aws.region"
 REGION = "region"
+PARTITION = "aws.partition"
 AGENT = "aws.agent"
 OPERATION = "aws.operation"

releasenotes/notes/add-partition-tag-b423877806ab7271.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    AWS: adds ``aws.partition`` tag onto AWS traces based on the region for the boto, botocore, and aiobotocore integrations.
+

tests/contrib/aiobotocore/test.py

@@ -38,6 +38,7 @@ async def test_traced_client(tracer):
     assert span.get_tag("aws.agent") == "aiobotocore"
     assert span.get_tag("aws.region") == "us-west-2"
     assert span.get_tag("region") == "us-west-2"
+    assert span.get_tag("aws.partition") == "aws"
     assert span.get_tag("aws.operation") == "DescribeInstances"
     assert_span_http_status_code(span, 200)
     assert span.get_metric("retry_attempts") == 0
@@ -203,6 +204,7 @@ async def test_sqs_client(tracer):
     assert_is_measured(span)
     assert span.get_tag("aws.region") == "us-west-2"
     assert span.get_tag("region") == "us-west-2"
+    assert span.get_tag("aws.partition") == "aws"
     assert span.get_tag("aws.operation") == "ListQueues"
     assert_span_http_status_code(span, 200)
     assert span.service == "aws.sqs"
@@ -225,6 +227,7 @@ async def test_kinesis_client(tracer):
     assert_is_measured(span)
     assert span.get_tag("aws.region") == "us-west-2"
     assert span.get_tag("region") == "us-west-2"
+    assert span.get_tag("aws.partition") == "aws"
     assert span.get_tag("aws.operation") == "ListStreams"
     assert_span_http_status_code(span, 200)
     assert span.service == "aws.kinesis"

tests/contrib/boto/test.py

@@ -52,6 +52,7 @@ def test_ec2_client(self):
         self.assertEqual(span.get_tag(http.METHOD), "POST")
         self.assertEqual(span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(span.get_tag("region"), "us-west-2")
+        self.assertEqual(span.get_tag("aws.partition"), "aws")
         self.assertEqual(span.get_tag("component"), "boto")
         self.assertEqual(span.get_tag("span.kind"), "client")
 
@@ -67,6 +68,7 @@ def test_ec2_client(self):
         self.assertEqual(span.get_tag(http.METHOD), "POST")
         self.assertEqual(span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(span.get_tag("region"), "us-west-2")
+        self.assertEqual(span.get_tag("aws.partition"), "aws")
         self.assertEqual(span.get_tag("component"), "boto")
         self.assertEqual(span.get_tag("span.kind"), "client")
         self.assertEqual(span.service, "test-boto-tracing.ec2")
@@ -495,6 +497,7 @@ def test_lambda_client(self):
         self.assertEqual(span.get_tag(http.METHOD), "GET")
         self.assertEqual(span.get_tag("aws.region"), "us-east-2")
         self.assertEqual(span.get_tag("region"), "us-east-2")
+        self.assertEqual(span.get_tag("aws.partition"), "aws")
         self.assertEqual(span.get_tag("aws.operation"), "list_functions")
         self.assertEqual(span.get_tag("component"), "boto")
         self.assertEqual(span.get_tag("span.kind"), "client")
@@ -612,6 +615,7 @@ def test_sts_client(self):
         assert_is_measured(span)
         self.assertEqual(span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(span.get_tag("region"), "us-west-2")
+        self.assertEqual(span.get_tag("aws.partition"), "aws")
         self.assertEqual(span.get_tag("aws.operation"), "GetFederationToken")
         self.assertEqual(span.get_tag("component"), "boto")
         self.assertEqual(span.get_tag("span.kind"), "client")
@@ -750,6 +754,7 @@ def test_elasticache_client(self):
         span = spans[0]
         self.assertEqual(span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(span.get_tag("region"), "us-west-2")
+        self.assertEqual(span.get_tag("aws.partition"), "aws")
         self.assertEqual(span.get_tag("component"), "boto")
         self.assertEqual(span.get_tag("span.kind"), "client")
         self.assertEqual(span.service, "test-boto-tracing.elasticache")
@@ -783,6 +788,7 @@ def test_ec2_client_ot(self):
         self.assertEqual(dd_span.get_tag(http.METHOD), "POST")
         self.assertEqual(dd_span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(dd_span.get_tag("region"), "us-west-2")
+        self.assertEqual(dd_span.get_tag("aws.partition"), "aws")
 
         with ot_tracer.start_active_span("ot_span"):
             ec2.run_instances(21)
@@ -800,6 +806,7 @@ def test_ec2_client_ot(self):
         self.assertEqual(dd_span.get_tag(http.METHOD), "POST")
         self.assertEqual(dd_span.get_tag("aws.region"), "us-west-2")
         self.assertEqual(dd_span.get_tag("region"), "us-west-2")
+        self.assertEqual(dd_span.get_tag("aws.partition"), "aws")
         self.assertEqual(dd_span.get_tag("component"), "boto")
         self.assertEqual(dd_span.get_tag("span.kind"), "client")
         self.assertEqual(dd_span.service, "test-boto-tracing.ec2")

tests/contrib/botocore/test.py

@@ -164,6 +164,7 @@ def test_traced_client(self):
         assert span.get_tag("aws.agent") == "botocore"
         assert span.get_tag("aws.region") == "us-west-2"
         assert span.get_tag("region") == "us-west-2"
+        assert span.get_tag("aws.partition") == "aws"
         assert span.get_tag("aws.operation") == "DescribeInstances"
         assert span.get_tag("aws.requestid") == "fdcdcab1-ae5c-489e-9c33-4637c5dda355"
         assert span.get_tag("component") == "botocore"
@@ -718,6 +719,7 @@ def _test_sqs_client(self):
         assert len(spans) == 1
         assert span.get_tag("aws.region") == "us-east-1"
         assert span.get_tag("region") == "us-east-1"
+        assert span.get_tag("aws.partition") == "aws"
         assert span.get_tag("aws.operation") == "CreateQueue"
         assert span.get_tag("component") == "botocore"
         assert_is_measured(span)
@@ -771,6 +773,7 @@ def test_sqs_send_message_distributed_tracing_off(self):
             assert len(spans) == 1
             assert span.get_tag("aws.region") == "us-east-1"
             assert span.get_tag("region") == "us-east-1"
+            assert span.get_tag("aws.partition") == "aws"
             assert span.get_tag("aws.operation") == "SendMessage"
             assert span.get_tag("params.MessageBody") is None
             assert span.get_tag("component") == "botocore"