Skip to content

Commit

Permalink
Add download_attachment() method (#381)
Browse files Browse the repository at this point in the history
  • Loading branch information
leplatrem authored Nov 22, 2024
1 parent acd0571 commit 24853fc
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 1 deletion.
6 changes: 5 additions & 1 deletion README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,11 @@ The history of a bucket can also be purged with:
Attachments
-----------

If the `kinto-attachment plugin <https://github.com/Kinto/kinto-attachment/>`_ is enabled, it is possible to add attachments on records:
If the `kinto-attachment plugin <https://github.com/Kinto/kinto-attachment/>`_ is enabled, it is possible to fetch, add, or remove attachments on records:

.. code-block:: python
filepath = client.download_attachment(record_obj)
.. code-block:: python
Expand Down
27 changes: 27 additions & 0 deletions src/kinto_http/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -880,6 +880,33 @@ def purge_history(self, *, bucket=None, safe=True, if_match=None) -> List[Dict]:
resp, _ = self.session.request("delete", endpoint, headers=headers)
return resp["data"]

@retry_timeout
def download_attachment(
self,
record,
filepath=None,
chunk_size=8 * 1024,
):
if "attachment" not in record:
raise ValueError("Specified record has no attachment")

server_info = self.server_info()
base_url = server_info["capabilities"]["attachments"]["base_url"]
location = record["attachment"]["location"]
url = base_url + location

if filepath is None:
filepath = record["attachment"]["filename"]
elif os.path.isdir(filepath):
filepath = os.path.join(filepath, record["attachment"]["filename"])

with open(filepath, "wb") as f:
with requests.get(url, stream=True) as r:
r.raise_for_status()
for chunk in r.iter_content(chunk_size=chunk_size):
f.write(chunk)
return filepath

@retry_timeout
def add_attachment(
self,
Expand Down
29 changes: 29 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import re

import pytest
Expand Down Expand Up @@ -1394,6 +1395,34 @@ def test_purging_of_history(client_setup: Client):
client.session.request.assert_called_with("delete", url, headers=None)


def test_download_attachment(client_setup: Client, mocker: MockerFixture):
client = client_setup

client.session.request.return_value = (
{"capabilities": {"attachments": {"base_url": "https://cdn/"}}},
{},
)

mock_requests_get = mocker.patch("kinto_http.requests.get")
mock_response = mocker.MagicMock()
mock_response.iter_content = mocker.MagicMock(return_value=[b"chunk1", b"chunk2", b"chunk3"])
mock_response.raise_for_status = mocker.MagicMock()
mock_requests_get.return_value.__enter__.return_value = mock_response

with pytest.raises(ValueError):
client.download_attachment({})

record = {"attachment": {"location": "file.bin", "filename": "local.bin"}}

path = client.download_attachment(record)
assert path == "local.bin"
with open(path) as f:
assert f.read() == "chunk1chunk2chunk3"

path = client.download_attachment(record, filepath="/tmp")
assert os.path.exists("/tmp/local.bin")


def test_add_attachment_guesses_mimetype(record_setup: Client, tmp_path):
client = record_setup
mock_response(client.session)
Expand Down

0 comments on commit 24853fc

Please sign in to comment.