def _check_status(self, resp, expect=200):
if resp.status_code == 400:
r = json.loads(resp.text)
raise InvalidSearch(r['message'])
if resp.status_code == 401:
raise AuthError('unauthorized')
if resp.status_code == 404:
raise NotFound('not found')
if resp.status_code == 408:
raise TimeoutError('timeout')
if resp.status_code == 422:
msg = json.loads(resp.text)
raise SubmissionFailed(msg['message'])
if resp.status_code == 429:
raise CIFBusy('RateLimit exceeded')
if resp.status_code in [500, 501, 502, 503, 504]:
raise CIFBusy('system seems busy..')
if resp.status_code != expect:
msg = 'unknown: %s' % resp.content
raise RuntimeError(msg)
def _recv(self, decode=True):
mtype, data = Msg().recv(self.socket)
if not decode:
return data
data = json.loads(data)
if data.get('message') == 'unauthorized':
raise AuthError()
if data.get('message') == 'busy':
raise CIFBusy()
if data.get('message') == 'invalid search':
raise InvalidSearch()
if data.get('status') != 'success':
raise RuntimeError(data.get('message'))
try:
data['data'] = zlib.decompress(data['data'])
except (zlib.error, TypeError):
pass
return data.get('data')
def _post(self, uri, data, expect=201):
if not uri.startswith('http'):
uri = "%s/%s/" % (self.remote, uri)
data = json.dumps(data)
if self.nowait:
uri = '{}?nowait=1'.format(uri)
if isinstance(data, str):
data = data.encode('utf-8')
headers = {'Content-Type': 'application/json'}
logger.debug('submitting')
resp = self.session.post(uri,
data=data,
verify=self.verify_ssl,
headers=headers,
timeout=self.timeout)
logger.debug(resp.content)
logger.debug(resp.status_code)
n = RETRIES
try:
self._check_status(resp, expect=expect)
n = 0
except Exception as e:
if resp.status_code == 429 or resp.status_code in [
500, 501, 502, 503, 504
]:
logger.error(e)
else:
raise e
while n != 0:
logger.info('setting random retry interval to spread out the load')
logger.info('retrying in %.00fs' % RETRIES_DELAY)
sleep(RETRIES_DELAY)
resp = self.session.post(uri,
data=data,
verify=self.verify_ssl,
headers=headers,
timeout=self.timeout)
if resp.status_code in [200, 201]:
break
if n == 0:
raise CIFBusy('system seems busy.. try again later')
return json.loads(resp.content.decode('utf-8'))
def _recv(self, decode=True, close=True):
mtype, data = Msg().recv(self.socket)
if close:
self.socket.close()
if not decode:
return data
data = json.loads(data)
if data.get('message') == 'unauthorized':
raise AuthError()
if data.get('message') == 'busy':
raise CIFBusy()
if data.get('message') == 'invalid search':
raise InvalidSearch()
if data.get('status') != 'success':
raise RuntimeError(data.get('message'))
if data.get('data') is None:
raise RuntimeError('invalid response')
if isinstance(data.get('data'), bool):
return data['data']
# is this a straight up elasticsearch string?
if data['data'] == '{}':
return []
if isinstance(data['data'], basestring) and data['data'].startswith('{"hits":{"hits":[{"_source":'):
data['data'] = json.loads(data['data'])
data['data'] = [r['_source'] for r in data['data']['hits']['hits']]
try:
data['data'] = zlib.decompress(data['data'])
except (zlib.error, TypeError):
pass
return data.get('data')
def __get(self, uri, params):
resp = self.session.get(uri,
params=params,
verify=self.verify_ssl,
timeout=self.timeout)
logger.debug(resp.text)
n = RETRIES
try:
self._check_status(resp, expect=200)
n = 0
return resp
except Exception as e:
if resp.status_code == 429 or resp.status_code in [
500, 501, 502, 503, 504
]:
logger.error(e)
else:
raise e
for nn in (1, n):
logger.warning(
'setting random retry interval to spread out the load')
logger.warning('retrying in %.00fs' % RETRIES_DELAY)
sleep(RETRIES_DELAY)
resp = self.session.get(uri,
params=params,
verify=self.verify_ssl,
timeout=self.timeout)
if resp.status_code == 200:
break
if nn == n:
raise CIFBusy('system seems busy.. try again later')
return resp
def _get(self, uri, params={}, retry=True):
if not uri.startswith('http'):
uri = self.remote + uri
resp = self.session.get(uri,
params=params,
verify=self.verify_ssl,
timeout=self.timeout)
n = RETRIES
try:
self._check_status(resp, expect=200)
n = 0
except Exception as e:
if resp.status_code == 429 or resp.status_code in [
500, 501, 502, 503, 504
]:
logger.error(e)
else:
raise e
while n != 0:
logger.warning(
'setting random retry interval to spread out the load')
logger.warning('retrying in %.00fs' % RETRIES_DELAY)
sleep(RETRIES_DELAY)
resp = self.session.get(uri,
params=params,
verify=self.verify_ssl,
timeout=self.timeout)
if resp.status_code == 200:
break
if n == 0:
raise CIFBusy('system seems busy.. try again later')
data = resp.content
s = (int(resp.headers['Content-Length']) / 1024 / 1024)
logger.info('processing %.2f megs' % s)
msgs = json.loads(data.decode('utf-8'))
if msgs.get('data') and msgs['data'] == '{}':
msgs['data'] = []
if msgs.get('data') and isinstance(
msgs['data'], basestring) and msgs['data'].startswith(
'{"hits":{"hits":[{"_source":'):
msgs['data'] = json.loads(msgs['data'])
msgs['data'] = [r['_source'] for r in msgs['data']['hits']['hits']]
if not msgs.get('status') and not msgs.get('message') == 'success':
raise RuntimeError(msgs)
if msgs.get('status') and msgs['status'] == 'failed':
raise InvalidSearch(msgs['message'])
if isinstance(msgs.get('data'), list):
for m in msgs['data']:
if m.get('message'):
try:
m['message'] = b64decode(m['message'])
except Exception as e:
pass
return msgs