|
| 1 | +# Copyright 2010 New Relic, Inc. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | + |
| 16 | +import logging |
| 17 | + |
| 18 | +from newrelic.api.application import application_instance |
| 19 | +from newrelic.api.function_trace import FunctionTrace |
| 20 | +from newrelic.api.transaction import current_transaction |
| 21 | +from newrelic.api.web_transaction import WebTransaction |
| 22 | +from newrelic.common.object_wrapper import wrap_function_wrapper |
| 23 | + |
| 24 | +_logger = logging.getLogger(__name__) |
| 25 | + |
| 26 | +CLIENT_ATTRIBUTES_DEPLOY_RESOURCE_LOG_MSG = "Exception occurred in PyZeebe instrumentation: Failed to extract resource count/file for method `deploy_resource`. Report this issue to New Relic support." |
| 27 | + |
| 28 | + |
| 29 | +# Adds client method params as txn or span attributes |
| 30 | +def _add_client_input_attributes(method_name, trace, args, kwargs): |
| 31 | + bpmn_id = extract_agent_attribute_from_methods( |
| 32 | + args, kwargs, method_name, ("run_process", "run_process_with_result"), "bpmn_process_id", 0 |
| 33 | + ) |
| 34 | + if bpmn_id: |
| 35 | + trace._add_agent_attribute("zeebe.client.bpmnProcessId", bpmn_id) |
| 36 | + |
| 37 | + msg_name = extract_agent_attribute_from_methods(args, kwargs, method_name, ("publish_message"), "name", 0) |
| 38 | + if msg_name: |
| 39 | + trace._add_agent_attribute("zeebe.client.messageName", msg_name) |
| 40 | + |
| 41 | + correlation_key = extract_agent_attribute_from_methods( |
| 42 | + args, kwargs, method_name, ("publish_message"), "correlation_key", 1 |
| 43 | + ) |
| 44 | + if correlation_key: |
| 45 | + trace._add_agent_attribute("zeebe.client.correlationKey", correlation_key) |
| 46 | + |
| 47 | + message_id = extract_agent_attribute_from_methods(args, kwargs, method_name, ("publish_message"), "message_id", 4) |
| 48 | + if message_id: |
| 49 | + trace._add_agent_attribute("zeebe.client.messageId", message_id) |
| 50 | + |
| 51 | + resource = extract_agent_attribute_from_methods(args, {}, method_name, ("deploy_resource"), None, 0) |
| 52 | + if resource: |
| 53 | + try: |
| 54 | + trace._add_agent_attribute("zeebe.client.resourceFile", resource) |
| 55 | + trace._add_agent_attribute("zeebe.client.resourceCount", len(list(args))) |
| 56 | + except Exception: |
| 57 | + _logger.warning(CLIENT_ATTRIBUTES_DEPLOY_RESOURCE_LOG_MSG, exc_info=True) |
| 58 | + |
| 59 | + |
| 60 | +def extract_agent_attribute_from_methods(args, kwargs, method_name, methods, param, index): |
| 61 | + try: |
| 62 | + if method_name in methods: |
| 63 | + value = kwargs.get(param) |
| 64 | + if not value and args and len(args) > index: |
| 65 | + value = args[index] |
| 66 | + return value |
| 67 | + except Exception: |
| 68 | + _logger.warning( |
| 69 | + "Exception occurred in PyZeebe instrumentation: failed to extract %s from %s. Report this issue to New Relic support.", |
| 70 | + param, |
| 71 | + method_name, |
| 72 | + exc_info=True, |
| 73 | + ) |
| 74 | + |
| 75 | + |
| 76 | +# Async wrapper that instruments router/worker annotations` |
| 77 | +async def _nr_wrapper_execute_one_job(wrapped, instance, args, kwargs): |
| 78 | + job = args[0] if args else kwargs.get("job") |
| 79 | + process_id = getattr(job, "bpmn_process_id", None) or "UnknownProcess" |
| 80 | + task_type = getattr(job, "type", None) or "UnknownType" |
| 81 | + txn_name = f"{process_id}/{task_type}" |
| 82 | + |
| 83 | + with WebTransaction(application_instance(), txn_name, group="ZeebeTask") as txn: |
| 84 | + if job is not None: |
| 85 | + if hasattr(job, "key"): |
| 86 | + txn.add_custom_attribute("zeebe.job.key", job.key) |
| 87 | + if hasattr(job, "type"): |
| 88 | + txn.add_custom_attribute("zeebe.job.type", job.type) |
| 89 | + if hasattr(job, "bpmn_process_id"): |
| 90 | + txn.add_custom_attribute("zeebe.job.bpmnProcessId", job.bpmn_process_id) |
| 91 | + if hasattr(job, "process_instance_key"): |
| 92 | + txn.add_custom_attribute("zeebe.job.processInstanceKey", job.process_instance_key) |
| 93 | + if hasattr(job, "element_id"): |
| 94 | + txn.add_custom_attribute("zeebe.job.elementId", job.element_id) |
| 95 | + |
| 96 | + return await wrapped(*args, **kwargs) |
| 97 | + |
| 98 | + |
| 99 | +# Async wrapper that instruments a ZeebeClient method. |
| 100 | +def _nr_client_wrapper(method_name): |
| 101 | + async def _client_wrapper(wrapped, instance, args, kwargs): |
| 102 | + txn = current_transaction() |
| 103 | + if not txn: |
| 104 | + return await wrapped(*args, **kwargs) |
| 105 | + |
| 106 | + with FunctionTrace(name=method_name, group="ZeebeClient") as trace: |
| 107 | + _add_client_input_attributes(method_name, trace, args, kwargs) |
| 108 | + return await wrapped(*args, **kwargs) |
| 109 | + |
| 110 | + return _client_wrapper |
| 111 | + |
| 112 | + |
| 113 | +# Instrument JobExecutor.execute_one_job to create a background transaction per job (invoked from @router.task or @worker.task annotations) |
| 114 | +def instrument_pyzeebe_worker_job_executor(module): |
| 115 | + if hasattr(module, "JobExecutor"): |
| 116 | + wrap_function_wrapper(module, "JobExecutor.execute_one_job", _nr_wrapper_execute_one_job) |
| 117 | + |
| 118 | + |
| 119 | +# Instrument ZeebeClient methods to trace client calls. |
| 120 | +def instrument_pyzeebe_client_client(module): |
| 121 | + target_methods = ("run_process", "run_process_with_result", "deploy_resource", "publish_message") |
| 122 | + |
| 123 | + for method_name in target_methods: |
| 124 | + if hasattr(module, "ZeebeClient"): |
| 125 | + if hasattr(module.ZeebeClient, method_name): |
| 126 | + wrap_function_wrapper(module, f"ZeebeClient.{method_name}", _nr_client_wrapper(method_name)) |
![dependabot[bot]](https://avatars.githubusercontent.com/in/29110?v=4&size=40){kind=avatar}
![mergify[bot]](https://avatars.githubusercontent.com/in/10562?v=4&size=40){kind=avatar}
0 commit comments