Skip to content

Commit 3a941ff

Browse files
SimOnPanwSimOnPanw
authored
Merge pull request #183 from mgappa/refactor/173_use_proper_retry
Refactor/173 use proper retry
2 parents ee0ca4b + 7bc8578 commit 3a941ff

File tree

9 files changed

+1955
-134
lines changed

9 files changed

+1955
-134
lines changed

‎prismacloud/api/cwpp/cwpp.py‎

Lines changed: 80 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import time
55

66
import requests
7+
from requests.adapters import HTTPAdapter, Retry
8+
79

810
class PrismaCloudAPICWPPMixin():
911
""" Requests and Output """
@@ -16,7 +18,8 @@ def login_compute(self):
1618
# Login via CWP.
1719
self.login('https://%s/api/v1/authenticate' % self.api_compute)
1820
else:
19-
self.error_and_exit(418, "Specify a Prisma Cloud URL or Prisma Cloud Compute URL")
21+
self.error_and_exit(
22+
418, "Specify a Prisma Cloud URL or Prisma Cloud Compute URL")
2023
self.debug_print('New API Token: %s' % self.token)
2124

2225
def extend_login_compute(self):
@@ -43,78 +46,91 @@ def execute_compute(self, action, endpoint, query_params=None, body_params=None,
4346
limit = 50
4447
more = False
4548
results = []
46-
while offset == 0 or more is True:
47-
if int(time.time() - self.token_timer) > self.token_limit:
48-
self.extend_login_compute()
49-
if paginated:
50-
url = 'https://%s/%s?limit=%s&offset=%s' % (self.api_compute, endpoint, limit, offset)
51-
else:
52-
url = 'https://%s/%s' % (self.api_compute, endpoint)
53-
if self.token:
54-
if self.api:
55-
# Authenticate via CSPM
56-
request_headers['x-redlock-auth'] = self.token
49+
50+
retries = Retry(total=self.retry_number,
51+
status_forcelist=self.retry_status_codes, raise_on_status=False)
52+
53+
with requests.Session() as session:
54+
session.mount('https://%s/%s' % (self.api_compute,
55+
endpoint), adapter=HTTPAdapter(max_retries=retries))
56+
57+
while offset == 0 or more is True:
58+
if int(time.time() - self.token_timer) > self.token_limit:
59+
self.extend_login_compute()
60+
if paginated:
61+
url = 'https://%s/%s?limit=%s&offset=%s' % (
62+
self.api_compute, endpoint, limit, offset)
5763
else:
58-
# Authenticate via CWP
59-
request_headers['Authorization'] = "Bearer %s" % self.token
60-
self.debug_print('API URL: %s' % url)
61-
self.debug_print('API Request Headers: (%s)' % request_headers)
62-
self.debug_print('API Query Params: %s' % query_params)
63-
self.debug_print('API Body Params: %s' % body_params_json)
64-
# Add User-Agent to the headers
65-
request_headers['User-Agent'] = self.user_agent
66-
api_response = requests.request(action, url, headers=request_headers, params=query_params, data=body_params_json, verify=self.verify, timeout=self.timeout)
67-
self.debug_print('API Response Status Code: (%s)' % api_response.status_code)
68-
self.debug_print('API Response Headers: (%s)' % api_response.headers)
69-
if api_response.status_code in self.retry_status_codes:
70-
for exponential_wait in self.retry_waits:
71-
time.sleep(exponential_wait)
72-
api_response = requests.request(action, url, headers=request_headers, params=query_params, data=body_params_json, verify=self.verify, timeout=self.timeout)
73-
if api_response.ok:
74-
break # retry loop
75-
if api_response.ok:
76-
if not api_response.content:
77-
return None
78-
if api_response.headers.get('Content-Type') == 'application/x-gzip':
79-
return api_response.content
80-
if api_response.headers.get('Content-Type') == 'text/csv':
81-
return api_response.content.decode('utf-8')
82-
try:
83-
result = json.loads(api_response.content)
84-
#if result is None:
85-
# self.logger.error('JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
86-
# if force:
87-
# return results # or continue
88-
# self.error_and_exit(api_response.status_code, 'JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
89-
except ValueError:
90-
self.logger.error('JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
91-
if force:
92-
return results # or continue
93-
self.error_and_exit(api_response.status_code, 'JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
94-
if 'Total-Count' in api_response.headers:
95-
self.debug_print('Retrieving Next Page of Results: Offset/Total Count: %s/%s' % (offset, api_response.headers['Total-Count']))
96-
total_count = int(api_response.headers['Total-Count'])
97-
if total_count > 0:
98-
results.extend(result)
99-
offset += limit
100-
more = bool(offset < total_count)
64+
url = 'https://%s/%s' % (self.api_compute, endpoint)
65+
if self.token:
66+
if self.api:
67+
# Authenticate via CSPM
68+
request_headers['x-redlock-auth'] = self.token
69+
else:
70+
# Authenticate via CWP
71+
request_headers['Authorization'] = "Bearer %s" % self.token
72+
self.debug_print('API URL: %s' % url)
73+
self.debug_print('API Request Headers: (%s)' % request_headers)
74+
self.debug_print('API Query Params: %s' % query_params)
75+
self.debug_print('API Body Params: %s' % body_params_json)
76+
# Add User-Agent to the headers
77+
request_headers['User-Agent'] = self.user_agent
78+
api_response = session.request(action, url, headers=request_headers, params=query_params,
79+
data=body_params_json, verify=self.verify, timeout=self.timeout)
80+
self.debug_print('API Response Status Code: (%s)' %
81+
api_response.status_code)
82+
self.debug_print('API Response Headers: (%s)' %
83+
api_response.headers)
84+
if api_response.ok:
85+
if not api_response.content:
86+
return None
87+
if api_response.headers.get('Content-Type') == 'application/x-gzip':
88+
return api_response.content
89+
if api_response.headers.get('Content-Type') == 'text/csv':
90+
return api_response.content.decode('utf-8')
91+
try:
92+
result = json.loads(api_response.content)
93+
# if result is None:
94+
# self.logger.error('JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
95+
# if force:
96+
# return results # or continue
97+
# self.error_and_exit(api_response.status_code, 'JSON returned None, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (url, query_params, body_params, api_response.content))
98+
except ValueError:
99+
self.logger.error('JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (
100+
url, query_params, body_params, api_response.content))
101+
if force:
102+
return results # or continue
103+
self.error_and_exit(api_response.status_code, 'JSON raised ValueError, API: (%s) with query params: (%s) and body params: (%s) parsing response: (%s)' % (
104+
url, query_params, body_params, api_response.content))
105+
if 'Total-Count' in api_response.headers:
106+
self.debug_print('Retrieving Next Page of Results: Offset/Total Count: %s/%s' % (
107+
offset, api_response.headers['Total-Count']))
108+
total_count = int(api_response.headers['Total-Count'])
109+
if total_count > 0:
110+
results.extend(result)
111+
offset += limit
112+
more = bool(offset < total_count)
113+
else:
114+
return result
101115
else:
102-
return result
103-
else:
104-
self.logger.error('API: (%s) responded with a status of: (%s), with query: (%s) and body params: (%s)' % (url, api_response.status_code, query_params, body_params))
105-
if force:
106-
return results
107-
self.error_and_exit(api_response.status_code, 'API: (%s) with query params: (%s) and body params: (%s) responded with an error and this response:\n%s' % (url, query_params, body_params, api_response.text))
108-
return results
116+
self.logger.error('API: (%s) responded with a status of: (%s), with query: (%s) and body params: (%s)' % (
117+
url, api_response.status_code, query_params, body_params))
118+
if force:
119+
return results
120+
self.error_and_exit(api_response.status_code, 'API: (%s) with query params: (%s) and body params: (%s) responded with an error and this response:\n%s' % (
121+
url, query_params, body_params, api_response.text))
122+
return results
109123

110124
# The Compute API setting is optional.
111125

112126
def validate_api_compute(self):
113127
if not self.api_compute:
114-
self.error_and_exit(500, 'Please specify a Prisma Cloud Compute API URL.')
128+
self.error_and_exit(
129+
500, 'Please specify a Prisma Cloud Compute API URL.')
115130

116131
# Exit handler (Error).
117132

118133
@classmethod
119134
def error_and_exit(cls, error_code, error_message='', system_message=''):
120-
raise SystemExit('\n\nStatus Code: %s\n%s\n%s\n' % (error_code, error_message, system_message))
135+
raise SystemExit('\n\nStatus Code: %s\n%s\n%s\n' %
136+
(error_code, error_message, system_message))

‎prismacloud/api/pc_lib_api.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def __init__(self):
4343
self.token_limit = 590 # aka 9 minutes
4444
self.retry_status_codes = [425, 429, 500, 502, 503, 504]
4545
self.retry_waits = [1, 2, 4, 8, 16, 32]
46+
self.retry_number = 6
4647
self.max_workers = 8
4748
#
4849
self.error_log = 'error.log'

‎pyproject.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
11
[build-system]
22
requires = ["setuptools"]
3-
build-backend = "setuptools.build_meta"
3+
build-backend = "setuptools.build_meta"
4+
5+
[project.optional-dependencies]
6+
test = ["coverage==7.6.10", "responses==0.25.3"]

‎requirements_test.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
coverage==7.6.10
2+
responses==0.25.3

‎setup.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,5 +34,8 @@
3434
'requests',
3535
'update_checker'
3636
],
37+
extras_require={
38+
'test': ['coverage==7.6.10', 'responses==0.25.3']
39+
},
3740
python_requires='>=3.6'
3841
)

0 commit comments

Comments
 (0)