diff --git a/loki_logger_handler/loki_logger_handler.py b/loki_logger_handler/loki_logger_handler.py index f261ca0..48f0baf 100644 --- a/loki_logger_handler/loki_logger_handler.py +++ b/loki_logger_handler/loki_logger_handler.py @@ -34,8 +34,8 @@ def __init__( enable_self_errors=False, enable_structured_loki_metadata=False, loki_metadata=None, - loki_metadata_keys=None - + loki_metadata_keys=None, + requests_timeout=None, ): """ Initialize the LokiLoggerHandler object. @@ -54,6 +54,7 @@ def __init__( enable_structured_loki_metadata (bool, optional): Whether to include structured loki_metadata in the logs. Defaults to False. Only supported for Loki 3.0 and above loki_metadata (dict, optional): Default loki_metadata values. Defaults to None. Only supported for Loki 3.0 and above loki_metadata_keys (arrray, optional): Specific log record keys to extract as loki_metadata. Only supported for Loki 3.0 and above + requests_timeout (int or tuple, optional): Timeout for requests to the Loki server. Defaults to None, which is the default timeout of the requests library. """ super(LokiLoggerHandler, self).__init__() @@ -72,7 +73,7 @@ def __init__( self.debug_logger.addHandler(console_handler) self.request = LokiRequest( - url=url, compressed=compressed, auth=auth, additional_headers=additional_headers or {} + url=url, compressed=compressed, auth=auth, additional_headers=additional_headers or {}, timeout=requests_timeout ) self.buffer = queue.Queue() diff --git a/loki_logger_handler/loki_request.py b/loki_logger_handler/loki_request.py index b1e1bc3..0f6abc4 100644 --- a/loki_logger_handler/loki_request.py +++ b/loki_logger_handler/loki_request.py @@ -15,7 +15,7 @@ class LokiRequest: session (requests.Session): The session used for making HTTP requests. """ - def __init__(self, url, compressed=False, auth=None, additional_headers=None): + def __init__(self, url, compressed=False, auth=None, additional_headers=None, timeout=None): """ Initialize the LokiRequest object with the server URL, compression option, and additional headers. @@ -32,6 +32,7 @@ def __init__(self, url, compressed=False, auth=None, additional_headers=None): self.headers = additional_headers if additional_headers is not None else {} self.headers["Content-Type"] = "application/json" self.session = requests.Session() + self.timeout = timeout def send(self, data): """ @@ -49,7 +50,7 @@ def send(self, data): self.headers["Content-Encoding"] = "gzip" data = gzip.compress(data.encode("utf-8")) - response = self.session.post(self.url, data=data, auth=self.auth, headers=self.headers) + response = self.session.post(self.url, data=data, auth=self.auth, headers=self.headers, timeout=self.timeout) response.raise_for_status() except requests.RequestException as e: