|
2 | 2 |
|
3 | 3 | from contextlib import asynccontextmanager, contextmanager
|
4 | 4 | from numbers import Number
|
5 |
| -from typing import Any, AsyncGenerator, Generator, Literal, Sequence, Union |
| 5 | +from typing import Any, AsyncGenerator, Generator, Literal, Optional, Sequence, Union |
6 | 6 |
|
7 | 7 | from aws_lambda_powertools.shared import constants
|
8 | 8 | from aws_lambda_powertools.shared.lazy_import import LazyLoader
|
@@ -69,14 +69,26 @@ def record_exception(self, exception: BaseException, **kwargs):
|
69 | 69 |
|
70 | 70 |
|
71 | 71 | class AwsXrayProvider(BaseProvider):
|
72 |
| - def __init__(self, xray_recorder=None): |
73 |
| - if not xray_recorder: |
74 |
| - from aws_xray_sdk.core import xray_recorder |
| 72 | + |
| 73 | + def __init__( |
| 74 | + self, |
| 75 | + service: str = "", |
| 76 | + disabled: Optional[bool] = None, |
| 77 | + auto_patch: Optional[bool] = None, |
| 78 | + patch_modules: Optional[Sequence[str]] = None, |
| 79 | + ): |
| 80 | + from aws_xray_sdk.core import xray_recorder # type: ignore |
75 | 81 |
|
76 | 82 | self.recorder = xray_recorder
|
77 | 83 | self.in_subsegment = self.recorder.in_subsegment
|
78 | 84 | self.in_subsegment_async = self.recorder.in_subsegment_async
|
79 | 85 |
|
| 86 | + self.service = service |
| 87 | + |
| 88 | + super().__init__( |
| 89 | + service=self.service, |
| 90 | + ) |
| 91 | + |
80 | 92 | @contextmanager
|
81 | 93 | def trace(self, name: str, **kwargs) -> Generator[XraySpan, None, None]:
|
82 | 94 | with self.in_subsegment(name=name, **kwargs) as sub_segment:
|
|
0 commit comments