TDL-19841 request timeout and backoff implementation

tap_braintree/__init__.py

@@ -2,10 +2,14 @@
 import sys
 import json
 import braintree
+from requests import request
 import singer
 from singer import utils
 from tap_braintree.discover import discover
 from tap_braintree.sync import sync as _sync
+import datetime
+
+REQUEST_TIME_OUT = 300
 
 REQUIRED_CONFIG_KEYS = [
     "merchant_id", 
@@ -40,12 +44,23 @@ def main():
 
     environment = getattr(braintree.Environment, config.pop("environment", "Production"))
 
+    try:
+        # Take value of request_timeout if provided in config else take default value
+        request_timeout = float(config.get("request_timeout", REQUEST_TIME_OUT))
+
+        if request_timeout == 0:
+            # Raise error when request_timeout is given as 0 in config
+            raise ValueError()
+    except ValueError:
+        raise ValueError('Please provide a value greater than 0 for the request_timeout parameter in config')
+
     gateway = braintree.BraintreeGateway(
         braintree.Configuration(
             environment,
             merchant_id = config['merchant_id'],
             public_key= config["public_key"],
-            private_key=config["private_key"]
+            private_key=config["private_key"],
+            timeout = request_timeout
         )
     )
     try:

tap_braintree/streams.py

@@ -1,10 +1,17 @@
 import pytz
 import singer
 import braintree
+import backoff
 from singer import utils
 from datetime import datetime, timedelta
 from .transform import transform_row
 
+from braintree.exceptions.too_many_requests_error import TooManyRequestsError
+from braintree.exceptions.server_error import ServerError
+from braintree.exceptions.service_unavailable_error import ServiceUnavailableError
+from braintree.exceptions.gateway_timeout_error import GatewayTimeoutError
+
+
 TRAILING_DAYS = timedelta(days=30)
 DEFAULT_TIMESTAMP = "1970-01-01T00:00:00Z"
 LOGGER = singer.get_logger()
@@ -86,6 +93,14 @@ class SyncWithoutWindow(Stream):
     replication_keys = "updated_at"
     replication_method = "INCREMENTAL"
 
+    # Backoff the request for 5 times when ConnectionError, TooManyRequestsError (status code = 429),
+    # ServerError(status code = 500) , ServiceUnavailableError (status code = 503) ,
+    # or GatewayTimeoutError (status code = 504) occurs
+    @backoff.on_exception(
+        backoff.expo,
+        (ConnectionError, TooManyRequestsError, ServerError, ServiceUnavailableError, GatewayTimeoutError),
+        max_tries=5,
+        factor=2)
     def sync(self, gateway, config, schema, state, selected_streams):
         """
         Sync function for incremental stream without window logic
@@ -133,6 +148,18 @@ class SyncWithWindow(Stream):
     replication_keys = "created_at"
     replication_method = "INCREMENTAL"
 
+    # Backoff the request for 5 times when ConnectionError, TooManyRequestsError (status code = 429),
+    # ServerError(status code = 500) , ServiceUnavailableError (status code = 503) ,
+    # or GatewayTimeoutError (status code = 504) occurs
+    @backoff.on_exception(
+        backoff.expo,
+        (ConnectionError, TooManyRequestsError, ServerError, ServiceUnavailableError, GatewayTimeoutError),
+        max_tries=5,
+        factor=2)
+    def GetRecords(self, gateway, start, end):
+        """Return records between given date window"""
+        return self.sdk_call(gateway, start, end)
+
     def sync(self, gateway, config, schema, state, selected_streams):
         """
         Sync function for incremental stream with window logic
@@ -163,7 +190,7 @@ def sync(self, gateway, config, schema, state, selected_streams):
         for start, end in self.daterange(period_start, period_end):
             end = min(end, period_end)
 
-            data = self.sdk_call(gateway, start, end)
+            data = self.GetRecords(gateway, start, end)
             time_extracted = utils.now()
 
             row_written_count = 0
@@ -247,6 +274,14 @@ class FullTableSync(Stream):
     sdk_call = None
     key_properties = ["id"]
 
+    # Backoff the request for 5 times when ConnectionError, TooManyRequestsError (status code = 429),
+    # ServerError(status code = 500) , ServiceUnavailableError (status code = 503) ,
+    # or GatewayTimeoutError (status code = 504) occurs
+    @backoff.on_exception(
+        backoff.expo,
+        (ConnectionError, TooManyRequestsError, ServerError, ServiceUnavailableError, GatewayTimeoutError),
+        max_tries=5,
+        factor=2)
     def sync(self, gateway, config, schema, state, selected_streams):
         """
         Sync function for full_table stream

tests/unittests/test_backoff.py

@@ -0,0 +1,57 @@
+import unittest
+from unittest import mock
+from parameterized import parameterized
+from braintree.exceptions.too_many_requests_error import TooManyRequestsError
+from braintree.exceptions.server_error import ServerError
+from braintree.exceptions.service_unavailable_error import ServiceUnavailableError
+from braintree.exceptions.gateway_timeout_error import GatewayTimeoutError
+from tap_braintree.streams import AddOn, Transaction
+
+
+class TestBackoff(unittest.TestCase):
+    '''
+    Test that backoff logic works properly
+    '''
+
+    gateway = "test"
+    config = {"start_date":"2022-08-01T00:00:00Z"}
+    schema = {}
+    state = {}
+
+    @parameterized.expand([
+        ['connection_error', ConnectionError, 5],
+        ['too_many_requests_error', TooManyRequestsError, 5],
+        ['server_error', ServerError, 5],
+        ['service_unavailable_error', ServiceUnavailableError, 5],
+        ['gateway_timeout_error', GatewayTimeoutError, 5]
+        ])
+    @mock.patch("tap_braintree.streams.AddOn.sdk_call")
+    @mock.patch("time.sleep")
+    def test_backoff_for_sync_without_window(self, name, test_exception, expected_count, mocked_time, mocked_sdk_call):
+        '''Test function to verify working of backoff for sync of incremental streams without window'''
+
+        stream_obj = AddOn()
+        mocked_sdk_call.side_effect = test_exception('exception')
+
+        with self.assertRaises(test_exception):
+            stream_obj.sync(self.gateway, self.config, self.schema, self.state, ["add_ons"])
+        self.assertEqual(mocked_sdk_call.call_count, expected_count)
+
+    @parameterized.expand([
+        ['connection_error', ConnectionError, 5],
+        ['too_many_requests_error', TooManyRequestsError, 5],
+        ['server_error', ServerError, 5],
+        ['service_unavailable_error', ServiceUnavailableError, 5],
+        ['gateway_timeout_error', GatewayTimeoutError, 5]
+        ])
+    @mock.patch("tap_braintree.streams.Transaction.sdk_call")
+    @mock.patch("time.sleep")
+    def test_backoff_for_sync_with_window(self, name, test_exception, expected_count, mocked_time, mocked_sdk_call):
+        '''Test function to verify working of backoff for sync of incremental streams with window'''
+
+        stream_obj = Transaction()
+        mocked_sdk_call.side_effect = test_exception('exception')
+
+        with self.assertRaises(test_exception):
+            stream_obj.GetRecords(self.gateway, self.config["start_date"], self.config["start_date"])
+        self.assertEqual(mocked_sdk_call.call_count, expected_count)