diff --git a/log_forwarder/clients.py b/log_forwarder/clients.py index 0c8b774..e008923 100644 --- a/log_forwarder/clients.py +++ b/log_forwarder/clients.py @@ -1,7 +1,9 @@ import time +import random import boto3 import urllib.request +import urllib.error import gzip import logging from typing import Dict @@ -10,6 +12,27 @@ logger = logging.getLogger() +class BrontoClientError(Exception): + """Base exception for BrontoClient errors""" + pass + + +class BrontoClientNonRetryableError(BrontoClientError): + """Exception raised for HTTP errors that should not be retried""" + def __init__(self, status_code: int, reason: str): + self.status_code = status_code + self.reason = reason + super().__init__(f'Non-retryable HTTP error {status_code}: {reason}') + + +class BrontoClientMaxAttemptsError(BrontoClientError): + """Exception raised when maximum retry attempts are reached""" + def __init__(self, attempts: int, last_status: int = None): + self.attempts = attempts + self.last_status = last_status + super().__init__(f'Max retry attempts ({attempts}) reached') + + class S3Client: def __init__(self, filepath): @@ -54,12 +77,21 @@ def reset(self): class BrontoClient: - def __init__(self, api_key, ingestion_endpoint, dataset, collection, client_type, tags: Dict[str, str]): + NON_RETRYABLE_STATUS_CODES = {400, 401, 403, 404, 405, 410, 422} + DEFAULT_MAX_ATTEMPTS = 5 + DEFAULT_BASE_DELAY_SEC = 2 + DEFAULT_TIMEOUT_SEC = 30 + + def __init__(self, api_key, ingestion_endpoint, dataset, collection, client_type, tags: Dict[str, str], + max_attempts: int = None, base_delay_sec: int = None, timeout_sec: int = None): self.api_key = api_key self.dataset = dataset self.collection = collection self.client_type = client_type self.ingestion_endpoint = ingestion_endpoint + self.max_attempts = max_attempts if max_attempts is not None else self.DEFAULT_MAX_ATTEMPTS + self.base_delay_sec = base_delay_sec if base_delay_sec is not None else self.DEFAULT_BASE_DELAY_SEC + self.timeout_sec = timeout_sec if timeout_sec is not None else self.DEFAULT_TIMEOUT_SEC self.formatted_tags = ','.join([f'{key}={value}' for key, value in tags.items()]) self.headers = { 'Content-Encoding': 'gzip', @@ -75,24 +107,54 @@ def __init__(self, api_key, ingestion_endpoint, dataset, collection, client_type if self.client_type is not None: self.headers.update({'x-bronto-client': self.client_type}) + def _calculate_delay(self, attempt: int) -> float: + """Calculate exponential backoff delay with jitter""" + base_delay = self.base_delay_sec * (2 ** (attempt - 1)) + jitter = random.uniform(0, base_delay * 0.1) + return base_delay + jitter + def _send_batch(self, compressed_batch): request = urllib.request.Request(self.ingestion_endpoint, data=compressed_batch, headers=self.headers) - attempt = 0 - max_attempts = 5 - with urllib.request.urlopen(request) as resp: - if resp.status != 200 and attempt < max_attempts: - attempt += 1 - delay_sec = attempt * 10 - logger.warning('Data sending failed. attempt=%s, max_attempts=%s, status=%s, reason=%s', - attempt, max_attempts, resp.status, resp.reason) + last_status = None + + for attempt in range(1, self.max_attempts + 1): + try: + with urllib.request.urlopen(request, timeout=self.timeout_sec) as resp: + last_status = resp.status + if resp.status == 200: + logger.info('data sent successfully. collection=%s, dataset=%s, tags=%s', + self.collection, self.dataset, self.formatted_tags) + return + elif resp.status in self.NON_RETRYABLE_STATUS_CODES: + logger.error('Non-retryable error encountered. status=%s, reason=%s', + resp.status, resp.reason) + raise BrontoClientNonRetryableError(resp.status, resp.reason) + else: + logger.warning('Data sending failed. attempt=%s, max_attempts=%s, status=%s, reason=%s', + attempt, self.max_attempts, resp.status, resp.reason) + except urllib.error.HTTPError as e: + last_status = e.code + if e.code in self.NON_RETRYABLE_STATUS_CODES: + logger.error('Non-retryable HTTP error. status=%s, reason=%s', e.code, e.reason) + raise BrontoClientNonRetryableError(e.code, str(e.reason)) + logger.warning('HTTP error occurred. attempt=%s, max_attempts=%s, status=%s, reason=%s', + attempt, self.max_attempts, e.code, e.reason) + except urllib.error.URLError as e: + logger.warning('URL error occurred. attempt=%s, max_attempts=%s, reason=%s', + attempt, self.max_attempts, e.reason) + except Exception as e: + logger.warning('Unexpected error occurred. attempt=%s, max_attempts=%s, error=%s', + attempt, self.max_attempts, str(e)) + + # If not the last attempt, wait before retrying + if attempt < self.max_attempts: + delay_sec = self._calculate_delay(attempt) + logger.info('Retrying in %.2f seconds...', delay_sec) time.sleep(delay_sec) - self._send_batch(compressed_batch) - elif resp.status == 200: - logger.info('data sent successfully. collection=%s, dataset=%s, tags=%s', self.collection, - self.dataset, self.formatted_tags) - else: - logger.error('max attempts reached. attempt=%s, max_attempts=%s', attempt, max_attempts) - raise Exception('BrontoClientMaxAttemptReached') + + # All attempts exhausted + logger.error('Max attempts reached. attempts=%s', self.max_attempts) + raise BrontoClientMaxAttemptsError(self.max_attempts, last_status) def send_data(self, batch, attributes=None): data = batch.get_formatted_data({} if attributes is None else attributes) diff --git a/tests/test_bronto_client.py b/tests/test_bronto_client.py new file mode 100644 index 0000000..2db52d8 --- /dev/null +++ b/tests/test_bronto_client.py @@ -0,0 +1,319 @@ +import unittest +from unittest.mock import Mock, patch, MagicMock +import urllib.error +import gzip + +from clients import ( + BrontoClient, + BrontoClientError, + BrontoClientNonRetryableError, + BrontoClientMaxAttemptsError, + Batch +) + + +class TestBrontoClient(unittest.TestCase): + + def setUp(self): + """Set up test fixtures""" + self.api_key = "test-api-key" + self.ingestion_endpoint = "https://test.bronto.com/ingest" + self.dataset = "test-dataset" + self.collection = "test-collection" + self.client_type = "test-client" + self.tags = {"env": "test", "version": "1.0"} + + self.client = BrontoClient( + api_key=self.api_key, + ingestion_endpoint=self.ingestion_endpoint, + dataset=self.dataset, + collection=self.collection, + client_type=self.client_type, + tags=self.tags + ) + + def test_client_initialization(self): + """Test that client initializes with correct default values""" + self.assertEqual(self.client.api_key, self.api_key) + self.assertEqual(self.client.ingestion_endpoint, self.ingestion_endpoint) + self.assertEqual(self.client.dataset, self.dataset) + self.assertEqual(self.client.collection, self.collection) + self.assertEqual(self.client.client_type, self.client_type) + self.assertEqual(self.client.max_attempts, BrontoClient.DEFAULT_MAX_ATTEMPTS) + self.assertEqual(self.client.base_delay_sec, BrontoClient.DEFAULT_BASE_DELAY_SEC) + self.assertEqual(self.client.timeout_sec, BrontoClient.DEFAULT_TIMEOUT_SEC) + + def test_client_initialization_with_custom_params(self): + """Test that client can be initialized with custom retry parameters""" + custom_client = BrontoClient( + api_key=self.api_key, + ingestion_endpoint=self.ingestion_endpoint, + dataset=self.dataset, + collection=self.collection, + client_type=self.client_type, + tags=self.tags, + max_attempts=3, + base_delay_sec=5, + timeout_sec=60 + ) + self.assertEqual(custom_client.max_attempts, 3) + self.assertEqual(custom_client.base_delay_sec, 5) + self.assertEqual(custom_client.timeout_sec, 60) + + def test_headers_are_set_correctly(self): + """Test that HTTP headers are constructed correctly""" + self.assertEqual(self.client.headers['x-bronto-api-key'], self.api_key) + self.assertEqual(self.client.headers['x-bronto-service-name'], self.dataset) + self.assertEqual(self.client.headers['x-bronto-service-namespace'], self.collection) + self.assertEqual(self.client.headers['x-bronto-client'], self.client_type) + self.assertIn('env=test', self.client.headers['x-bronto-tags']) + self.assertIn('version=1.0', self.client.headers['x-bronto-tags']) + + def test_calculate_delay_exponential_backoff(self): + """Test that delay calculation uses exponential backoff""" + # First attempt: base_delay * 2^0 = 2 * 1 = 2 + delay1 = self.client._calculate_delay(1) + self.assertGreaterEqual(delay1, 2.0) + self.assertLess(delay1, 3.0) # 10% jitter + + # Second attempt: base_delay * 2^1 = 2 * 2 = 4 + delay2 = self.client._calculate_delay(2) + self.assertGreaterEqual(delay2, 4.0) + self.assertLess(delay2, 5.0) # 10% jitter + + # Third attempt: base_delay * 2^2 = 2 * 4 = 8 + delay3 = self.client._calculate_delay(3) + self.assertGreaterEqual(delay3, 8.0) + self.assertLess(delay3, 9.0) # 10% jitter + + @patch('clients.urllib.request.urlopen') + def test_send_batch_success(self, mock_urlopen): + """Test successful batch sending""" + mock_response = MagicMock() + mock_response.status = 200 + mock_response.reason = 'OK' + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + batch = Batch(1000) + batch.add("test log line") + + # Should not raise any exception + self.client.send_data(batch) + + # Verify urlopen was called once + self.assertEqual(mock_urlopen.call_count, 1) + + @patch('clients.urllib.request.urlopen') + def test_send_batch_non_retryable_403(self, mock_urlopen): + """Test that 403 errors are not retried""" + mock_urlopen.side_effect = urllib.error.HTTPError( + url=self.ingestion_endpoint, + code=403, + msg='Forbidden', + hdrs={}, + fp=None + ) + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientNonRetryableError) as context: + self.client.send_data(batch) + + self.assertEqual(context.exception.status_code, 403) + # Should only try once, not retry + self.assertEqual(mock_urlopen.call_count, 1) + + @patch('clients.time.sleep') + @patch('clients.urllib.request.urlopen') + def test_send_batch_non_retryable_401(self, mock_urlopen, mock_sleep): + """Test that 401 errors are not retried""" + mock_urlopen.side_effect = urllib.error.HTTPError( + url=self.ingestion_endpoint, + code=401, + msg='Unauthorized', + hdrs={}, + fp=None + ) + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientNonRetryableError) as context: + self.client.send_data(batch) + + self.assertEqual(context.exception.status_code, 401) + self.assertEqual(mock_urlopen.call_count, 1) + + @patch('clients.urllib.request.urlopen') + def test_send_batch_http_error_non_retryable(self, mock_urlopen): + """Test that HTTPError with non-retryable status codes are not retried""" + mock_urlopen.side_effect = urllib.error.HTTPError( + url=self.ingestion_endpoint, + code=404, + msg='Not Found', + hdrs={}, + fp=None + ) + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientNonRetryableError) as context: + self.client.send_data(batch) + + self.assertEqual(context.exception.status_code, 404) + self.assertEqual(mock_urlopen.call_count, 1) + + @patch('clients.time.sleep') + @patch('clients.urllib.request.urlopen') + def test_send_batch_retries_on_500(self, mock_urlopen, mock_sleep): + """Test that 500 errors are retried""" + mock_response = MagicMock() + mock_response.status = 500 + mock_response.reason = 'Internal Server Error' + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientMaxAttemptsError) as context: + self.client.send_data(batch) + + # Should try max_attempts times (default 5) + self.assertEqual(mock_urlopen.call_count, 5) + # Should sleep 4 times (between attempts) + self.assertEqual(mock_sleep.call_count, 4) + self.assertEqual(context.exception.attempts, 5) + self.assertEqual(context.exception.last_status, 500) + + @patch('clients.time.sleep') + @patch('clients.urllib.request.urlopen') + def test_send_batch_retries_then_succeeds(self, mock_urlopen, mock_sleep): + """Test that retries eventually succeed""" + # First 2 calls fail with 503, third succeeds + mock_response_fail = MagicMock() + mock_response_fail.status = 503 + mock_response_fail.reason = 'Service Unavailable' + mock_response_fail.__enter__ = Mock(return_value=mock_response_fail) + mock_response_fail.__exit__ = Mock(return_value=False) + + mock_response_success = MagicMock() + mock_response_success.status = 200 + mock_response_success.reason = 'OK' + mock_response_success.__enter__ = Mock(return_value=mock_response_success) + mock_response_success.__exit__ = Mock(return_value=False) + + mock_urlopen.side_effect = [mock_response_fail, mock_response_fail, mock_response_success] + + batch = Batch(1000) + batch.add("test log line") + + # Should succeed without raising exception + self.client.send_data(batch) + + # Should have tried 3 times + self.assertEqual(mock_urlopen.call_count, 3) + # Should have slept 2 times (between attempts) + self.assertEqual(mock_sleep.call_count, 2) + + @patch('clients.time.sleep') + @patch('clients.urllib.request.urlopen') + def test_send_batch_url_error_retries(self, mock_urlopen, mock_sleep): + """Test that URLError causes retries""" + mock_urlopen.side_effect = urllib.error.URLError('Connection refused') + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientMaxAttemptsError): + self.client.send_data(batch) + + # Should try max_attempts times + self.assertEqual(mock_urlopen.call_count, 5) + self.assertEqual(mock_sleep.call_count, 4) + + @patch('clients.time.sleep') + @patch('clients.urllib.request.urlopen') + def test_send_batch_timeout_retries(self, mock_urlopen, mock_sleep): + """Test that timeout errors cause retries""" + import socket + mock_urlopen.side_effect = socket.timeout('Connection timed out') + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientMaxAttemptsError): + self.client.send_data(batch) + + # Should try max_attempts times + self.assertEqual(mock_urlopen.call_count, 5) + self.assertEqual(mock_sleep.call_count, 4) + + @patch('clients.urllib.request.urlopen') + def test_send_batch_with_custom_max_attempts(self, mock_urlopen): + """Test that custom max_attempts is respected""" + custom_client = BrontoClient( + api_key=self.api_key, + ingestion_endpoint=self.ingestion_endpoint, + dataset=self.dataset, + collection=self.collection, + client_type=self.client_type, + tags=self.tags, + max_attempts=2 + ) + + mock_response = MagicMock() + mock_response.status = 500 + mock_response.reason = 'Internal Server Error' + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + batch = Batch(1000) + batch.add("test log line") + + with self.assertRaises(BrontoClientMaxAttemptsError) as context: + custom_client.send_data(batch) + + # Should only try 2 times + self.assertEqual(mock_urlopen.call_count, 2) + self.assertEqual(context.exception.attempts, 2) + + @patch('clients.urllib.request.urlopen') + def test_send_data_compresses_batch(self, mock_urlopen): + """Test that send_data properly compresses the batch data""" + mock_response = MagicMock() + mock_response.status = 200 + mock_response.reason = 'OK' + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + batch = Batch(1000) + test_data = "test log line" + batch.add(test_data) + + self.client.send_data(batch) + + # Verify the request was made + call_args = mock_urlopen.call_args + request = call_args[0][0] + + # Verify data is compressed + compressed_data = request.data + decompressed = gzip.decompress(compressed_data).decode() + self.assertIn(test_data, decompressed) + + def test_non_retryable_status_codes(self): + """Test that NON_RETRYABLE_STATUS_CODES contains expected codes""" + expected_codes = {400, 401, 403, 404, 405, 410, 422} + self.assertEqual(BrontoClient.NON_RETRYABLE_STATUS_CODES, expected_codes) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file