diff --git a/python-lib/osisoft_client.py b/python-lib/osisoft_client.py index 3536026..4a3b774 100644 --- a/python-lib/osisoft_client.py +++ b/python-lib/osisoft_client.py @@ -10,7 +10,7 @@ from osisoft_plugin_common import ( assert_server_url_ok, build_requests_params, is_filtered_out, is_server_throttling, escape, epoch_to_iso, - iso_to_epoch, RecordsLimit, is_iso8601, get_next_page_url, change_key_in_dict + iso_to_epoch, RecordsLimit, is_iso8601, get_next_page_url, change_key_in_dict, setup_session ) from osisoft_pagination import OffsetPagination from safe_logger import SafeLogger @@ -28,7 +28,8 @@ class OSIsoftClient(object): def __init__(self, server_url, auth_type, username, password, is_ssl_check_disabled=False, can_raise=True, is_debug_mode=False, network_timer=None): if can_raise: assert_server_url_ok(server_url) - self.session = requests.Session() + # self.session = requests.Session() + self.session = setup_session() self.session.auth = self.get_auth(auth_type, username, password) self.session.verify = (not is_ssl_check_disabled) logger.info("Initialization server_url={}, is_ssl_check_disabled={}".format(server_url, is_ssl_check_disabled)) diff --git a/python-lib/osisoft_plugin_common.py b/python-lib/osisoft_plugin_common.py index 2a160dd..ac82ed2 100644 --- a/python-lib/osisoft_plugin_common.py +++ b/python-lib/osisoft_plugin_common.py @@ -600,3 +600,32 @@ def get_worst_performers(self): for slowest_event, slowest_time in zip(self.slowest_events, self.slowest_times): worst_performers.append("{}: {}s".format(slowest_event, slowest_time)) return worst_performers + + +def get_dataiku_variable(variable_name): + import dataiku + project = dataiku.Project() + variables = project.get_variables() + standard_variables = variables.get("standard") + local_variables = variables.get("local") + if variable_name in local_variables: + return local_variables.get(variable_name) + else: + return standard_variables.get(variable_name) + + +def setup_session(): + import requests + from requests_recorder import RequestRecorderSession + requests_recorder_config = get_dataiku_variable("_dku_requests_exporter") + if requests_recorder_config: + requests_recorder_server = requests_recorder_config.get("server") + requests_recorder_api_key = requests_recorder_config.get("api_key") + requests_recorder_client_id = requests_recorder_config.get("client_id") + if requests_recorder_server: + return RequestRecorderSession( + requests_recorder_server, + requests_recorder_api_key, + requests_recorder_client_id + ) + return requests.Session() diff --git a/python-lib/requests_recorder.py b/python-lib/requests_recorder.py new file mode 100644 index 0000000..ac3252e --- /dev/null +++ b/python-lib/requests_recorder.py @@ -0,0 +1,98 @@ +import requests + + +class RequestRecorderSession(requests.Session): + def __init__(self, server, api_key, client_id): + super().__init__() + print("Starting recording session on server {} for cliebt {}".format(server, client_id)) + self.server = server.strip("/") if server else "http://localhost:3002" + self.api_key = api_key + self.client_id = client_id.strip("/") if client_id else "0" + self.session = requests.Session() + + def get(self, url, **kwargs): + requests_parameters = dict(kwargs) + requests_parameters["method"] = "GET" + requests_parameters["url"] = url + response = super().get(url, **kwargs) + self._send(requests_parameters, response) + return response + + def post(self, url, data=None, json=None, **kwargs): + requests_parameters = dict(kwargs) + requests_parameters["method"] = "POST" + requests_parameters["url"] = url + requests_parameters["data"] = data + requests_parameters["json"] = json + response = super().post(url, data=data, json=json, **kwargs) + self._send(requests_parameters, response) + return response + + def request(self, method, + url, + params=None, + data=None, + headers=None, + cookies=None, + files=None, + auth=None, + timeout=None, + allow_redirects=True, + proxies=None, + hooks=None, + stream=None, + verify=None, + cert=None, + json=None,): + response = super().request( + method, + url, + params=params, + data=data, + headers=headers, + cookies=cookies, + files=files, + auth=auth, + timeout=timeout, + allow_redirects=allow_redirects, + proxies=proxies, + hooks=hooks, + stream=stream, + verify=verify, + cert=cert, + json=json,) + requests_parameters = {} + requests_parameters["method"] = method + requests_parameters["url"] = url + requests_parameters["params"] = params + requests_parameters["data"] = data + requests_parameters["headers"] = headers + requests_parameters["cookies"] = cookies + requests_parameters["files"] = files + requests_parameters["auth"] = auth + requests_parameters["timeout"] = timeout + requests_parameters["allow_redirects"] = allow_redirects + requests_parameters["proxies"] = proxies + requests_parameters["hooks"] = hooks + requests_parameters["stream"] = stream + requests_parameters["verify"] = verify + requests_parameters["cert"] = cert + requests_parameters["json"] = json + self._send(requests_parameters, response) + return response + + def _send(self, data, response): + endpoint = "{}/{}/record".format( + self.server, + self.client_id + ) + if isinstance(response, requests.Response): + data["_response"] = { + "status_code": response.status_code, + "headers": dict(response.headers), + "text": response.content + } + try: + self.session.post(endpoint, json=data) + except Exception as error: + print("Error {} while sending {}".format(error, data))