Fix snapshot exclusive tagging (#4034)

* Fix exclusive tags for snapshots

* Simplify adding/removing a single tag

Author: schustmi

src/zenml/utils/tag_utils.py

@@ -127,38 +127,38 @@ def to_request(self) -> "TagRequest":
 
 @overload
 def add_tags(
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     run: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     artifact: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     artifact_version_id: UUID,
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     artifact_name: str,
     artifact_version: str,
 ) -> None: ...
@@ -167,7 +167,7 @@ def add_tags(
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     infer_artifact: bool = False,
     artifact_name: Optional[str] = None,
 ) -> None: ...
@@ -176,37 +176,37 @@ def add_tags(
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     pipeline: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     run_template: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     snapshot: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def add_tags(
     *,
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     deployment: Union[UUID, str],
 ) -> None: ...
 
 
 def add_tags(
-    tags: List[Union[str, Tag]],
+    tags: Union[str, Tag, List[Union[str, Tag]]],
     # Pipelines
     pipeline: Optional[Union[UUID, str]] = None,
     # Runs
@@ -253,6 +253,9 @@ def add_tags(
     resource_id = None
     resource_type = None
 
+    if isinstance(tags, (str, Tag)):
+        tags = [tags]
+
     # Tag a pipeline
     if pipeline is not None and all(
         v is None
@@ -509,7 +512,6 @@ def add_tags(
             """
         )
 
-    # Create tag resources and add tags
     if resource_id:
         for tag in tags:
             client.attach_tag(
@@ -520,70 +522,70 @@ def add_tags(
 
 @overload
 def remove_tags(
-    tags: List[str],
+    tags: Union[str, List[str]],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     pipeline: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     run: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     run_template: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     snapshot: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     deployment: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     artifact: Union[UUID, str],
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     artifact_version_id: UUID,
 ) -> None: ...
 
 
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     artifact_name: str,
     artifact_version: str,
 ) -> None: ...
@@ -592,14 +594,14 @@ def remove_tags(
 @overload
 def remove_tags(
     *,
-    tags: List[str],
+    tags: Union[str, List[str]],
     infer_artifact: bool = False,
     artifact_name: Optional[str] = None,
 ) -> None: ...
 
 
 def remove_tags(
-    tags: List[str],
+    tags: Union[str, List[str]],
     # Pipelines
     pipeline: Optional[Union[UUID, str]] = None,
     # Runs
@@ -645,6 +647,9 @@ def remove_tags(
     resource_id = None
     resource_type = None
 
+    if isinstance(tags, str):
+        tags = [tags]
+
     # Remove tags from a pipeline
     if pipeline is not None and all(
         v is None

src/zenml/zen_stores/sql_zen_store.py

@@ -13841,7 +13841,7 @@ def _create_tag_resource_schemas(
                             PipelineSnapshotFilter(
                                 id=f"notequals:{resource.id}",
                                 project=resource.project.id,
-                                pipeline_id=scope_id,
+                                pipeline=scope_id,
                                 tags=[tag_schema.name],
                             )
                         )
@@ -13884,9 +13884,8 @@ def _create_tag_resource_schemas(
                             "Can not attach exclusive tag to resource of type "
                             f"{resource_type.value} with ID: `{resource.id}`. "
                             "Exclusive tag functionality only works for "
-                            "templates, for pipeline runs (within the scope of "
-                            "pipelines) and for artifact versions (within the "
-                            "scope of artifacts)."
+                            "pipeline runs, pipeline snapshots, pipeline "
+                            "deployments and artifact versions"
                         )
 
                     # Check for duplicate IDs in any of the scope_ids list