Merge branch 'main' into upgrade-jupyterhub-helm-chart

Author: marcelovilla

src/_nebari/stages/infrastructure/__init__.py

@@ -49,18 +49,12 @@ class NodeGroup(schema.Base):
     def validate_taint_strings(cls, taints: list[Any]):
         if taints is None:
             return taints
-
-        TAINT_STR_REGEX = re.compile(r"(\w+)=(\w+):(\w+)")
         return_value = []
         for taint in taints:
             if not isinstance(taint, str):
                 return_value.append(taint)
             else:
-                match = TAINT_STR_REGEX.match(taint)
-                if not match:
-                    raise ValueError(f"Invalid taint string: {taint}")
-                key, taints, effect = match.groups()
-                parsed_taint = schema.Taint(key=key, value=taints, effect=effect)
+                parsed_taint = schema.Taint.from_string(taint)
                 return_value.append(parsed_taint)
 
         return return_value
@@ -73,7 +67,6 @@ def validate_taint_strings(cls, taints: list[Any]):
 
 
 def set_missing_taints_to_default_taints(node_groups: NodeGroup) -> NodeGroup:
-
     for node_group_name, node_group in node_groups.items():
         if node_group.taints is None:
             if node_group_name == "general":
@@ -239,21 +232,6 @@ def construct_aws_ami_type(
 
     return "AL2_x86_64"
 
-    @field_validator("node_taints", mode="before")
-    def convert_taints(cls, value: Optional[List[schema.Taint]]):
-        return [
-            dict(
-                key=taint.key,
-                value=taint.value,
-                effect={
-                    schema.TaintEffectEnum.NoSchedule: "NO_SCHEDULE",
-                    schema.TaintEffectEnum.PreferNoSchedule: "PREFER_NO_SCHEDULE",
-                    schema.TaintEffectEnum.NoExecute: "NO_EXECUTE",
-                }[taint.effect],
-            )
-            for taint in value
-        ]
-
 
 class AWSInputVars(schema.Base):
     name: str
@@ -656,7 +634,7 @@ def _check_input(cls, data: Any) -> Any:
                 or available_kms_keys[key_id[0]].Arn != data["eks_kms_arn"]
             ):
                 raise ValueError(
-                    f"Amazon Web Services KMS Key with ARN {data['eks_kms_arn']} not one of available/enabled keys={ [v.Arn for v in available_kms_keys.values() if v.KeyManager == 'CUSTOMER' and v.KeySpec == 'SYMMETRIC_DEFAULT']}"
+                    f"Amazon Web Services KMS Key with ARN {data['eks_kms_arn']} not one of available/enabled keys={[v.Arn for v in available_kms_keys.values() if v.KeyManager == 'CUSTOMER' and v.KeySpec == 'SYMMETRIC_DEFAULT']}"
                 )
             key_id = key_id[0]
             # Raise error if key is not a customer managed key

src/nebari/schema.py

@@ -1,5 +1,6 @@
 import enum
-from typing import Annotated
+import re
+from typing import Annotated, ClassVar
 
 import pydantic
 from pydantic import ConfigDict, Field, StringConstraints, field_validator
@@ -117,10 +118,32 @@ def to_yaml(cls, representer, node):
         return representer.represent_str(node.value)
 
 
+TAINT_KEY_REGEX = r"([a-zA-Z0-9][-a-zA-Z0-9_.]{0,251}[a-zA-Z0-9]?(?:/[a-zA-Z0-9][-a-zA-Z0-9_.]{0,251}[a-zA-Z0-9]?)?)"
+TAINT_VALUE_REGEX = r"([a-zA-Z0-9][-a-zA-Z0-9_.]{0,61}[a-zA-Z0-9]?)?"
+TAINT_STR_REGEX = rf"^{TAINT_KEY_REGEX}(?:={TAINT_VALUE_REGEX})?:({'|'.join([taint_effect.value for taint_effect in TaintEffectEnum])})$"
+
+
 class Taint(Base):
-    key: str
-    value: str
+    # Taint constraints listed at https://kubernetes.io/docs/reference/generated/kubectl/kubectl-commands#taint
+    key: Annotated[str, StringConstraints(pattern=TAINT_KEY_REGEX)]
+    value: Annotated[None | str, StringConstraints(pattern=TAINT_VALUE_REGEX)]
     effect: TaintEffectEnum
+    _TAINT_REGEX_PATTERN: ClassVar[re.Pattern] = re.compile(TAINT_STR_REGEX)
+
+    @classmethod
+    def from_string(cls, taint_string: str):
+        """
+        Given a taint string, parse it into the Taint object.
+        """
+        match = cls._TAINT_REGEX_PATTERN.match(taint_string)
+        if not match:
+            raise ValueError(f"Invalid taint string: {taint_string}")
+        key, value, effect = match.groups()
+        return Taint(
+            key=key,
+            value=value,
+            effect=effect,
+        )
 
 
 provider_enum_name_map: dict[ProviderEnum, str] = {

tests/tests_unit/test_stages.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from _nebari.stages.infrastructure import NodeGroup
 from _nebari.stages.terraform_state import TerraformStateStage
 from _nebari.utils import yaml
 from _nebari.version import __version__
@@ -115,3 +116,32 @@ def test_check_immutable_fields_change_dict_any(
 
     # This should not raise an exception
     terraform_state_stage.check_immutable_fields()
+
+
+@pytest.mark.parametrize(
+    "taints,keys,values,effects",
+    [
+        (["key1=value1:NoSchedule"], ["key1"], ["value1"], ["NoSchedule"]),
+        (["key1=value1:PreferNoSchedule"], ["key1"], ["value1"], ["PreferNoSchedule"]),
+        (["key1=value1:NoExecute"], ["key1"], ["value1"], ["NoExecute"]),
+        (
+            ["dedicated=special-user:NoSchedule"],
+            ["dedicated"],
+            ["special-user"],
+            ["NoSchedule"],
+        ),
+        (["dedicated:NoSchedule"], ["dedicated"], [None], ["NoSchedule"]),
+        (
+            ["key1=value1:NoExecute", "key2=value2:NoExecute"],
+            ["key1", "key2"],
+            ["value1", "value2"],
+            ["NoExecute", "NoExecute"],
+        ),
+    ],
+)
+def test_node_group_taints(taints, keys, values, effects):
+    ng = NodeGroup(instance="t3.medium", min_nodes=1, max_nodes=1, taints=taints)
+
+    assert [t.key for t in ng.taints] == keys
+    assert [t.value for t in ng.taints] == values
+    assert [t.effect for t in ng.taints] == effects