new: Add support for LKE node pool labels & taints (#448)

* WIP

* fix populate errors

* Finish unit tests

* Add integration test

* Update linode_api4/objects/lke.py

Co-authored-by: Zhiwei Liang <[email protected]>

---------

Co-authored-by: Zhiwei Liang <[email protected]>

linode_api4/groups/lke.py

@@ -3,13 +3,13 @@
 from linode_api4.errors import UnexpectedResponseError
 from linode_api4.groups import Group
 from linode_api4.objects import (
-    Base,
-    JSONObject,
     KubeVersion,
     LKECluster,
     LKEClusterControlPlaneOptions,
+    Type,
     drop_null_keys,
 )
+from linode_api4.objects.base import _flatten_request_body_recursive
 
 
 class LKEGroup(Group):
@@ -107,41 +107,22 @@ def cluster_create(
         :returns: The new LKE Cluster
         :rtype: LKECluster
         """
-        pools = []
-        if not isinstance(node_pools, list):
-            node_pools = [node_pools]
-
-        for c in node_pools:
-            if isinstance(c, dict):
-                new_pool = {
-                    "type": (
-                        c["type"].id
-                        if "type" in c and issubclass(type(c["type"]), Base)
-                        else c.get("type")
-                    ),
-                    "count": c.get("count"),
-                }
-
-                pools += [new_pool]
 
         params = {
             "label": label,
-            "region": region.id if issubclass(type(region), Base) else region,
-            "node_pools": pools,
-            "k8s_version": (
-                kube_version.id
-                if issubclass(type(kube_version), Base)
-                else kube_version
-            ),
-            "control_plane": (
-                control_plane.dict
-                if issubclass(type(control_plane), JSONObject)
-                else control_plane
+            "region": region,
+            "k8s_version": kube_version,
+            "node_pools": (
+                node_pools if isinstance(node_pools, list) else [node_pools]
             ),
+            "control_plane": control_plane,
         }
         params.update(kwargs)
 
-        result = self.client.post("/lke/clusters", data=drop_null_keys(params))
+        result = self.client.post(
+            "/lke/clusters",
+            data=_flatten_request_body_recursive(drop_null_keys(params)),
+        )
 
         if "id" not in result:
             raise UnexpectedResponseError(
@@ -150,7 +131,7 @@ def cluster_create(
 
         return LKECluster(self.client, result["id"], result)
 
-    def node_pool(self, node_type, node_count):
+    def node_pool(self, node_type: Union[Type, str], node_count: int, **kwargs):
         """
         Returns a dict that is suitable for passing into the `node_pools` array
         of :any:`cluster_create`.  This is a convenience method, and need not be
@@ -160,11 +141,17 @@ def node_pool(self, node_type, node_count):
         :type node_type: Type or str
         :param node_count: The number of nodes to create in this node pool.
         :type node_count: int
+        :param kwargs: Other attributes to create this node pool with.
+        :type kwargs: Any
 
         :returns: A dict describing the desired node pool.
         :rtype: dict
         """
-        return {
+        result = {
             "type": node_type,
             "count": node_count,
         }
+
+        result.update(kwargs)
+
+        return result

linode_api4/objects/base.py

@@ -277,6 +277,8 @@ def save(self, force=True) -> bool:
                 ):
                     data[key] = None
 
+            # Ensure we serialize any values that may not be already serialized
+            data = _flatten_request_body_recursive(data)
         else:
             data = self._serialize()
 
@@ -343,10 +345,7 @@ def _serialize(self):
 
         # Resolve the underlying IDs of results
         for k, v in result.items():
-            if isinstance(v, Base):
-                result[k] = v.id
-            elif isinstance(v, MappedObject) or issubclass(type(v), JSONObject):
-                result[k] = v.dict
+            result[k] = _flatten_request_body_recursive(v)
 
         return result
 
@@ -502,3 +501,29 @@ def make_instance(cls, id, client, parent_id=None, json=None):
         :returns: A new instance of this type, populated with json
         """
         return Base.make(id, client, cls, parent_id=parent_id, json=json)
+
+
+def _flatten_request_body_recursive(data: Any) -> Any:
+    """
+    This is a helper recursively flatten the given data for use in an API request body.
+
+    NOTE: This helper does NOT raise an error if an attribute is
+    not known to be JSON serializable.
+
+    :param data: Arbitrary data to flatten.
+    :return: The serialized data.
+    """
+
+    if isinstance(data, dict):
+        return {k: _flatten_request_body_recursive(v) for k, v in data.items()}
+
+    if isinstance(data, list):
+        return [_flatten_request_body_recursive(v) for v in data]
+
+    if isinstance(data, Base):
+        return data.id
+
+    if isinstance(data, MappedObject) or issubclass(type(data), JSONObject):
+        return data.dict
+
+    return data

linode_api4/objects/lke.py

@@ -29,6 +29,18 @@ class KubeVersion(Base):
     }
 
 
+@dataclass
+class LKENodePoolTaint(JSONObject):
+    """
+    LKENodePoolTaint represents the structure of a single taint that can be
+    applied to a node pool.
+    """
+
+    key: Optional[str] = None
+    value: Optional[str] = None
+    effect: Optional[str] = None
+
+
 @dataclass
 class LKEClusterControlPlaneACLAddressesOptions(JSONObject):
     """
@@ -139,37 +151,51 @@ class LKENodePool(DerivedBase):
         ),  # this is formatted in _populate below
         "autoscaler": Property(mutable=True),
         "tags": Property(mutable=True, unordered=True),
+        "labels": Property(mutable=True),
+        "taints": Property(mutable=True),
     }
 
+    def _parse_raw_node(
+        self, raw_node: Union[LKENodePoolNode, dict, str]
+    ) -> LKENodePoolNode:
+        """
+        Builds a list of LKENodePoolNode objects given a node pool response's JSON.
+        """
+        if isinstance(raw_node, LKENodePoolNode):
+            return raw_node
+
+        if isinstance(raw_node, dict):
+            node_id = raw_node.get("id")
+            if node_id is None:
+                raise ValueError("Node dictionary does not contain 'id' key")
+
+            return LKENodePoolNode(self._client, raw_node)
+
+        if isinstance(raw_node, str):
+            return self._client.load(
+                LKENodePoolNode, target_id=raw_node, target_parent_id=self.id
+            )
+
+        raise TypeError("Unsupported node type: {}".format(type(raw_node)))
+
     def _populate(self, json):
         """
         Parse Nodes into more useful LKENodePoolNode objects
         """
+
         if json is not None and json != {}:
-            new_nodes = []
-            for c in json["nodes"]:
-                if isinstance(c, LKENodePoolNode):
-                    new_nodes.append(c)
-                elif isinstance(c, dict):
-                    node_id = c.get("id")
-                    if node_id is not None:
-                        new_nodes.append(LKENodePoolNode(self._client, c))
-                    else:
-                        raise ValueError(
-                            "Node dictionary does not contain 'id' key"
-                        )
-                elif isinstance(c, str):
-                    node_details = self._client.get(
-                        LKENodePool.api_endpoint.format(
-                            cluster_id=self.id, id=c
-                        )
-                    )
-                    new_nodes.append(
-                        LKENodePoolNode(self._client, node_details)
-                    )
-                else:
-                    raise TypeError("Unsupported node type: {}".format(type(c)))
-            json["nodes"] = new_nodes
+            json["nodes"] = [
+                self._parse_raw_node(node) for node in json.get("nodes", [])
+            ]
+
+            json["taints"] = [
+                (
+                    LKENodePoolTaint.from_json(taint)
+                    if not isinstance(taint, LKENodePoolTaint)
+                    else taint
+                )
+                for taint in json.get("taints", [])
+            ]
 
         super()._populate(json)
 
@@ -302,7 +328,14 @@ def control_plane_acl(self) -> LKEClusterControlPlaneACL:
 
         return LKEClusterControlPlaneACL.from_json(self._control_plane_acl)
 
-    def node_pool_create(self, node_type, node_count, **kwargs):
+    def node_pool_create(
+        self,
+        node_type: Union[Type, str],
+        node_count: int,
+        labels: Dict[str, str] = None,
+        taints: List[Union[LKENodePoolTaint, Dict[str, Any]]] = None,
+        **kwargs,
+    ):
         """
         Creates a new :any:`LKENodePool` for this cluster.
 
@@ -312,6 +345,10 @@ def node_pool_create(self, node_type, node_count, **kwargs):
         :type node_type: :any:`Type` or str
         :param node_count: The number of nodes to create in this pool.
         :type node_count: int
+        :param labels: A dict mapping labels to their values to apply to this pool.
+        :type labels: Dict[str, str]
+        :param taints: A list of taints to apply to this pool.
+        :type taints: List of :any:`LKENodePoolTaint` or dict
         :param kwargs: Any other arguments to pass to the API.  See the API docs
                        for possible values.
 
@@ -322,6 +359,13 @@ def node_pool_create(self, node_type, node_count, **kwargs):
             "type": node_type,
             "count": node_count,
         }
+
+        if labels is not None:
+            params["labels"] = labels
+
+        if taints is not None:
+            params["taints"] = taints
+
         params.update(kwargs)
 
         result = self._client.post(

test/fixtures/lke_clusters_18881_pools_456.json

@@ -23,6 +23,17 @@
       "example tag",
       "another example"
     ],
+    "taints": [
+      {
+        "key": "foo",
+        "value": "bar",
+        "effect": "NoSchedule"
+      }
+    ],
+    "labels": {
+      "foo": "bar",
+      "bar": "foo"
+    },
     "type": "g6-standard-4",
     "disk_encryption": "enabled"
   }