diff --git a/README.rst b/README.rst index 6f5dfb4..4f05405 100644 --- a/README.rst +++ b/README.rst @@ -418,7 +418,11 @@ The history of a bucket can also be purged with: Attachments ----------- -If the `kinto-attachment plugin `_ is enabled, it is possible to add attachments on records: +If the `kinto-attachment plugin `_ 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 diff --git a/src/kinto_http/client.py b/src/kinto_http/client.py index c4793a1..c7e178e 100644 --- a/src/kinto_http/client.py +++ b/src/kinto_http/client.py @@ -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, diff --git a/tests/test_client.py b/tests/test_client.py index b15620a..e2bce35 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,3 +1,4 @@ +import os import re import pytest @@ -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)