add BYOC tags list/update/replace commands [MILK-337]

aiven/client/cli.py

@@ -20,7 +20,7 @@
 from datetime import datetime, timedelta, timezone
 from decimal import Decimal
 from http import HTTPStatus
-from typing import Any, Callable, Final, IO, Mapping, Protocol, Sequence
+from typing import Any, Callable, Final, IO, Mapping, Optional, Protocol, Sequence, TypeVar
 from urllib.parse import urlparse
 
 import errno
@@ -33,6 +33,8 @@
 import sys
 import time
 
+S = TypeVar("S", str, Optional[str])  # Must be exactly str or str | None
+
 USER_GROUP_COLUMNS = [
     "user_group_name",
     "user_group_id",
@@ -5951,6 +5953,7 @@ def byoc__update(self) -> None:
             cloud_region=self.args.cloud_region,
             reserved_cidr=self.args.reserved_cidr,
             display_name=self.args.display_name,
+            tags=None,
         )
         self.print_response(output)
 
@@ -6057,6 +6060,52 @@ def byoc__cloud__permissions__remove(self) -> None:
             )
         )
 
+    @staticmethod
+    def add_prefix_to_keys(prefix: str, tags: Mapping[str, S]) -> Mapping[str, S]:
+        return {f"{prefix}{k}": v for (k, v) in tags.items()}
+
+    @staticmethod
+    def remove_prefix_from_keys(prefix: str, tags: Mapping[str, str]) -> Mapping[str, str]:
+        return {(k.partition(prefix)[-1] if k.startswith(prefix) else k): v for (k, v) in tags.items()}
+
+    @arg.json
+    @arg("--organization-id", required=True, help="Identifier of the organization of the custom cloud environment")
+    @arg("--byoc-id", required=True, help="Identifier of the custom cloud environment that defines the BYOC cloud")
+    def byoc__tags__list(self) -> None:
+        """List BYOC tags"""
+        tags = self.client.list_byoc_tags(organization_id=self.args.organization_id, byoc_id=self.args.byoc_id)
+        # Remove the "byoc_resource_tag:" prefix from BYOC tag keys to print them as expected by the end user.
+        self._print_tags({"tags": self.remove_prefix_from_keys("byoc_resource_tag:", tags.get("tags", {}))})
+
+    @arg.json
+    @arg("--organization-id", required=True, help="Identifier of the organization of the custom cloud environment")
+    @arg("--byoc-id", required=True, help="Identifier of the custom cloud environment that defines the BYOC cloud")
+    @arg("--add-tag", help="Add a new tag (key=value)", action="append", default=[])
+    @arg("--remove-tag", help="Remove the named tag", action="append", default=[])
+    def byoc__tags__update(self) -> None:
+        """Add or remove BYOC tags"""
+        response = self.client.update_byoc_tags(
+            organization_id=self.args.organization_id,
+            byoc_id=self.args.byoc_id,
+            # Add the "byoc_resource_tag:" prefix to BYOC tag keys to make them cascade to the Bastion service.
+            tag_updates=self.add_prefix_to_keys("byoc_resource_tag:", self._tag_update_body_from_args()),
+        )
+        print(response["message"])
+
+    @arg.json
+    @arg("--organization-id", required=True, help="Identifier of the organization of the custom cloud environment")
+    @arg("--byoc-id", required=True, help="Identifier of the custom cloud environment that defines the BYOC cloud")
+    @arg("--tag", help="Tag for service (key=value)", action="append", default=[])
+    def byoc__tags__replace(self) -> None:
+        """Replace BYOC tags, deleting any old ones first"""
+        response = self.client.replace_byoc_tags(
+            organization_id=self.args.organization_id,
+            byoc_id=self.args.byoc_id,
+            # Add the "byoc_resource_tag:" prefix to BYOC tag keys to make them cascade to the Bastion service.
+            tags=self.add_prefix_to_keys("byoc_resource_tag:", self._tag_replace_body_from_args()),
+        )
+        print(response["message"])
+
     @arg.json
     @arg.project
     @arg.service_name

aiven/client/client.py

@@ -2734,6 +2734,7 @@ def byoc_update(
         cloud_region: str | None,
         reserved_cidr: str | None,
         display_name: str | None,
+        tags: Mapping[str, str | None] | None,
     ) -> Mapping[Any, Any]:
         body = {
             key: value
@@ -2743,6 +2744,7 @@ def byoc_update(
                 "cloud_region": cloud_region,
                 "reserved_cidr": reserved_cidr,
                 "display_name": display_name,
+                "tags": tags,
             }.items()
             if value is not None
         }
@@ -2836,6 +2838,48 @@ def byoc_permissions_set(
             body={"accounts": accounts, "projects": projects},
         )
 
+    def list_byoc_tags(self, organization_id: str, byoc_id: str) -> Mapping:
+        output = self.byoc_update(
+            organization_id=organization_id,
+            byoc_id=byoc_id,
+            # Putting all arguments to `None` makes `byoc_update()` behave like a `GET BYOC BY ID` API which does not exist.
+            deployment_model=None,
+            cloud_provider=None,
+            cloud_region=None,
+            reserved_cidr=None,
+            display_name=None,
+            tags=None,
+        )
+        return {"tags": output.get("custom_cloud_environment", {}).get("tags", {})}
+
+    def update_byoc_tags(self, organization_id: str, byoc_id: str, tag_updates: Mapping[str, str | None]) -> Mapping:
+        self.byoc_update(
+            organization_id=organization_id,
+            byoc_id=byoc_id,
+            deployment_model=None,
+            cloud_provider=None,
+            cloud_region=None,
+            reserved_cidr=None,
+            display_name=None,
+            tags=tag_updates,
+        )
+        # There have been no errors raised
+        return {"message": "tags updated"}
+
+    def replace_byoc_tags(self, organization_id: str, byoc_id: str, tags: Mapping[str, str]) -> Mapping:
+        self.byoc_update(
+            organization_id=organization_id,
+            byoc_id=byoc_id,
+            deployment_model=None,
+            cloud_provider=None,
+            cloud_region=None,
+            reserved_cidr=None,
+            display_name=None,
+            tags=tags,
+        )
+        # There have been no errors raised
+        return {"message": "tags updated"}
+
     def alloydbomni_google_cloud_private_key_set(self, *, project: str, service: str, private_key: str) -> dict[str, Any]:
         return self.verify(
             self.post,

tests/test_cli.py

@@ -1787,6 +1787,7 @@ def test_byoc_update() -> None:
         cloud_region="eu-west-2",
         reserved_cidr="10.1.0.0/24",
         display_name="Another name",
+        tags=None,
     )
 
 
@@ -1864,3 +1865,130 @@ def test_byoc_delete() -> None:
         organization_id="org123456789a",
         byoc_id="d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
     )
+
+
+def test_add_prefix_to_keys() -> None:
+    prefix = "byoc_resource_tag:"
+    tags = {
+        "key_1": "value_1",
+        "key_2": "",
+        "key_3": None,
+        "byoc_resource_tag:key_4": "value_4",
+        "key_5": "byoc_resource_tag:keep-the-whole-value-5",
+    }
+    expected_output = {
+        "byoc_resource_tag:key_1": "value_1",
+        "byoc_resource_tag:key_2": "",
+        "byoc_resource_tag:key_3": None,
+        "byoc_resource_tag:byoc_resource_tag:key_4": "value_4",
+        "byoc_resource_tag:key_5": "byoc_resource_tag:keep-the-whole-value-5",
+    }
+    output = AivenCLI.add_prefix_to_keys(prefix, tags)
+    assert output == expected_output
+
+
+def test_remove_prefix_from_keys() -> None:
+    prefix = "byoc_resource_tag:"
+    tags = {
+        "byoc_resource_tag:key_1": "value_1",
+        "byoc_resource_tag:key_2": "",
+        "byoc_resource_tag:byoc_resource_tag:key_3": "value_3",
+        "key_4": "value_4",
+        "byoc_resource_tag:key_5": "byoc_resource_tag:keep-the-whole-value-5",
+    }
+    expected_output = {
+        "key_1": "value_1",
+        "key_2": "",
+        "byoc_resource_tag:key_3": "value_3",
+        "key_4": "value_4",
+        "key_5": "byoc_resource_tag:keep-the-whole-value-5",
+    }
+    output = AivenCLI.remove_prefix_from_keys(prefix, tags)
+    assert output == expected_output
+
+
+def test_byoc_tags_list() -> None:
+    aiven_client = mock.Mock(spec_set=AivenClient)
+    aiven_client.list_byoc_tags.return_value = {
+        "tags": {
+            "byoc_resource_tag:key_1": "value_1",
+            "byoc_resource_tag:key_2": "",
+            "byoc_resource_tag:key_3": "value_3",
+            "byoc_resource_tag:key_4": "",
+            "byoc_resource_tag:key_5": "byoc_resource_tag:keep-the-whole-value-5",
+        },
+    }
+    args = [
+        "byoc",
+        "tags",
+        "list",
+        "--organization-id=org123456789a",
+        "--byoc-id=d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+    ]
+    build_aiven_cli(aiven_client).run(args=args)
+    aiven_client.list_byoc_tags.assert_called_once_with(
+        organization_id="org123456789a",
+        byoc_id="d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+    )
+
+
+def test_byoc_tags_update() -> None:
+    aiven_client = mock.Mock(spec_set=AivenClient)
+    aiven_client.update_byoc_tags.return_value = {"message": "tags updated"}
+    args = [
+        "byoc",
+        "tags",
+        "update",
+        "--organization-id=org123456789a",
+        "--byoc-id=d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+        "--add-tag",
+        "key_1=value_1",
+        "--add-tag",
+        "key_2=",
+        "--remove-tag",
+        "key_3",
+        "--remove-tag",
+        "byoc_resource_tag:key_4",
+        "--add-tag",
+        "key_5=byoc_resource_tag:keep-the-whole-value-5",
+    ]
+    build_aiven_cli(aiven_client).run(args=args)
+    aiven_client.update_byoc_tags.assert_called_once_with(
+        organization_id="org123456789a",
+        byoc_id="d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+        tag_updates={
+            "byoc_resource_tag:key_1": "value_1",
+            "byoc_resource_tag:key_2": "",
+            "byoc_resource_tag:key_3": None,
+            "byoc_resource_tag:byoc_resource_tag:key_4": None,
+            "byoc_resource_tag:key_5": "byoc_resource_tag:keep-the-whole-value-5",
+        },
+    )
+
+
+def test_byoc_tags_replace() -> None:
+    aiven_client = mock.Mock(spec_set=AivenClient)
+    aiven_client.replace_byoc_tags.return_value = {"message": "tags updated"}
+    args = [
+        "byoc",
+        "tags",
+        "replace",
+        "--organization-id=org123456789a",
+        "--byoc-id=d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+        "--tag",
+        "key_1=value_1",
+        "--tag",
+        "key_2=",
+        "--tag",
+        "byoc_resource_tag:key_3=byoc_resource_tag:keep-the-whole-value-3",
+    ]
+    build_aiven_cli(aiven_client).run(args=args)
+    aiven_client.replace_byoc_tags.assert_called_once_with(
+        organization_id="org123456789a",
+        byoc_id="d6a490ad-f43d-49d8-b3e5-45bc5dbfb387",
+        tags={
+            "byoc_resource_tag:key_1": "value_1",
+            "byoc_resource_tag:key_2": "",
+            "byoc_resource_tag:byoc_resource_tag:key_3": "byoc_resource_tag:keep-the-whole-value-3",
+        },
+    )