feat: unify handling of storage constraints in deploy

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():