Skip to content

Commit 2520bd5

Browse files
committed
add enable env for tracing
1 parent 3afb629 commit 2520bd5

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py

+14-3
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@
4949
)
5050

5151

52+
class RestTraceMode(ConfigEnum):
53+
ENABLED = auto()
54+
DISABLED = auto()
55+
56+
5257
class RestSinkEndpoint(ConfigEnum):
5358
RESTLI = auto()
5459
OPENAPI = auto()
@@ -75,9 +80,16 @@ class RestSinkMode(ConfigEnum):
7580
)
7681

7782

83+
_DEFAULT_REST_TRACE_MODE = pydantic.parse_obj_as(
84+
RestTraceMode,
85+
os.getenv("DATAHUB_REST_TRACE_MODE", RestTraceMode.DISABLED),
86+
)
87+
88+
7889
class DatahubRestSinkConfig(DatahubClientConfig):
7990
mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
8091
endpoint: RestSinkEndpoint = _DEFAULT_REST_SINK_ENDPOINT
92+
default_trace_mode: RestTraceMode = _DEFAULT_REST_TRACE_MODE
8193

8294
# These only apply in async modes.
8395
max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -184,9 +196,8 @@ def _make_emitter(cls, config: DatahubRestSinkConfig) -> DataHubRestEmitter:
184196
ca_certificate_path=config.ca_certificate_path,
185197
client_certificate_path=config.client_certificate_path,
186198
disable_ssl_verification=config.disable_ssl_verification,
187-
openapi_ingestion=True
188-
if config.endpoint == RestSinkEndpoint.OPENAPI
189-
else False,
199+
openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
200+
default_trace_mode=config.default_trace_mode == RestTraceMode.ENABLED,
190201
)
191202

192203
@property

0 commit comments

Comments
 (0)