Merge pull request #1213 from james-garner-canonical/2024-11/feat/unify-handling-of-storage-constraints-in-deploy

#1213

#### Description

This PR unifies storage constraint parsing into a single method (`juju.constraints.parse_storage_constraints`), which is called in a single place (`juju.model.Model._deploy`), allowing both bundle and model deployments to specify storage constraints using either the [juju storage constraint directive format](https://juju.is/docs/juju/storage-constraint) (e.g. `{'label': 'ebs,100G'}`) or pre-parsed dictionaries (e.g. `{'label': {'count': 1, 'pool': 'ebs', 'size': 102400}`).

#### QA Steps

Unit tests have been updated to reflect the new parsing location. Integration tests have been added to verify that model deployment can request storage using either format. The existing bundle integration tests should continue to pass.


#### Notes & Discussion

This PR resolves the issues with storage constraint parsing identified in:
- #1052
- #1075

#1052 was initially addressed in #1053, which was included in the [3.5.2.0](https://github.com/juju/python-libjuju/releases/tag/3.5.2.0) release. This allowed bundle deployments (using `juju.bundle.AddApplicationChange.run`) to correctly handle storage constraints specified as strings in the [juju storage constraint directive format](https://juju.is/docs/juju/storage-constraint).

Unfortunately, this erroneously required that model deployments (using `juju.model.Model.deploy`) also use this string format, instead of the parsed dictionary format that was previously accepted. This was noticed in #1075 and initially fixed in #1105, which was merged into `main` but never released. This fix moved parsing of storage constraint strings to bundle deployment exclusively.

Due to the interim period in which `3.5.2` has (incorrectly) required model deployments to use the string format, I think the best fix at this point is to allow both bundle deployments and model deployments to use either format, and parse them into the parsed dictionary format in a single place in `juju.model.Model._deploy` (the private method called by both bundle and model deployments).

After merging, let's look at getting these changes out in a `3.5.2.2` bugfix release.

Author: jujubot

juju/bundle.py

@@ -1,5 +1,6 @@
 # Copyright 2023 Canonical Ltd.
 # Licensed under the Apache V2, see LICENCE file for details.
+from __future__ import annotations
 
 import base64
 import io
@@ -8,7 +9,7 @@
 import zipfile
 from contextlib import closing
 from pathlib import Path
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Mapping, cast
 
 import requests
 import yaml
@@ -17,12 +18,15 @@
 from . import jasyncio, utils
 from .client import client
 from .constraints import parse as parse_constraints
-from .constraints import parse_storage_constraint
 from .errors import JujuError
 from .origin import Channel, Source
 from .url import URL, Schema
 from .utils import get_base_from_origin_or_channel
 
+if TYPE_CHECKING:
+    from .constraints import StorageConstraintDict
+    from .model import Model
+
 log = logging.getLogger(__name__)
 
 
@@ -599,7 +603,8 @@ class AddApplicationChange(ChangeInfo):
     :options: holds application options.
     :constraints: optional application constraints.
     :storage: optional storage constraints, in the form of `{label: constraint}`.
-        The label is a string specified by the charm, while the constraint is a string following
+        The label is a string specified by the charm, while the constraint is
+        either a constraints.StorageConstraintDict, or a string following
         `the juju storage constraint directive format <https://juju.is/docs/juju/storage-constraint>`_,
         specifying the storage pool, number of volumes, and size of each volume.
     :devices: optional devices constraints.
@@ -610,7 +615,7 @@ class AddApplicationChange(ChangeInfo):
         application's charm.
     """
 
-    storage: Optional[Dict[str, str]]
+    storage: Mapping[str, str | StorageConstraintDict] | None = None
 
     @staticmethod
     def method():
@@ -627,7 +632,8 @@ async def run(self, context):
         # NB: this should really be handled by the controller when generating the
         # bundle change plan, and this short-term workaround may be missing some
         # aspects of the logic which the CLI client contains to handle edge cases.
-        if self.application in context.model.applications:
+        model = cast("Model", context.model)  # pyright: ignore[reportUnknownMemberType]
+        if self.application in model.applications:
             log.debug("Skipping %s; already in model", self.application)
             return self.application
 
@@ -637,9 +643,9 @@ async def run(self, context):
         if self.options is not None:
             options = self.options
         if context.trusted:
-            if context.model.info.agent_version < client.Number.from_json("2.4.0"):
+            if model.info.agent_version < client.Number.from_json("2.4.0"):
                 raise NotImplementedError(
-                    f"trusted is not supported on model version {context.model.info.agent_version}"
+                    f"trusted is not supported on model version {model.info.agent_version}"
                 )
             options["trust"] = "true"
 
@@ -676,22 +682,19 @@ async def run(self, context):
             .get("resources", {})
         )
         if Schema.CHARM_HUB.matches(url.schema):
-            resources = await context.model._add_charmhub_resources(
+            resources = await model._add_charmhub_resources(
                 self.application, charm, origin, overrides=self.resources
             )
 
-        await context.model._deploy(
+        await model._deploy(
             charm_url=charm,
             application=self.application,
             series=self.series,
             config=options,
             constraints=self.constraints,
             endpoint_bindings=self.endpoint_bindings,
             resources=resources,
-            storage={
-                label: parse_storage_constraint(constraint)
-                for label, constraint in (self.storage or {}).items()
-            },
+            storage=self.storage,
             channel=self.channel,
             devices=self.devices,
             num_units=self.num_units,

juju/constraints.py

@@ -19,7 +19,7 @@
 #
 
 import re
-from typing import Dict, List, Optional, TypedDict, Union
+from typing import Dict, List, Mapping, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -190,6 +190,25 @@ def parse_storage_constraint(constraint: str) -> StorageConstraintDict:
     return storage
 
 
+def parse_storage_constraints(
+    constraints: Optional[Mapping[str, Union[str, StorageConstraintDict]]] = None,
+) -> Dict[str, StorageConstraintDict]:
+    if constraints is None:
+        return {}
+    parsed: dict[str, StorageConstraintDict] = {}
+    for label, storage_constraint in constraints.items():
+        if isinstance(storage_constraint, str):
+            parsed[label] = parse_storage_constraint(storage_constraint)
+        elif isinstance(storage_constraint, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+            parsed[label] = storage_constraint
+        else:
+            raise ValueError(
+                f"Unexpected constraint {storage_constraint!r}"
+                f" for label {label!r} in {constraints}"
+            )
+    return parsed
+
+
 DEVICE = re.compile(
     # original regex:
     # '^(?P<count>[0-9]+)?(?:^|,)(?P<type>[^,]+)(?:$|,(?!$))(?P<attrs>(?:[^=]+=[^;]+)+)*$'

juju/model.py

@@ -32,6 +32,7 @@
 from .client import client, connection, connector
 from .client.overrides import Caveat, Macaroon
 from .constraints import parse as parse_constraints
+from .constraints import parse_storage_constraints
 from .controller import ConnectedController, Controller
 from .delta import get_entity_class, get_entity_delta
 from .errors import (
@@ -61,6 +62,7 @@
 if TYPE_CHECKING:
     from .application import Application
     from .client._definitions import FullStatus
+    from .constraints import StorageConstraintDict
     from .machine import Machine
     from .relation import Relation
     from .remoteapplication import ApplicationOffer, RemoteApplication
@@ -1788,7 +1790,7 @@ async def deploy(
         resources=None,
         series=None,
         revision=None,
-        storage=None,
+        storage: Mapping[str, str | StorageConstraintDict] | None = None,
         to=None,
         devices=None,
         trust=False,
@@ -1813,7 +1815,11 @@ async def deploy(
         :param str series: Series on which to deploy DEPRECATED: use --base (with Juju 3.1)
         :param int revision: specifying a revision requires a channel for future upgrades for charms.
             For bundles, revision and channel are mutually exclusive.
-        :param dict storage: Storage constraints TODO how do these look?
+        :param dict storage: optional storage constraints, in the form of `{label: constraint}`.
+            The label is a string specified by the charm, while the constraint is
+            a constraints.StorageConstraintsDict, or a string following
+            `the juju storage constraint directive format <https://juju.is/docs/juju/storage-constraint>`_,
+            specifying the storage pool, number of volumes, and size of each volume.
         :param to: Placement directive as a string. For example:
 
             '23' - place on machine 23
@@ -1829,8 +1835,6 @@ async def deploy(
         :param str[] attach_storage: Existing storage to attach to the deployed unit
             (not available on k8s models)
         """
-        if storage:
-            storage = {k: client.Constraints(**v) for k, v in storage.items()}
         if trust and (self.info.agent_version < client.Number.from_json("2.4.0")):
             raise NotImplementedError(
                 f"trusted is not supported on model version {self.info.agent_version}"
@@ -2251,7 +2255,7 @@ async def _deploy(
         constraints,
         endpoint_bindings,
         resources,
-        storage,
+        storage: Mapping[str, str | StorageConstraintDict] | None,
         channel=None,
         num_units=None,
         placement=None,
@@ -2261,9 +2265,18 @@ async def _deploy(
         force=False,
         server_side_deploy=False,
     ):
-        """Logic shared between `Model.deploy` and `BundleHandler.deploy`."""
+        """Logic shared between `Model.deploy` and `BundleHandler.deploy`.
+
+        :param dict storage: optional storage constraints, in the form of `{label: constraint}`.
+            The label is a string specified by the charm, while the constraint is
+            either a constraints.StorageConstraintDict, or a string following
+            `the juju storage constraint directive format <https://juju.is/docs/juju/storage-constraint>`_,
+            specifying the storage pool, number of volumes, and size of each volume.
+        """
         log.info("Deploying %s", charm_url)
 
+        storage = parse_storage_constraints(storage)
+
         trust = config.get("trust", False)
         # stringify all config values for API, and convert to YAML
         config = {k: str(v) for k, v in config.items()}

tests/integration/test_model.py

@@ -45,6 +45,42 @@ async def test_model_name():
         await model.disconnect()
 
 
+@base.bootstrapped
+async def test_deploy_with_storage_unparsed():
+    async with base.CleanModel() as model:
+        await model.deploy(
+            "postgresql",
+            storage={"pgdata": "1G"},
+        )
+        await model.wait_for_idle(status="active")
+        storages = await model.list_storage()
+        assert len(storages) == 1
+        [storage] = storages
+        # size information isn't exposed, so can't assert on that
+        assert storage["owner-tag"].startswith(tag.unit("postgresql"))
+        assert storage["storage-tag"].startswith(tag.storage("pgdata"))
+        assert storage["life"] == "alive"
+        assert storage["status"].status == "attached"
+
+
+@base.bootstrapped
+async def test_deploy_with_storage_preparsed():
+    async with base.CleanModel() as model:
+        await model.deploy(
+            "postgresql",
+            storage={"pgdata": {"size": 1024, "count": 1}},
+        )
+        await model.wait_for_idle(status="active")
+        storages = await model.list_storage()
+        assert len(storages) == 1
+        [storage] = storages
+        # size information isn't exposed, so can't assert on that
+        assert storage["owner-tag"].startswith(tag.unit("postgresql"))
+        assert storage["storage-tag"].startswith(tag.storage("pgdata"))
+        assert storage["life"] == "alive"
+        assert storage["status"].status == "attached"
+
+
 @base.bootstrapped
 @pytest.mark.bundle
 async def test_deploy_local_bundle_dir():