diff --git a/Makefile b/Makefile index 3e138d4..e4a1ea4 100644 --- a/Makefile +++ b/Makefile @@ -53,6 +53,7 @@ build-dep-fedora: python3-mypy \ python3-pytest \ python3-requests \ + python3-requests-toolbelt \ python3-types-requests \ python3-setuptools_scm \ rpmdevtools \ diff --git a/README.rst b/README.rst index b091f43..2429212 100644 --- a/README.rst +++ b/README.rst @@ -343,9 +343,25 @@ Delete a service integration:: $ avn service integration-delete --project +Custom Files +-------------------- + +Listing files:: + + $ anv service custom-file list --project + +Reading file:: + + $ anv service custom-file get --project --file_id [--target_filepath ] [--stdout_write] + + +Uploading new files:: + $ avn service custom-file upload --project --file_type --file_path --file_name +Updating existing files:: + $ avn service custom-file update --project --file_path --file_id .. _teams: Working with Teams diff --git a/aiven/client/cli.py b/aiven/client/cli.py index 26b95b3..28dd199 100644 --- a/aiven/client/cli.py +++ b/aiven/client/cli.py @@ -5546,6 +5546,77 @@ def service__flink__cancel_application_deployments(self) -> None: ), ) + @arg.project + @arg.service_name + def service__custom_file__list(self) -> None: + """List all the custom files for the specified service""" + self.print_response( + self.client.custom_file_list( + project=self.args.project, + service=self.args.service_name, + ), + ) + + @arg.project + @arg.service_name + @arg("--file_id", help="A file ID, usually an uuid4 string", required=True) + @arg("--target_filepath", help="Where to save the file downloaded") + @arg("--stdout_write", help="Write to the stdout instead of file", action="store_true") + def service__custom_file__get(self) -> None: + """ + Download the custom file for the specified service and ID + Requires at least `--target_filepath` or `--stdout_write`, but able to handle both + """ + if not (self.args.target_filepath or self.args.stdout_write): + raise argx.UserError("You need to specify `--target_filepath` or `--stdout_write` or both") + + result = self.client.custom_file_get( + project=self.args.project, + service=self.args.service_name, + file_id=self.args.file_id, + ) + + if self.args.target_filepath: + with open(self.args.target_filepath, "wb") as f: + f.write(result) + + if self.args.stdout_write: + print(result.decode("utf-8")) + + @arg.project + @arg.service_name + @arg("--file_path", help="A path to the file", required=True) + @arg("--file_type", choices=["synonyms", "stopwords", "wordnet"], required=True) + @arg("--file_name", help="A name for the file", required=True) + def service__custom_file__upload(self) -> None: + """Upload custom file for the specified service""" + with open(self.args.file_path, "rb") as f: + self.print_response( + self.client.custom_file_upload( + project=self.args.project, + service=self.args.service_name, + file_type=self.args.file_type, + file_name=self.args.file_name, + file_object=f, + ), + ) + + @arg.project + @arg.service_name + @arg("--file_path", help="A path to the file", required=True) + @arg("--file_id", help="A file ID, usually an uuid4 string", required=True) + def service__custom_file__update(self) -> None: + """Update custom file for the specified service and ID""" + with open(self.args.file_path, "rb") as f: + self.print_response( + self.client.custom_file_update( + project=self.args.project, + service=self.args.service_name, + file_id=self.args.file_id, + file_object=f, + ), + ) + @arg.json @arg("name", help="Name of the organization to create") @arg.force diff --git a/aiven/client/client.py b/aiven/client/client.py index fa89278..bf4867f 100644 --- a/aiven/client/client.py +++ b/aiven/client/client.py @@ -8,7 +8,8 @@ from .session import get_requests_session from http import HTTPStatus from requests import Response -from typing import Any, Callable, Collection, Mapping, Sequence, TypedDict +from requests_toolbelt import MultipartEncoder # type: ignore +from typing import Any, BinaryIO, Callable, Collection, Mapping, Sequence, TypedDict from urllib.parse import quote import json @@ -79,6 +80,10 @@ def _execute(self, func: Callable, method: str, path: str, body: Any, params: An headers["content-type"] = "application/json" data = json.dumps(body) log_data = json.dumps(body, sort_keys=True, indent=4) + elif isinstance(body, MultipartEncoder): + headers["content-type"] = body.content_type + data = body + log_data = data else: headers["content-type"] = "application/octet-stream" data = body @@ -174,11 +179,28 @@ def verify( ) time.sleep(0.2) + return self._process_response(response=response, op=op, path=path, result_key=result_key) + + @staticmethod + def build_path(*parts: str) -> str: + return "/" + "/".join(quote(part, safe="") for part in parts) + + def _process_response( + self, + response: Response, + op: Callable[..., Response], + path: str, + result_key: str | None = None, + ) -> Mapping | bytes: # Check API is actually returning data or not if response.status_code == HTTPStatus.NO_CONTENT or len(response.content) == 0: return {} + if response.headers.get("Content-Type") == "application/octet-stream": + return response.content + result = response.json() + if result.get("error"): raise ResponseError( "server returned error: {op} {base_url}{path} {result}".format( @@ -190,10 +212,6 @@ def verify( return result[result_key] return result - @staticmethod - def build_path(*parts: str) -> str: - return "/" + "/".join(quote(part, safe="") for part in parts) - class AivenClient(AivenClientBase): """Aiven Client with high-level operations""" @@ -2138,6 +2156,54 @@ def clickhouse_database_list(self, project: str, service: str) -> Mapping: path = self.build_path("project", project, "service", service, "clickhouse", "db") return self.verify(self.get, path, result_key="databases") + def custom_file_list(self, project: str, service: str) -> Mapping: + path = self.build_path("project", project, "service", service, "file") + return self.verify(self.get, path) + + def custom_file_get(self, project: str, service: str, file_id: str) -> bytes: + path = self.build_path("project", project, "service", service, "file", file_id) + return self.verify(self.get, path) + + def custom_file_upload( + self, + project: str, + service: str, + file_type: str, + file_object: BinaryIO, + file_name: str, + update: bool = False, + ) -> Mapping: + path = self.build_path("project", project, "service", service, "file") + return self.verify( + self.post, + path, + body=MultipartEncoder( + fields={ + "file": (file_name, file_object, "application/octet-stream"), + "filetype": file_type, + "filename": file_name, + } + ), + ) + + def custom_file_update( + self, + project: str, + service: str, + file_object: BinaryIO, + file_id: str, + ) -> Mapping: + path = self.build_path("project", project, "service", service, "file", file_id) + return self.verify( + self.put, + path, + body=MultipartEncoder( + fields={ + "file": (file_id, file_object, "application/octet-stream"), + } + ), + ) + def flink_list_applications( self, *, diff --git a/pyproject.toml b/pyproject.toml index 5d3c8d6..21b6146 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dynamic = ["version"] dependencies = [ 'requests>=2.2.1; sys_platform == "linux"', 'requests>=2.9.1; sys_platform != "linux"', + "requests-toolbelt>=0.9.0", "certifi>=2015.11.20.1", ] [project.optional-dependencies] diff --git a/tests/test_cli.py b/tests/test_cli.py index 334f0d2..4b9ed06 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -21,6 +21,7 @@ import pytest import random import string +import tempfile import uuid EXIT_CODE_INVALID_USAGE = 2 @@ -1467,3 +1468,214 @@ def test_project_update__parent_id_as_org_id_requested_correctly() -> None: project=project_name, tech_emails=None, ) + + +def test_custom_files_list(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = { + "custom_files": [ + { + "create_time": "2023-01-01T00:00:00Z", + "file_id": "a4d83aaa-35ae-51b4-b500-81e743bfe906", + "filename": "foo1", + "filesize": 25, + "filetype": "synonyms", + "service_reference": "custom/synonyms/foo1", + "update_time": "2023-01-01T00:00:00Z", + } + ] + } + + aiven_client.custom_file_list.return_value = return_value + + build_aiven_cli(aiven_client).run(args=["service", "custom-file", "list", "--project", "test", "foo"]) + aiven_client.custom_file_list.assert_called_with(project="test", service="foo") + captured = capsys.readouterr() + assert json.loads(captured.out) == return_value + + +def test_custom_files_get_stdout(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = b"foo => bar" + + aiven_client.custom_file_get.return_value = return_value + + build_aiven_cli(aiven_client).run( + args=[ + "service", + "custom-file", + "get", + "--project", + "test", + "--file_id", + "a4d83aaa-35ae-51b4-b500-81e743bfe906", + "--stdout_write", + "foo", + ] + ) + aiven_client.custom_file_get.assert_called_with( + project="test", service="foo", file_id="a4d83aaa-35ae-51b4-b500-81e743bfe906" + ) + captured = capsys.readouterr() + assert captured.out.strip() == return_value.decode("utf-8") + + +def test_custom_files_get_file(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = b"foo => bar" + + aiven_client.custom_file_get.return_value = return_value + with tempfile.NamedTemporaryFile(delete=False) as f_temp: + # Closing file so Windows could let cli open it + file_name = f_temp.name + build_aiven_cli(aiven_client).run( + args=[ + "service", + "custom-file", + "get", + "--project", + "test", + "--file_id", + "a4d83aaa-35ae-51b4-b500-81e743bfe906", + "--target_filepath", + file_name, + "foo", + ] + ) + aiven_client.custom_file_get.assert_called_with( + project="test", service="foo", file_id="a4d83aaa-35ae-51b4-b500-81e743bfe906" + ) + captured = capsys.readouterr() + assert captured.out.strip() == "" + with open(file_name, "rb") as f: + assert f.read() == return_value + + +def test_custom_files_get_both(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = b"foo => bar" + + aiven_client.custom_file_get.return_value = return_value + with tempfile.NamedTemporaryFile(delete=False) as f: + # Closing file so Windows could let cli open it + file_name = f.name + build_aiven_cli(aiven_client).run( + args=[ + "service", + "custom-file", + "get", + "--project", + "test", + "--file_id", + "a4d83aaa-35ae-51b4-b500-81e743bfe906", + "--target_filepath", + file_name, + "--stdout_write", + "foo", + ] + ) + aiven_client.custom_file_get.assert_called_with( + project="test", service="foo", file_id="a4d83aaa-35ae-51b4-b500-81e743bfe906" + ) + captured = capsys.readouterr() + assert captured.out.strip() == return_value.decode("utf-8") + with open(file_name, "rb") as f: + assert f.read() == return_value + + +def test_custom_files_get_none(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = b"foo => bar" + + aiven_client.custom_file_get.return_value = return_value + aiven_cli = build_aiven_cli(aiven_client) + aiven_cli.run( + args=[ + "service", + "custom-file", + "get", + "--project", + "test", + "--file_id", + "a4d83aaa-35ae-51b4-b500-81e743bfe906", + "foo", + ] + ) + aiven_client.custom_file_get.assert_not_called() + captured = capsys.readouterr() + assert captured.out.strip() == "" + with pytest.raises(argx.UserError): + aiven_cli.service__custom_file__get() + + +def test_custom_files_upload(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = { + "file_id": "d19878d5-2726-5773-81cb-ff61562c892c", + "message": "created", + "service_reference": "custom/synonyms/foo2", + } + + aiven_client.custom_file_upload.return_value = return_value + with tempfile.NamedTemporaryFile(delete=False) as f: + # Closing file so Windows could let cli open it + file_name = f.name + build_aiven_cli(aiven_client).run( + args=[ + "service", + "custom-file", + "upload", + "--project", + "test", + "--file_path", + file_name, + "--file_type", + "synonyms", + "--file_name", + "foo2", + "foo", + ] + ) + # Can't check args as the file is reopened + aiven_client.custom_file_upload.assert_called_once() + captured = capsys.readouterr() + assert json.loads(captured.out.strip()) == return_value + + +def test_custom_files_update(capsys: CaptureFixture[str]) -> None: + aiven_client = mock.Mock(spec_set=AivenClient) + + return_value = { + "file_id": "d19878d5-2726-5773-81cb-ff61562c892c", + "message": "updated", + "service_reference": "custom/synonyms/foo2", + } + + aiven_client.custom_file_update.return_value = return_value + with tempfile.NamedTemporaryFile(delete=False) as f: + # Closing file so Windows could let cli open it + file_name = f.name + build_aiven_cli(aiven_client).run( + args=[ + "service", + "custom-file", + "update", + "--project", + "test", + "--file_path", + file_name, + "--file_id", + "d19878d5-2726-5773-81cb-ff61562c892c", + "foo", + ] + ) + # Can't check args as the file is reopened + aiven_client.custom_file_update.assert_called_once() + captured = capsys.readouterr() + assert json.loads(captured.out.strip()) == return_value diff --git a/tests/test_client.py b/tests/test_client.py index 90c9140..dcbfdfa 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -5,6 +5,7 @@ from __future__ import annotations from aiven.client import AivenClient +from aiven.client.client import ResponseError from http import HTTPStatus from typing import Any from unittest import mock @@ -14,14 +15,23 @@ class MockResponse: - def __init__(self, status_code: int, json_data: dict[str, Any] | None = None, headers: dict[str, str] | None = None): + def __init__( + self, + status_code: int, + json_data: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + content: bytes | None = None, + ): self.status_code = status_code self.json_data = json_data - if json_data is not None: + if content is not None: + self.content = content + elif json_data is not None: self.content = json.dumps(json_data).encode("utf-8") else: self.content = b"" self.headers = {} if headers is None else headers + self.text = self.content.decode("utf-8") def json(self) -> Any: return self.json_data @@ -41,3 +51,94 @@ def test_no_content_returned_from_api(response: MockResponse) -> None: assert aiven_client.verify(aiven_client.patch, "/") == {} assert aiven_client.verify(aiven_client.put, "/") == {} assert aiven_client.verify(aiven_client.delete, "/") == {} + + +@pytest.mark.parametrize( + "response,expected_result", + [ + ( + MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/octet-stream"}, + content=b"foo", + ), + b"foo", + ), + ( + MockResponse(status_code=HTTPStatus.OK, headers={"Content-Type": "application/json"}, json_data={}), + {}, + ), + ( + MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/json"}, + json_data={}, + content=b"foo", + ), + {}, + ), + ( + MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/json"}, + json_data={"foo": "bar"}, + content=b"foo", + ), + {"foo": "bar"}, + ), + ( + MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/octet-stream"}, + json_data={"foo": "bar"}, + content=b"foo", + ), + b"foo", + ), + ], +) +def test_response_processing(response: MockResponse, expected_result: Any) -> None: + def operation() -> MockResponse: + return response + + aiven_client = AivenClient("") + assert aiven_client._process_response(response=response, op=operation, path="") == expected_result # type: ignore + + +def test_response_processing_result_key() -> None: + expected_value = 2 + response = MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/json"}, + json_data={"foo": 1, "bar": expected_value, "spam": 3}, + ) + + def operation() -> MockResponse: + return response + + aiven_client = AivenClient("") + assert ( + aiven_client._process_response(response=response, op=operation, path="", result_key="bar") # type: ignore + == expected_value + ) + + +def test_response_processing_error_raise() -> None: + response = MockResponse( + status_code=HTTPStatus.OK, + headers={"Content-Type": "application/json"}, + json_data={"foo": 1, "bar": 2, "spam": 3, "error": "test error"}, + ) + + def operation() -> MockResponse: + """test operation""" + return response + + aiven_client = AivenClient("test_base_url") + + with pytest.raises(ResponseError) as e: + aiven_client._process_response(response=response, op=operation, path="", result_key="bar") # type: ignore + assert ( + str(e) + == "server returned error: test operation test_base_url {'foo': 1, 'bar': 2, 'spam': 3, 'error': 'test error'}" + )