Skip to content

Commit f807dfe

Browse files
samwelkandasamwelkanda
committed
Add check for non-buffered or streaming requests
1 parent 479ae16 commit f807dfe

3 files changed

Lines changed: 51 additions & 19 deletions

File tree

‎msgraphcore/middleware/options/middleware_control.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@ def get(self, middleware_option_name):
1515

1616
def get_middleware_options(self, func):
1717
def wrapper(*args, **kwargs):
18-
# Get middleware options from **kwargs
1918
self._reset_middleware_options()
19+
# Get middleware options from **kwargs
2020
scopes = kwargs.pop('scopes', None)
2121
if scopes:
2222
# Set middleware options, for use by middleware in the middleware pipeline

‎msgraphcore/middleware/retry_middleware.py‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ class RetryMiddleware(BaseMiddleware):
2323

2424
def __init__(self, retry_configs={}):
2525
super().__init__()
26-
self.total_retries: int = retry_configs.pop('retry_total', self.DEFAULT_TOTAL_RETRIES)
26+
self.total_retries: int = min(
27+
retry_configs.pop('retry_total', self.DEFAULT_TOTAL_RETRIES), self.MAX_TOTAL_RETRIES
28+
)
2729
self.backoff_factor: float = retry_configs.pop(
2830
'retry_backoff_factor', self.DEFAULT_BACKOFF_FACTOR
2931
)
@@ -33,7 +35,9 @@ def __init__(self, retry_configs={}):
3335
status_codes: [int] = retry_configs.pop('retry_on_status_codes', [])
3436

3537
self._retry_on_status_codes = set(status_codes) | self._DEFAULT_RETRY_CODES
36-
self._allowed_methods = frozenset(['HEAD', 'GET', 'PUT', 'DELETE', 'OPTIONS', 'TRACE'])
38+
self._allowed_methods = frozenset(
39+
['HEAD', 'GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS']
40+
)
3741
self._respect_retry_after_header = True
3842

3943
@classmethod
@@ -53,7 +57,7 @@ def get_retry_options(self):
5357
if retry_config_options:
5458
return {
5559
'total':
56-
retry_config_options.retry_total
60+
min(retry_config_options.retry_total, self.MAX_TOTAL_RETRIES)
5761
if retry_config_options.retry_total is not None else self.total_retries,
5862
'backoff':
5963
retry_config_options.retry_backoff_factor
@@ -119,6 +123,8 @@ def should_retry(self, retry_options, response):
119123
"""
120124
if not self._is_method_retryable(retry_options, response.request):
121125
return False
126+
if not self._is_request_payload_buffered(response):
127+
return False
122128
return retry_options['total'] and response.status_code in retry_options['retry_codes']
123129

124130
def _is_method_retryable(self, retry_options, request):
@@ -130,6 +136,18 @@ def _is_method_retryable(self, retry_options, request):
130136
return False
131137
return True
132138

139+
def _is_request_payload_buffered(self, response):
140+
"""
141+
Checks if the request payload is buffered/rewindable.
142+
Payloads with forward only streams will return false and have the responses
143+
returned without any retry attempt.
144+
"""
145+
if response.request.method.upper() in frozenset(['HEAD', 'GET', 'DELETE', 'OPTIONS']):
146+
return True
147+
if response.request.headers['Content-Type'] == "application/octet-stream":
148+
return False
149+
return True
150+
133151
def increment_counter(self, retry_options):
134152
"""
135153
Increment the retry counters on every valid retry

‎tests/unit/test_retry_handler.py‎

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def test_no_config():
2020
assert retry_handler.backoff_max == retry_handler.MAXIMUM_BACKOFF
2121
assert retry_handler.backoff_factor == retry_handler.DEFAULT_BACKOFF_FACTOR
2222
assert retry_handler._allowed_methods == frozenset(
23-
['HEAD', 'GET', 'PUT', 'DELETE', 'OPTIONS', 'TRACE']
23+
['HEAD', 'GET', 'PUT', 'POST', 'PATCH', 'DELETE', 'OPTIONS']
2424
)
2525
assert retry_handler._respect_retry_after_header
2626
assert retry_handler._retry_on_status_codes == retry_handler._DEFAULT_RETRY_CODES
@@ -72,20 +72,6 @@ def test_method_retryable_with_valid_method():
7272
assert retry_handler._is_method_retryable(settings, response.request)
7373

7474

75-
@responses.activate
76-
def test_method_retryable_with_invalid_method():
77-
"""
78-
Test if method is retryable with a non retryable request method
79-
"""
80-
responses.add(responses.POST, BASE_URL, body="Test", status=503)
81-
response = requests.post(BASE_URL, data={"left": 1, "right": 3})
82-
83-
retry_handler = RetryMiddleware(retry_configs={})
84-
settings = retry_handler.get_retry_options()
85-
86-
assert not retry_handler._is_method_retryable(settings, response.request)
87-
88-
8975
@responses.activate
9076
def test_should_retry_valid():
9177
"""
@@ -114,6 +100,34 @@ def test_should_retry_invalid():
114100
assert not retry_handler.should_retry(settings, response)
115101

116102

103+
@responses.activate
104+
def test_is_request_payload_buffered_valid():
105+
"""
106+
Test for _is_request_payload_buffered helper method.
107+
Should return true request payload is buffered/rewindable.
108+
"""
109+
responses.add(responses.GET, BASE_URL, status=429)
110+
response = requests.get(BASE_URL)
111+
112+
retry_handler = RetryMiddleware(retry_configs={})
113+
114+
assert retry_handler._is_request_payload_buffered(response)
115+
116+
117+
@responses.activate
118+
def test_is_request_payload_buffered_invalid():
119+
"""
120+
Test for _is_request_payload_buffered helper method.
121+
Should return false if request payload is forward streamed.
122+
"""
123+
responses.add(responses.POST, BASE_URL, status=429)
124+
response = requests.post(BASE_URL, headers={'Content-Type': "application/octet-stream"})
125+
126+
retry_handler = RetryMiddleware(retry_configs={})
127+
128+
assert not retry_handler._is_request_payload_buffered(response)
129+
130+
117131
def test_retries_exhausted():
118132
"""
119133
Test that the retries exhausted method works correctly when total_retries is greater than zero

0 commit comments

Comments
 (0)