Skip to content

mesh-1786: generic retrieve chunk #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
matt-mercer merged 1 commit into from
Dec 8, 2023
Merged

mesh-1786: generic retrieve chunk #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 46 additions & 32 deletions mesh_client/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def _get_version(*names: str) -> str:
LIVE_CA_CERT = os.path.join(_PACKAGE_DIR, "nhs-live-ca-bundle.pem")


_OPTIONAL_HEADERS = {
_OPTIONAL_HEADERS: Dict[str, str] = {
"workflow_id": "Mex-WorkflowID",
"filename": "Mex-FileName",
"local_id": "Mex-LocalID",
Expand All @@ -85,6 +85,11 @@ def _get_version(*names: str) -> str:
"content_type": "Content-Type",
}


def optional_header_map() -> Dict[str, str]:
return _OPTIONAL_HEADERS.copy()


_SEND_HEADERS = {**_OPTIONAL_HEADERS}
_SEND_HEADERS.update(
{
Expand Down Expand Up @@ -345,7 +350,8 @@ def __init__( # noqa: C901
and whether messages should be compressed, transparently, before
sending.
"""

if isinstance(shared_key, str):
shared_key = shared_key.encode(encoding="utf-8")
shared_key = shared_key or get_shared_key_from_environ()

self._mailbox = mailbox
Expand Down Expand Up @@ -534,44 +540,52 @@ def retrieve_message(self, message_id: str) -> "Message":
https://digital.nhs.uk/developer/api-catalogue/message-exchange-for-social-care-and-health-api#get-/messageexchange/-mailbox_id-/inbox/-message_id-
"""
message_id = getattr(message_id, "_msg_id", message_id)
response = self._session.get(f"{self.mailbox_url}/inbox/{q(message_id)}", stream=True, timeout=self._timeout)
response.raise_for_status()
response = self.retrieve_message_chunk(message_id, chunk_num=1)
return Message(message_id, response, self)

def retrieve_message_chunk(self, message_id: str, chunk_num: Union[int, str]) -> Response:
def retrieve_message_chunk(self, message_id: str, chunk_num: int) -> Response:
"""
get a chunk
https://digital.nhs.uk/developer/api-catalogue/message-exchange-for-social-care-and-health-api#get-/messageexchange/-mailbox_id-/inbox/-message_id-/-chunk_number-
Args:
message_id (str): message id to receive
chunk_num (int): chunk number
message_id: message id to receive
chunk_num: chunk number

Returns:
Response: http response
"""
chunk_num = int(chunk_num)
uri = (
f"{self.mailbox_url}/inbox/{q(message_id)}/{chunk_num}"
if chunk_num > 1
else f"{self.mailbox_url}/inbox/{q(message_id)}"
)

response = self._session.get(
f"{self.mailbox_url}/inbox/{q(message_id)}/{chunk_num}",
uri,
stream=True,
timeout=self._timeout,
)
response.raise_for_status()
return response

@staticmethod
def _headers(recipient: str, chunk_no: int, total_chunks: int, compress: bool, **kwargs) -> Dict[str, str]:
def _headers_for_chunk(
recipient: str, chunk_num: int, total_chunks: int, compress: bool, **kwargs
) -> Dict[str, str]:
content_type = kwargs.get("content_type", "application/octet-stream")
if chunk_no > 1:
if chunk_num > 1:
headers = {
"Content-Type": content_type,
"Mex-Chunk-Range": f"{chunk_no}:{total_chunks}",
"Mex-Chunk-Range": f"{chunk_num}:{total_chunks}",
}
if compress:
headers["Content-Encoding"] = "gzip"
return headers

headers = {
"Content-Type": content_type,
"Mex-Chunk-Range": f"{chunk_no}:{total_chunks}",
"Mex-Chunk-Range": f"{chunk_num}:{total_chunks}",
"Mex-To": recipient,
}

Expand All @@ -593,7 +607,7 @@ def send_chunk(
self,
recipient: str,
chunk,
chunk_no: int,
chunk_num: int,
total_chunks: int,
compress: Optional[bool] = None,
message_id: Optional[str] = None,
Expand All @@ -604,36 +618,36 @@ def send_chunk(
Args:
recipient (str): mailbox id to send to
chunk: data to send
chunk_no: chunk number >= 1
chunk_num: chunk number >= 1
total_chunks: total number of chunks >= 1
compress: instruction to compress this (overrides the 'transparent_compress' at the mailbox level)
message_id: message id, required if chunk_no > 1
message_id: message id, required if chunk_num > 1
**kwargs: optional mesh header args workflow_id, filename, local_id, subject, encrypted,
compressed, checksum, partner_id, content_type

Returns:
Response: raw http response
"""

chunk_num = int(chunk_num)
compress = self._transparent_compress if compress is None else compress

def maybe_compressed(maybe_compress: bytes):
if not compress:
return maybe_compress
return GzipCompressStream(maybe_compress)

headers = self._headers(
recipient=recipient, chunk_no=chunk_no, total_chunks=total_chunks, compress=compress, **kwargs
headers = self._headers_for_chunk(
recipient=recipient, chunk_num=chunk_num, total_chunks=total_chunks, compress=compress, **kwargs
)

data = maybe_compressed(chunk)

buffer = data
if chunk_no > 1 and self._retries:
if chunk_num > 1 and self._retries:
# urllib3 body_pos requires a seekable stream to allow rewinding on retry ( we never retry the first chunk )
buffer = BytesIO(data.read() if hasattr(data, "read") else data)

if chunk_no == 1:
if chunk_num == 1:
assert not message_id, "message_id should not be sent with the first chunk"

response = self._session.post(
Expand All @@ -651,7 +665,7 @@ def maybe_compressed(maybe_compress: bytes):
assert message_id, "message_id is required for chunks number >= 2"

response = self._session.post(
f"{self.mailbox_url}/outbox/{q(message_id)}/{chunk_no}",
f"{self.mailbox_url}/outbox/{q(message_id)}/{chunk_num}",
data=buffer,
headers=headers,
timeout=self._timeout,
Expand Down Expand Up @@ -715,7 +729,7 @@ def send_message(
total_chunks = len(chunks)

response1 = self.send_chunk(
recipient=recipient, chunk=first_chunk, chunk_no=1, total_chunks=total_chunks, compress=compress, **kwargs
recipient=recipient, chunk=first_chunk, chunk_num=1, total_chunks=total_chunks, compress=compress, **kwargs
)

response_dict = response1.json()
Expand All @@ -734,7 +748,7 @@ def send_message(
_response = self.send_chunk(
recipient=recipient,
chunk=chunk,
chunk_no=chunk_num,
chunk_num=chunk_num,
message_id=message_id,
total_chunks=total_chunks,
compress=compress,
Expand Down Expand Up @@ -937,10 +951,10 @@ def __init__(self, msg_id: str, response, client):
def maybe_decompress(resp):
return GzipDecompressStream(resp.raw) if resp.headers.get("Content-Encoding") == "gzip" else resp.raw

self._response = CombineStreams(
self._stream = CombineStreams(
chain(
[maybe_decompress(response)],
(maybe_decompress(client.retrieve_message_chunk(msg_id, str(i + 2))) for i in range(chunk_count - 1)),
(maybe_decompress(client.retrieve_message_chunk(msg_id, i + 2)) for i in range(chunk_count - 1)),
)
)

Expand All @@ -957,27 +971,27 @@ def read(self, n=None) -> bytes:
Read up to n bytes from the message, or read the remainder of the
message, if n is not provided.
"""
return self._response.read(n)
return self._stream.read(n)

def readline(self) -> bytes:
"""
Read a single line from the message
"""
return self._response.readline()
return self._stream.readline()

def readlines(self) -> List[bytes]:
"""
Read all lines from the message
"""
return self._response.readlines()
return self._stream.readlines()

def close(self):
"""Close the stream underlying this message"""
if hasattr(self._response, "close"):
if hasattr(self._stream, "close"):
try:
self._response.close()
self._stream.close()
finally:
self._response = None # type: ignore[assignment]
self._stream = None # type: ignore[assignment]

def acknowledge(self):
"""
Expand Down Expand Up @@ -1014,7 +1028,7 @@ def __iter__(self):
"""
Iterate through lines of the message
"""
return iter(self._response)
return iter(self._stream)


class Message(_BaseMessage, _MessageAttrs):
Expand Down
47 changes: 39 additions & 8 deletions poetry.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ types-pkg-resources = "^0.1.3"
coverage = "^7.2.7"
pytest = "^7.4.0"
importlib-metadata = {version = ">=4.11.4", python = "<3.12"}
#importlib-resources = {version = "*", python = "<3.9"}
pytest-httpserver = {version = "^1.0.8", python = ">=3.8,<4.0"}
ruff = "^0.1.5"
boto3 = "^1.33.7"
Expand Down
22 changes: 11 additions & 11 deletions tests/mesh_noretry_send_chunk_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,14 +61,14 @@ def test_chunk_retries(httpserver: HTTPServer, alice: MeshClient, bob: MeshClien
def send_chunk_handler(request: Request):
last_path = request.path.split("/")[-1]

chunk_no = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_no] += 1
chunk_num = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_num] += 1

if chunk_no == 1:
if chunk_num == 1:
received_chunks.append(request.data)
return json_response(cast(SendMessageResponse_v2, {"message_id": message_id}), status=202)

if chunk_no == 2 and chunk_call_counts[chunk_no] < 4:
if chunk_num == 2 and chunk_call_counts[chunk_num] < 4:
return plain_response("", status=502)

received_chunks.append(request.data)
Expand Down Expand Up @@ -101,10 +101,10 @@ def test_chunk_all_retries_fail(httpserver: HTTPServer, alice: MeshClient, bob:
def send_chunk_handler(request: Request):
last_path = request.path.split("/")[-1]

chunk_no = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_no] += 1
chunk_num = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_num] += 1

if chunk_no == 1:
if chunk_num == 1:
received_chunks.append(request.data)
return json_response(cast(SendMessageResponse_v2, {"message_id": message_id}), status=202)

Expand Down Expand Up @@ -139,14 +139,14 @@ def test_chunk_retries_with_file(httpserver: HTTPServer, alice: MeshClient, bob:
def send_chunk_handler(request: Request):
last_path = request.path.split("/")[-1]

chunk_no = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_no] += 1
chunk_num = int(last_path) if last_path.isdigit() else 1
chunk_call_counts[chunk_num] += 1

if chunk_no == 1:
if chunk_num == 1:
received_chunks.append(request.data)
return json_response(cast(SendMessageResponse_v2, {"message_id": message_id}), status=202)

if chunk_no == 2 and chunk_call_counts[chunk_no] < 4:
if chunk_num == 2 and chunk_call_counts[chunk_num] < 4:
return plain_response("", status=502)

received_chunks.append(request.data)
Expand Down
Loading