Skip to content

Commit 03b5457

Browse files
Ferenc-Ferenc-
committed
feat: Add opt-in exit spans to boto3 and aiohttp client
Signed-off-by: Ferenc Géczi <[email protected]>
1 parent 7dfb921 commit 03b5457

File tree

8 files changed

+258
-51
lines changed

8 files changed

+258
-51
lines changed

‎instana/instrumentation/aiohttp/client.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from ...log import logger
99
from ...singletons import agent, async_tracer
1010
from ...util.secrets import strip_secrets_from_query
11+
from ...util.traceutils import tracing_is_off
1112

1213
try:
1314
import aiohttp
@@ -16,14 +17,12 @@
1617

1718
async def stan_request_start(session, trace_config_ctx, params):
1819
try:
19-
parent_span = async_tracer.active_span
20-
2120
# If we're not tracing, just return
22-
if parent_span is None:
21+
if tracing_is_off():
2322
trace_config_ctx.scope = None
2423
return
2524

26-
scope = async_tracer.start_active_span("aiohttp-client", child_of=parent_span)
25+
scope = async_tracer.start_active_span("aiohttp-client", child_of=async_tracer.active_span)
2726
trace_config_ctx.scope = scope
2827

2928
async_tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, params.headers)

‎instana/instrumentation/boto3_inst.py‎

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
from ..log import logger
1010
from ..singletons import tracer, agent
11-
from ..util.traceutils import get_active_tracer
11+
from ..util.traceutils import get_tracer_tuple, tracing_is_off
1212

1313
try:
1414
import opentracing as ot
@@ -47,28 +47,20 @@ def lambda_inject_context(payload, scope):
4747

4848
@wrapt.patch_function_wrapper("botocore.auth", "SigV4Auth.add_auth")
4949
def emit_add_auth_with_instana(wrapped, instance, args, kwargs):
50-
active_tracer = get_active_tracer()
51-
52-
# If we're not tracing, just return;
53-
if active_tracer is None:
54-
return wrapped(*args, **kwargs)
55-
56-
span = active_tracer.active_span
57-
extract_custom_headers(span, args[0].headers)
58-
50+
if not tracing_is_off() and tracer.active_span:
51+
extract_custom_headers(tracer.active_span, args[0].headers)
5952
return wrapped(*args, **kwargs)
6053

6154

6255
@wrapt.patch_function_wrapper('botocore.client', 'BaseClient._make_api_call')
6356
def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
64-
# pylint: disable=protected-access
65-
active_tracer = get_active_tracer()
66-
6757
# If we're not tracing, just return
68-
if active_tracer is None:
58+
if tracing_is_off():
6959
return wrapped(*arg_list, **kwargs)
7060

71-
with active_tracer.start_active_span("boto3", child_of=active_tracer.active_span) as scope:
61+
tracer, parent_span, _ = get_tracer_tuple()
62+
63+
with tracer.start_active_span("boto3", child_of=parent_span) as scope:
7264
try:
7365
operation = arg_list[0]
7466
payload = arg_list[1]
@@ -110,18 +102,17 @@ def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
110102

111103

112104
def s3_inject_method_with_instana(wrapped, instance, arg_list, kwargs):
105+
# If we're not tracing, just return
106+
if tracing_is_off():
107+
return wrapped(*arg_list, **kwargs)
108+
113109
fas = inspect.getfullargspec(wrapped)
114110
fas_args = fas.args
115111
fas_args.remove('self')
116112

117-
# pylint: disable=protected-access
118-
active_tracer = get_active_tracer()
119-
120-
# If we're not tracing, just return
121-
if active_tracer is None:
122-
return wrapped(*arg_list, **kwargs)
113+
tracer, parent_span, _ = get_tracer_tuple()
123114

124-
with active_tracer.start_active_span("boto3", child_of=active_tracer.active_span) as scope:
115+
with tracer.start_active_span("boto3", child_of=parent_span) as scope:
125116
try:
126117
operation = wrapped.__name__
127118
scope.span.set_tag('op', operation)

‎tests/clients/boto3/test_boto3_lambda.py‎

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def setUp(self):
2828
def tearDown(self):
2929
# Stop Moto after each test
3030
self.mock.stop()
31-
31+
agent.options.allow_exit_as_root = False
3232

3333
def test_lambda_invoke(self):
3434
with tracer.start_active_span('test'):
@@ -66,6 +66,32 @@ def test_lambda_invoke(self):
6666
self.assertEqual(boto_span.data['http']['method'], 'POST')
6767
self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
6868

69+
def test_lambda_invoke_as_root_exit_span(self):
70+
agent.options.allow_exit_as_root = True
71+
result = self.aws_lambda.invoke(FunctionName=self.function_name, Payload=json.dumps({"message": "success"}))
72+
73+
self.assertEqual(result["StatusCode"], 200)
74+
result_payload = json.loads(result["Payload"].read().decode("utf-8"))
75+
self.assertIn("message", result_payload)
76+
self.assertEqual("success", result_payload["message"])
77+
78+
spans = self.recorder.queued_spans()
79+
self.assertEqual(1, len(spans))
80+
boto_span = spans[0]
81+
self.assertTrue(boto_span)
82+
self.assertEqual(boto_span.n, "boto3")
83+
self.assertIsNone(boto_span.p)
84+
self.assertIsNone(boto_span.ec)
85+
86+
self.assertEqual(boto_span.data['boto3']['op'], 'Invoke')
87+
endpoint = f'https://lambda.{self.lambda_region}.amazonaws.com'
88+
self.assertEqual(boto_span.data['boto3']['ep'], endpoint)
89+
self.assertEqual(boto_span.data['boto3']['reg'], self.lambda_region)
90+
self.assertIn('FunctionName', boto_span.data['boto3']['payload'])
91+
self.assertEqual(boto_span.data['boto3']['payload']['FunctionName'], self.function_name)
92+
self.assertEqual(boto_span.data['http']['status'], 200)
93+
self.assertEqual(boto_span.data['http']['method'], 'POST')
94+
self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
6995

7096
def test_request_header_capture_before_call(self):
7197
original_extra_http_headers = agent.options.extra_http_headers
@@ -125,7 +151,7 @@ def add_custom_header_before_call(params, **kwargs):
125151
self.assertEqual("this", boto_span.data["http"]["header"]["X-Capture-This"])
126152
self.assertIn("X-Capture-That", boto_span.data["http"]["header"])
127153
self.assertEqual("that", boto_span.data["http"]["header"]["X-Capture-That"])
128-
154+
129155
agent.options.extra_http_headers = original_extra_http_headers
130156

131157

@@ -188,7 +214,7 @@ def add_custom_header_before_sign(request, **kwargs):
188214
self.assertEqual("Value1", boto_span.data["http"]["header"]["X-Custom-1"])
189215
self.assertIn("X-Custom-2", boto_span.data["http"]["header"])
190216
self.assertEqual("Value2", boto_span.data["http"]["header"]["X-Custom-2"])
191-
217+
192218
agent.options.extra_http_headers = original_extra_http_headers
193219

194220

@@ -198,7 +224,7 @@ def test_response_header_capture(self):
198224

199225
# Access the event system on the S3 client
200226
event_system = self.aws_lambda.meta.events
201-
227+
202228
response_headers = {
203229
"X-Capture-This-Too": "this too",
204230
"X-Capture-That-Too": "that too",
@@ -250,5 +276,5 @@ def modify_after_call_args(parsed, **kwargs):
250276
self.assertEqual("this too", boto_span.data["http"]["header"]["X-Capture-This-Too"])
251277
self.assertIn("X-Capture-That-Too", boto_span.data["http"]["header"])
252278
self.assertEqual("that too", boto_span.data["http"]["header"]["X-Capture-That-Too"])
253-
279+
254280
agent.options.extra_http_headers = original_extra_http_headers

‎tests/clients/boto3/test_boto3_s3.py‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ def setUp(self):
3232
def tearDown(self):
3333
# Stop Moto after each test
3434
self.mock.stop()
35+
agent.options.allow_exit_as_root = False
3536

3637

3738
def test_vanilla_create_bucket(self):
@@ -76,6 +77,32 @@ def test_s3_create_bucket(self):
7677
self.assertEqual(boto_span.data['http']['url'], 'https://s3.amazonaws.com:443/CreateBucket')
7778

7879

80+
def test_s3_create_bucket_as_root_exit_span(self):
81+
agent.options.allow_exit_as_root = True
82+
self.s3.create_bucket(Bucket="aws_bucket_name")
83+
84+
agent.options.allow_exit_as_root = False
85+
result = self.s3.list_buckets()
86+
self.assertEqual(1, len(result['Buckets']))
87+
self.assertEqual(result['Buckets'][0]['Name'], 'aws_bucket_name')
88+
89+
spans = self.recorder.queued_spans()
90+
self.assertEqual(1, len(spans))
91+
boto_span = spans[0]
92+
self.assertTrue(boto_span)
93+
self.assertEqual(boto_span.n, "boto3")
94+
self.assertIsNone(boto_span.p)
95+
self.assertIsNone(boto_span.ec)
96+
97+
self.assertEqual(boto_span.data['boto3']['op'], 'CreateBucket')
98+
self.assertEqual(boto_span.data['boto3']['ep'], 'https://s3.amazonaws.com')
99+
self.assertEqual(boto_span.data['boto3']['reg'], 'us-east-1')
100+
self.assertDictEqual(boto_span.data['boto3']['payload'], {'Bucket': 'aws_bucket_name'})
101+
self.assertEqual(boto_span.data['http']['status'], 200)
102+
self.assertEqual(boto_span.data['http']['method'], 'POST')
103+
self.assertEqual(boto_span.data['http']['url'], 'https://s3.amazonaws.com:443/CreateBucket')
104+
105+
79106
def test_s3_list_buckets(self):
80107
with tracer.start_active_span('test'):
81108
result = self.s3.list_buckets()
@@ -282,7 +309,7 @@ def add_custom_header_before_call(params, **kwargs):
282309

283310
with tracer.start_active_span('test'):
284311
self.s3.create_bucket(Bucket="aws_bucket_name")
285-
312+
286313
result = self.s3.list_buckets()
287314
self.assertEqual(1, len(result['Buckets']))
288315
self.assertEqual(result['Buckets'][0]['Name'], 'aws_bucket_name')
@@ -316,7 +343,7 @@ def add_custom_header_before_call(params, **kwargs):
316343
self.assertEqual("this", boto_span.data["http"]["header"]["X-Capture-This"])
317344
self.assertIn("X-Capture-That", boto_span.data["http"]["header"])
318345
self.assertEqual("that", boto_span.data["http"]["header"]["X-Capture-That"])
319-
346+
320347
agent.options.extra_http_headers = original_extra_http_headers
321348

322349

@@ -343,7 +370,7 @@ def add_custom_header_before_sign(request, **kwargs):
343370

344371
with tracer.start_active_span('test'):
345372
self.s3.create_bucket(Bucket="aws_bucket_name")
346-
373+
347374
result = self.s3.list_buckets()
348375
self.assertEqual(1, len(result['Buckets']))
349376
self.assertEqual(result['Buckets'][0]['Name'], 'aws_bucket_name')
@@ -377,7 +404,7 @@ def add_custom_header_before_sign(request, **kwargs):
377404
self.assertEqual("Value1", boto_span.data["http"]["header"]["X-Custom-1"])
378405
self.assertIn("X-Custom-2", boto_span.data["http"]["header"])
379406
self.assertEqual("Value2", boto_span.data["http"]["header"]["X-Custom-2"])
380-
407+
381408
agent.options.extra_http_headers = original_extra_http_headers
382409

383410

@@ -388,7 +415,7 @@ def test_response_header_capture(self):
388415

389416
# Access the event system on the S3 client
390417
event_system = self.s3.meta.events
391-
418+
392419
response_headers = {
393420
"X-Capture-This-Too": "this too",
394421
"X-Capture-That-Too": "that too",
@@ -437,5 +464,5 @@ def modify_after_call_args(parsed, **kwargs):
437464
self.assertEqual("this too", boto_span.data["http"]["header"]["X-Capture-This-Too"])
438465
self.assertIn("X-Capture-That-Too", boto_span.data["http"]["header"])
439466
self.assertEqual("that too", boto_span.data["http"]["header"]["X-Capture-That-Too"])
440-
467+
441468
agent.options.extra_http_headers = original_extra_http_headers

‎tests/clients/boto3/test_boto3_secretsmanager.py‎

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def setUp(self):
2929
def tearDown(self):
3030
# Stop Moto after each test
3131
self.mock.stop()
32+
agent.options.allow_exit_as_root = False
3233

3334

3435
def test_vanilla_list_secrets(self):
@@ -78,6 +79,41 @@ def test_get_secret_value(self):
7879
self.assertEqual(boto_span.data['http']['url'], 'https://secretsmanager.us-east-1.amazonaws.com:443/GetSecretValue')
7980

8081

82+
def test_get_secret_value_as_root_exit_span(self):
83+
secret_id = 'Uber_Password'
84+
85+
response = self.secretsmanager.create_secret(
86+
Name=secret_id,
87+
SecretBinary=b'password1',
88+
SecretString='password1',
89+
)
90+
91+
self.assertEqual(response['Name'], secret_id)
92+
93+
agent.options.allow_exit_as_root = True
94+
result = self.secretsmanager.get_secret_value(SecretId=secret_id)
95+
96+
self.assertEqual(result['Name'], secret_id)
97+
98+
spans = self.recorder.queued_spans()
99+
self.assertEqual(1, len(spans))
100+
101+
boto_span = spans[0]
102+
self.assertTrue(boto_span)
103+
self.assertEqual(boto_span.n, "boto3")
104+
self.assertIsNone(boto_span.p)
105+
self.assertIsNone(boto_span.ec)
106+
107+
self.assertEqual(boto_span.data['boto3']['op'], 'GetSecretValue')
108+
self.assertEqual(boto_span.data['boto3']['ep'], 'https://secretsmanager.us-east-1.amazonaws.com')
109+
self.assertEqual(boto_span.data['boto3']['reg'], 'us-east-1')
110+
self.assertNotIn('payload', boto_span.data['boto3'])
111+
112+
self.assertEqual(boto_span.data['http']['status'], 200)
113+
self.assertEqual(boto_span.data['http']['method'], 'POST')
114+
self.assertEqual(boto_span.data['http']['url'], 'https://secretsmanager.us-east-1.amazonaws.com:443/GetSecretValue')
115+
116+
81117
def test_request_header_capture_before_call(self):
82118
secret_id = 'Uber_Password'
83119

@@ -141,10 +177,10 @@ def add_custom_header_before_call(params, **kwargs):
141177
self.assertEqual("this", boto_span.data["http"]["header"]["X-Capture-This"])
142178
self.assertIn("X-Capture-That", boto_span.data["http"]["header"])
143179
self.assertEqual("that", boto_span.data["http"]["header"]["X-Capture-That"])
144-
180+
145181
agent.options.extra_http_headers = original_extra_http_headers
146182

147-
183+
148184
def test_request_header_capture_before_sign(self):
149185
secret_id = 'Uber_Password'
150186

@@ -209,7 +245,7 @@ def add_custom_header_before_sign(request, **kwargs):
209245
self.assertEqual("Value1", boto_span.data["http"]["header"]["X-Custom-1"])
210246
self.assertIn("X-Custom-2", boto_span.data["http"]["header"])
211247
self.assertEqual("Value2", boto_span.data["http"]["header"]["X-Custom-2"])
212-
248+
213249
agent.options.extra_http_headers = original_extra_http_headers
214250

215251

@@ -229,7 +265,7 @@ def test_response_header_capture(self):
229265

230266
# Access the event system on the S3 client
231267
event_system = self.secretsmanager.meta.events
232-
268+
233269
response_headers = {
234270
"X-Capture-This-Too": "this too",
235271
"X-Capture-That-Too": "that too",
@@ -276,5 +312,5 @@ def modify_after_call_args(parsed, **kwargs):
276312
self.assertEqual("this too", boto_span.data["http"]["header"]["X-Capture-This-Too"])
277313
self.assertIn("X-Capture-That-Too", boto_span.data["http"]["header"])
278314
self.assertEqual("that too", boto_span.data["http"]["header"]["X-Capture-That-Too"])
279-
315+
280316
agent.options.extra_http_headers = original_extra_http_headers

0 commit comments

Comments
 (0)