def get_task_result():
response = self.call_api("get", ("tasks", namespace, task_id))
content_disposition = response.headers.get("Content-Disposition")
if content_disposition and content_disposition.startswith(
"attachment"):
# we were sent a file to download
filename = None
filename_re = r'attachment\s*;\s*filename\s*=\s*"(.+)"\s*$'
m = re.search(filename_re, content_disposition)
if m:
filename = m.group(1)
return FileDownloadResponse(
response.content, response.headers.get("Content-Type"),
filename)
# otherwise we should have received a json
reply = response.json()
if reply["status"] != "ok":
raise encapsia_api.EncapsiaApiError(response.text)
rest_api_result = reply["result"]
task_status = rest_api_result["status"]
task_result = rest_api_result["result"]
if task_status == "finished":
return task_result
elif task_status == "failed":
raise encapsia_api.EncapsiaApiError(rest_api_result)
else:
return NoResultYet
def call_api(
self,
method,
path_segments,
data=None,
json=None,
return_json=False,
check_json_status=False,
extra_headers=None,
expected_codes=(200, 201),
params=None,
stream=False,
):
headers = {
"Accept": "application/json",
"Authorization": "Bearer {}".format(self.token),
"User-Agent": f"encapsia-api/{encapsia_api.__version__}",
}
if json:
headers["Content-type"] = "application/json"
if extra_headers:
headers.update(extra_headers)
if path_segments:
segments = [self.url, self.version]
if isinstance(path_segments, str):
segments.append(path_segments.lstrip("/"))
else:
segments.extend([s.lstrip("/") for s in path_segments])
else:
segments = [self.url]
url = "/".join(segments)
response = requests.request(
method,
url,
data=data,
json=json,
params=params,
headers=headers,
verify=True,
allow_redirects=False,
stream=stream,
)
if response.status_code not in expected_codes:
raise encapsia_api.EncapsiaApiError(
"{} {}\nFull response:\n{}".format(
response.status_code,
response.reason,
(response.content or "").strip(),
)
)
if not stream and return_json:
answer = response.json()
if check_json_status and answer["status"] != "ok":
raise encapsia_api.EncapsiaApiError(response.text)
return answer
else:
return response
def get_job_result():
with self.call_api("get", ("tasks", namespace, task_id)) as response:
reply = response.json()
if reply.get("status") != "ok":
raise encapsia_api.EncapsiaApiError(response.text)
rest_api_result = reply["result"]
task_status = rest_api_result["status"]
if task_status == "finished":
reply = self.get(("jobs", namespace, job_id))
joblog = reply["result"]["logs"][0]
assert joblog["status"] == "success"
result = joblog["output"]
if download:
filename = pathlib.Path(download)
with filename.open("wt") as f:
json.dump(result, f, indent=4)
return FileDownloadResponse(filename, "application/json")
else:
return result
elif task_status == "failed":
raise encapsia_api.EncapsiaApiFailedTaskError(
"Failed Task behind the Job. See Exception payload attribute.",
payload=rest_api_result,
)
else:
return NoResultYet
def download_blob_data(self, blob_id, fileobj=None):
"""Download blob data for given blob_id."""
extra_headers = {"Accept": "*/*"}
response = self.call_api(
"get",
("blobs", blob_id),
extra_headers=extra_headers,
expected_codes=(200, 302, 404),
stream=fileobj is not None,
)
if response.status_code in (302, 404):
return None
elif response.status_code == 200:
if fileobj is not None:
# NB Using shutil.copyfileobj is an attractive option, but does not
# decode the gzip and deflate transfer-encodings...
for chunk in response.iter_content(chunk_size=None):
fileobj.write(chunk)
return None
else:
return response.content
else:
raise encapsia_api.EncapsiaApiError(
"Unable to download blob {}: {}".format(
blob_id, response.status_code))
def get_result():
reply = self.get(("dbctl", "action", name, action_id))
rest_api_result = reply["result"]
action_status = rest_api_result["status"]
action_result = rest_api_result["result"]
if action_status == "finished":
return action_result
elif action_status == "failed":
raise encapsia_api.EncapsiaApiError(rest_api_result)
else:
return NoResultYet
def dbctl_download_data(self, handle, filename=None):
"""Download data and return (temp) filename."""
headers = {
"Accept": "*/*",
"Authorization": "Bearer {}".format(self.token)
}
url = "/".join([self.url, self.version, "dbctl/data", handle])
response = requests.get(url, headers=headers, verify=True, stream=True)
if response.status_code != 200:
raise encapsia_api.EncapsiaApiError("{} {}".format(
response.status_code, response.reason))
if filename is None:
_, filename = tempfile.mkstemp()
_stream_response_to_file(response, filename)
return filename
def discover_credentials(host=None):
"""Return (url, token) or raise EncapsiaApiError."""
if not host:
host = os.environ.get("ENCAPSIA_HOST")
if host:
store = CredentialsStore()
try:
url, token = store.get(host)
except KeyError:
raise encapsia_api.EncapsiaApiError(
f"Cannot find entry for '{host}' in encapsia credentials file."
)
else:
url, token = _get_env_var("ENCAPSIA_URL"), _get_env_var(
"ENCAPSIA_TOKEN")
return url, token
def download_blob_data(self, blob_id):
"""Download blob data for given blob_id."""
extra_headers = {"Accept": "*/*"}
response = self.call_api(
"get",
("blobs", blob_id),
extra_headers=extra_headers,
expected_codes=(200, 302, 404),
)
if response.status_code == 200:
return response.content
elif response.status_code in (302, 404):
return None
else:
raise encapsia_api.EncapsiaApiError(
"Unable to download blob {}: {}".format(blob_id, response.status_code)
)
def download_to_file(url, token, target_file=None, cleanup=True):
"""Context manager for downloading a fixed file to a target_file."""
if target_file is None:
filename = pathlib.Path(tempfile.mkstemp()[1])
else:
filename = pathlib.Path(target_file)
try:
headers = {"Accept": "*/*", "Authorization": "Bearer {}".format(token)}
response = requests.get(url, headers=headers, verify=True, stream=True)
if response.status_code != 200:
raise encapsia_api.EncapsiaApiError("{} {}".format(
response.status_code, response.reason))
stream_response_to_file(response, filename)
yield filename
finally:
if cleanup:
filename.unlink()
def dbctl_download_data(self, handle, filename=None):
"""Download data and return (temp) filename."""
headers = {
"Accept": "*/*",
"Authorization": "Bearer {}".format(self.token)
}
url = "/".join([self.url, self.version, "dbctl/data", handle])
response = requests.get(url, headers=headers, verify=True, stream=True)
if response.status_code != 200:
raise encapsia_api.EncapsiaApiError("{} {}".format(
response.status_code, response.reason))
if filename is None:
_, filename = tempfile.mkstemp()
with open(filename, "wb") as f:
for chunk in response.iter_content(chunk_size=1024):
f.write(chunk)
return filename
def get_task_result():
with self.call_api(
"get", ("tasks", namespace, task_id), stream=True
) as response:
if response.headers.get("Content-type") == "application/json":
reply = response.json()
if reply.get("status") != "ok":
raise encapsia_api.EncapsiaApiError(response.text)
rest_api_result = reply["result"]
task_status = rest_api_result["status"]
task_result = rest_api_result["result"]
if task_status == "finished":
if download:
filename = pathlib.Path(download)
with filename.open("wt") as f:
json.dump(task_result, f, indent=4)
return FileDownloadResponse(filename, "application/json")
else:
return task_result
elif task_status == "failed":
raise encapsia_api.EncapsiaApiFailedTaskError(
"Failed Task. See Exception payload attribute.",
payload=rest_api_result,
)
else:
return NoResultYet
else:
# Stream the response directly to the given file.
# Note we don't care whether this is JSON, CSV, or some other type.
if download:
filename = pathlib.Path(download)
stream_response_to_file(response, filename)
return FileDownloadResponse(
filename, response.headers.get("Content-type")
)
else:
return response.text
def _get_env_var(name):
try:
return os.environ[name]
except KeyError:
raise encapsia_api.EncapsiaApiError(
f"Environment variable {name} does not exist!")
