Missing CLI methods

Author: schustmi

src/zenml/cli/stack.py

@@ -676,6 +676,22 @@ def register_stack(
     type=str,
     required=False,
 )
+@click.option(
+    "--secret",
+    "secrets",
+    help="Secrets to attach to the stack.",
+    type=str,
+    required=False,
+    multiple=True,
+)
+@click.option(
+    "--remove-secret",
+    "remove_secrets",
+    help="Secrets to remove from the stack.",
+    type=str,
+    required=False,
+    multiple=True,
+)
 def update_stack(
     stack_name_or_id: Optional[str] = None,
     artifact_store: Optional[str] = None,
@@ -690,6 +706,8 @@ def update_stack(
     data_validator: Optional[str] = None,
     image_builder: Optional[str] = None,
     model_registry: Optional[str] = None,
+    secrets: List[str] = [],
+    remove_secrets: List[str] = [],
 ) -> None:
     """Update a stack.
 
@@ -708,6 +726,8 @@ def update_stack(
         data_validator: Name of the new data validator for this stack.
         image_builder: Name of the new image builder for this stack.
         model_registry: Name of the new model registry for this stack.
+        secrets: Secrets to attach to the stack.
+        remove_secrets: Secrets to remove from the stack.
     """
     client = Client()
 
@@ -746,6 +766,8 @@ def update_stack(
             updated_stack = client.update_stack(
                 name_id_or_prefix=stack_name_or_id,
                 component_updates=updates,
+                add_secrets=secrets,
+                remove_secrets=remove_secrets,
             )
 
         except (KeyError, IllegalOperationError) as err:

src/zenml/cli/stack_components.py

@@ -232,6 +232,14 @@ def generate_stack_component_register_command(
         required=False,
         type=str,
     )
+    @click.option(
+        "--secret",
+        "secrets",
+        help="Secrets to attach to the component.",
+        type=str,
+        required=False,
+        multiple=True,
+    )
     @click.argument("args", nargs=-1, type=click.UNPROCESSED)
     def register_stack_component_command(
         name: str,
@@ -240,6 +248,7 @@ def register_stack_component_command(
         labels: Optional[List[str]] = None,
         connector: Optional[str] = None,
         resource_id: Optional[str] = None,
+        secrets: List[str] = [],
     ) -> None:
         """Registers a stack component.
 
@@ -250,6 +259,7 @@ def register_stack_component_command(
             labels: Labels to be associated with the component.
             connector: Name of the service connector to connect the component to.
             resource_id: The resource ID to use with the connector.
+            secrets: Secrets to attach to the component.
         """
         client = Client()
 
@@ -277,6 +287,7 @@ def register_stack_component_command(
                 component_type=component_type,
                 configuration=parsed_args,
                 labels=parsed_labels,
+                secrets=secrets,
             )
 
             cli_utils.declare(
@@ -323,18 +334,38 @@ def generate_stack_component_update_command(
         "-l key1=value1 -l key2=value2.",
         multiple=True,
     )
+    @click.option(
+        "--secret",
+        "secrets",
+        help="Secrets to attach to the component.",
+        type=str,
+        required=False,
+        multiple=True,
+    )
+    @click.option(
+        "--remove-secret",
+        "remove_secrets",
+        help="Secrets to remove from the component.",
+        type=str,
+        required=False,
+        multiple=True,
+    )
     @click.argument("args", nargs=-1, type=click.UNPROCESSED)
     def update_stack_component_command(
         name_id_or_prefix: Optional[str],
         args: List[str],
         labels: Optional[List[str]] = None,
+        secrets: List[str] = [],
+        remove_secrets: List[str] = [],
     ) -> None:
         """Updates a stack component.
 
         Args:
             name_id_or_prefix: The name or id of the stack component to update.
             args: Additional arguments to pass to the update command.
             labels: Labels to be associated with the component.
+            secrets: Secrets to attach to the component.
+            remove_secrets: Secrets to remove from the component.
         """
         client = Client()
 
@@ -358,6 +389,8 @@ def update_stack_component_command(
                     component_type=component_type,
                     configuration=parsed_args,
                     labels=parsed_labels,
+                    add_secrets=secrets,
+                    remove_secrets=remove_secrets,
                 )
             except KeyError as err:
                 cli_utils.error(str(err))

src/zenml/models/v2/core/component.py

@@ -151,11 +151,6 @@ class ComponentUpdate(BaseUpdate):
         default=None,
         title="Environment variables to set when running on this component.",
     )
-    secrets: Optional[List[UUID]] = Field(
-        default=None,
-        title="Secrets to set as environment variables when running on this "
-        "component.",
-    )
     connector_resource_id: Optional[str] = Field(
         description="The ID of a specific resource instance to "
         "gain access to through the connector",

src/zenml/models/v2/core/stack.py

@@ -171,11 +171,6 @@ class StackUpdate(BaseUpdate):
         default=None,
         title="Environment variables to set when running on this stack.",
     )
-    secrets: Optional[List[UUID]] = Field(
-        default=None,
-        title="Secrets to set as environment variables when running on this "
-        "stack.",
-    )
     labels: Optional[Dict[str, Any]] = Field(
         default=None,
         title="The stack labels.",