Skip to content

Commit fa8d86d

Browse files
authored
Filter component by user name (#3126)
* Filter component by user name * Include full name in user filters * Formatting * Handle none value
1 parent f9980f5 commit fa8d86d

12 files changed

+72
-9
lines changed

src/zenml/client.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1906,6 +1906,7 @@ def list_stack_components(
19061906
user_id: Optional[Union[str, UUID]] = None,
19071907
connector_id: Optional[Union[str, UUID]] = None,
19081908
stack_id: Optional[Union[str, UUID]] = None,
1909+
user: Optional[Union[UUID, str]] = None,
19091910
hydrate: bool = False,
19101911
) -> Page[ComponentResponse]:
19111912
"""Lists all registered stack components.
@@ -1925,6 +1926,7 @@ def list_stack_components(
19251926
connector_id: The id of the connector to filter by.
19261927
stack_id: The id of the stack to filter by.
19271928
name: The name of the component to filter by.
1929+
user: The ID of name of the user to filter by.
19281930
hydrate: Flag deciding whether to hydrate the output model(s)
19291931
by including metadata fields in the response.
19301932
@@ -1946,6 +1948,7 @@ def list_stack_components(
19461948
id=id,
19471949
created=created,
19481950
updated=updated,
1951+
user=user,
19491952
)
19501953
component_filter_model.set_scope_workspace(self.active_workspace.id)
19511954

src/zenml/models/v2/base/filter.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -604,12 +604,15 @@ def generate_name_or_id_query_conditions(
604604
self,
605605
value: Union[UUID, str],
606606
table: Type["NamedSchema"],
607+
additional_columns: Optional[List[str]] = None,
607608
) -> "ColumnElement[bool]":
608609
"""Generate filter conditions for name or id of a table.
609610
610611
Args:
611612
value: The filter value.
612613
table: The table to filter.
614+
additional_columns: Additional table columns that should also
615+
filtered for the given value as part of the or condition.
613616
614617
Returns:
615618
The query conditions.
@@ -637,6 +640,12 @@ def generate_name_or_id_query_conditions(
637640
)
638641
conditions.append(filter_.generate_query_conditions(table=table))
639642

643+
for column in additional_columns or []:
644+
filter_ = FilterGenerator(table).define_filter(
645+
column=column, value=value, operator=operator
646+
)
647+
conditions.append(filter_.generate_query_conditions(table=table))
648+
640649
return or_(*conditions)
641650

642651
def generate_custom_query_conditions_for_column(

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -556,7 +556,9 @@ def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
556556
user_filter = and_(
557557
ArtifactVersionSchema.user_id == UserSchema.id,
558558
self.generate_name_or_id_query_conditions(
559-
value=self.user, table=UserSchema
559+
value=self.user,
560+
table=UserSchema,
561+
additional_columns=["full_name"],
560562
),
561563
)
562564
custom_filters.append(user_filter)

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,7 @@ class ComponentFilter(WorkspaceScopedFilter):
376376
*WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
377377
"scope_type",
378378
"stack_id",
379+
"user",
379380
]
380381
CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
381382
*WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -418,6 +419,10 @@ class ComponentFilter(WorkspaceScopedFilter):
418419
description="Stack of the stack component",
419420
union_mode="left_to_right",
420421
)
422+
user: Optional[Union[UUID, str]] = Field(
423+
default=None,
424+
description="Name/ID of the user that created the component.",
425+
)
421426

422427
def set_scope_type(self, component_type: str) -> None:
423428
"""Set the type of component on which to perform the filtering to scope the response.
@@ -464,3 +469,31 @@ def generate_filter(
464469
base_filter = operator(base_filter, stack_filter)
465470

466471
return base_filter
472+
473+
def get_custom_filters(self) -> List["ColumnElement[bool]"]:
474+
"""Get custom filters.
475+
476+
Returns:
477+
A list of custom filters.
478+
"""
479+
from sqlmodel import and_
480+
481+
from zenml.zen_stores.schemas import (
482+
StackComponentSchema,
483+
UserSchema,
484+
)
485+
486+
custom_filters = super().get_custom_filters()
487+
488+
if self.user:
489+
user_filter = and_(
490+
StackComponentSchema.user_id == UserSchema.id,
491+
self.generate_name_or_id_query_conditions(
492+
value=self.user,
493+
table=UserSchema,
494+
additional_columns=["full_name"],
495+
),
496+
)
497+
custom_filters.append(user_filter)
498+
499+
return custom_filters

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,9 @@ def get_custom_filters(
368368
user_filter = and_(
369369
ModelSchema.user_id == UserSchema.id,
370370
self.generate_name_or_id_query_conditions(
371-
value=self.user, table=UserSchema
371+
value=self.user,
372+
table=UserSchema,
373+
additional_columns=["full_name"],
372374
),
373375
)
374376
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -660,7 +660,9 @@ def get_custom_filters(
660660
user_filter = and_(
661661
ModelVersionSchema.user_id == UserSchema.id,
662662
self.generate_name_or_id_query_conditions(
663-
value=self.user, table=UserSchema
663+
value=self.user,
664+
table=UserSchema,
665+
additional_columns=["full_name"],
664666
),
665667
)
666668
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,9 @@ def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
294294
== ArtifactVersionSchema.id,
295295
ArtifactVersionSchema.user_id == UserSchema.id,
296296
self.generate_name_or_id_query_conditions(
297-
value=self.user, table=UserSchema
297+
value=self.user,
298+
table=UserSchema,
299+
additional_columns=["full_name"],
298300
),
299301
)
300302
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,9 @@ def get_custom_filters(self) -> List["ColumnElement[bool]"]:
218218
== PipelineRunSchema.id,
219219
PipelineRunSchema.user_id == UserSchema.id,
220220
self.generate_name_or_id_query_conditions(
221-
value=self.user, table=UserSchema
221+
value=self.user,
222+
table=UserSchema,
223+
additional_columns=["full_name"],
222224
),
223225
)
224226
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,9 @@ def get_custom_filters(
364364
user_filter = and_(
365365
PipelineSchema.user_id == UserSchema.id,
366366
self.generate_name_or_id_query_conditions(
367-
value=self.user, table=UserSchema
367+
value=self.user,
368+
table=UserSchema,
369+
additional_columns=["full_name"],
368370
),
369371
)
370372
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -779,7 +779,9 @@ def get_custom_filters(
779779
user_filter = and_(
780780
PipelineRunSchema.user_id == UserSchema.id,
781781
self.generate_name_or_id_query_conditions(
782-
value=self.user, table=UserSchema
782+
value=self.user,
783+
table=UserSchema,
784+
additional_columns=["full_name"],
783785
),
784786
)
785787
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,9 @@ def get_custom_filters(
413413
user_filter = and_(
414414
RunTemplateSchema.user_id == UserSchema.id,
415415
self.generate_name_or_id_query_conditions(
416-
value=self.user, table=UserSchema
416+
value=self.user,
417+
table=UserSchema,
418+
additional_columns=["full_name"],
417419
),
418420
)
419421
custom_filters.append(user_filter)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -381,7 +381,9 @@ def get_custom_filters(self) -> List["ColumnElement[bool]"]:
381381
user_filter = and_(
382382
StackSchema.user_id == UserSchema.id,
383383
self.generate_name_or_id_query_conditions(
384-
value=self.user, table=UserSchema
384+
value=self.user,
385+
table=UserSchema,
386+
additional_columns=["full_name"],
385387
),
386388
)
387389
custom_filters.append(user_filter)

0 commit comments

Comments
 (0)