Skip to content

Commit 3d3f1e5

Browse files
schustmischustmi
committed
Linting
1 parent af9af44 commit 3d3f1e5

File tree

13 files changed

+46
-30
lines changed

13 files changed

+46
-30
lines changed

‎src/zenml/client.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
List,
2929
Mapping,
3030
Optional,
31+
Sequence,
3132
Tuple,
3233
Type,
3334
TypeVar,
@@ -1183,7 +1184,7 @@ def create_stack(
11831184
components: Mapping[StackComponentType, Union[str, UUID]],
11841185
stack_spec_file: Optional[str] = None,
11851186
labels: Optional[Dict[str, Any]] = None,
1186-
secrets: Optional[List[Union[UUID, str]]] = None,
1187+
secrets: Optional[Sequence[Union[UUID, str]]] = None,
11871188
) -> StackResponse:
11881189
"""Registers a stack and its components.
11891190
@@ -1318,8 +1319,8 @@ def update_stack(
13181319
component_updates: Optional[
13191320
Dict[StackComponentType, List[Union[UUID, str]]]
13201321
] = None,
1321-
add_secrets: Optional[List[Union[UUID, str]]] = None,
1322-
remove_secrets: Optional[List[Union[UUID, str]]] = None,
1322+
add_secrets: Optional[Sequence[Union[UUID, str]]] = None,
1323+
remove_secrets: Optional[Sequence[Union[UUID, str]]] = None,
13231324
) -> StackResponse:
13241325
"""Updates a stack and its components.
13251326
@@ -2009,7 +2010,7 @@ def create_stack_component(
20092010
component_type: StackComponentType,
20102011
configuration: Dict[str, str],
20112012
labels: Optional[Dict[str, Any]] = None,
2012-
secrets: Optional[List[Union[UUID, str]]] = None,
2013+
secrets: Optional[Sequence[Union[UUID, str]]] = None,
20132014
) -> "ComponentResponse":
20142015
"""Registers a stack component.
20152016
@@ -2066,8 +2067,8 @@ def update_stack_component(
20662067
disconnect: Optional[bool] = None,
20672068
connector_id: Optional[UUID] = None,
20682069
connector_resource_id: Optional[str] = None,
2069-
add_secrets: Optional[List[Union[UUID, str]]] = None,
2070-
remove_secrets: Optional[List[Union[UUID, str]]] = None,
2070+
add_secrets: Optional[Sequence[Union[UUID, str]]] = None,
2071+
remove_secrets: Optional[Sequence[Union[UUID, str]]] = None,
20712072
) -> ComponentResponse:
20722073
"""Updates a stack component.
20732074

‎src/zenml/config/compiler.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,11 @@
2323
List,
2424
Mapping,
2525
Optional,
26+
Sequence,
2627
Tuple,
28+
Union,
2729
)
30+
from uuid import UUID
2831

2932
from zenml import __version__
3033
from zenml.config.base_settings import BaseSettings, ConfigurationLevel
@@ -466,7 +469,7 @@ def _compile_step_invocation(
466469
self,
467470
invocation: "StepInvocation",
468471
pipeline_environment: Optional[Dict[str, Any]],
469-
pipeline_secrets: List[str],
472+
pipeline_secrets: Sequence[Union[UUID, str]],
470473
pipeline_settings: Dict[str, "BaseSettings"],
471474
pipeline_extra: Dict[str, Any],
472475
stack: "Stack",

‎src/zenml/orchestrators/cache_utils.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,11 @@ def generate_cache_key(
107107
hash_.update(str(value).encode())
108108

109109
# User-defined secrets
110-
for name in sorted(step.config.secrets):
111-
hash_.update(name.encode())
110+
for secret_name_or_id in sorted(step.config.secrets):
111+
if isinstance(secret_name_or_id, UUID):
112+
hash_.update(secret_name_or_id.bytes)
113+
else:
114+
hash_.update(secret_name_or_id.encode())
112115

113116
return hash_.hexdigest()
114117

‎src/zenml/pipelines/pipeline_decorator.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
Union,
2525
overload,
2626
)
27+
from uuid import UUID
2728

2829
from zenml.logger import get_logger
2930

@@ -51,7 +52,7 @@ def pipeline(
5152
enable_artifact_metadata: Optional[bool] = None,
5253
enable_step_logs: Optional[bool] = None,
5354
environment: Optional[Dict[str, Any]] = None,
54-
secrets: Optional[List[str]] = None,
55+
secrets: Optional[List[Union[UUID, str]]] = None,
5556
enable_pipeline_logs: Optional[bool] = None,
5657
settings: Optional[Dict[str, "SettingsOrDict"]] = None,
5758
tags: Optional[List[Union[str, "Tag"]]] = None,
@@ -71,7 +72,7 @@ def pipeline(
7172
enable_artifact_metadata: Optional[bool] = None,
7273
enable_step_logs: Optional[bool] = None,
7374
environment: Optional[Dict[str, Any]] = None,
74-
secrets: Optional[List[str]] = None,
75+
secrets: Optional[List[Union[UUID, str]]] = None,
7576
enable_pipeline_logs: Optional[bool] = None,
7677
settings: Optional[Dict[str, "SettingsOrDict"]] = None,
7778
tags: Optional[List[Union[str, "Tag"]]] = None,

‎src/zenml/pipelines/pipeline_definition.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
List,
2929
Mapping,
3030
Optional,
31+
Sequence,
3132
Set,
3233
Tuple,
3334
TypeVar,
@@ -136,7 +137,7 @@ def __init__(
136137
enable_artifact_visualization: Optional[bool] = None,
137138
enable_step_logs: Optional[bool] = None,
138139
environment: Optional[Dict[str, Any]] = None,
139-
secrets: Optional[List[str]] = None,
140+
secrets: Optional[List[Union[UUID, str]]] = None,
140141
enable_pipeline_logs: Optional[bool] = None,
141142
settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
142143
tags: Optional[List[Union[str, "Tag"]]] = None,
@@ -310,7 +311,7 @@ def configure(
310311
enable_artifact_visualization: Optional[bool] = None,
311312
enable_step_logs: Optional[bool] = None,
312313
environment: Optional[Dict[str, Any]] = None,
313-
secrets: Optional[List[str]] = None,
314+
secrets: Optional[Sequence[Union[UUID, str]]] = None,
314315
enable_pipeline_logs: Optional[bool] = None,
315316
settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
316317
tags: Optional[List[Union[str, "Tag"]]] = None,
@@ -384,7 +385,7 @@ def configure(
384385
tags = self._configuration.tags + tags
385386

386387
if merge and secrets and self._configuration.secrets:
387-
secrets = self._configuration.secrets + secrets
388+
secrets = self._configuration.secrets + list(secrets)
388389

389390
values = dict_utils.remove_none_values(
390391
{

‎src/zenml/stack/stack.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -795,6 +795,7 @@ def validate_image_builder(self) -> None:
795795
user=Client().active_user.id,
796796
created=now,
797797
updated=now,
798+
secrets=[],
798799
)
799800

800801
self._image_builder = image_builder

‎src/zenml/steps/base_step.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -609,7 +609,7 @@ def configure(
609609
"OutputMaterializersSpecification"
610610
] = None,
611611
environment: Optional[Dict[str, Any]] = None,
612-
secrets: Optional[List[Union[str, UUID]]] = None,
612+
secrets: Optional[Sequence[Union[str, UUID]]] = None,
613613
settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
614614
extra: Optional[Dict[str, Any]] = None,
615615
on_failure: Optional["HookSpecification"] = None,
@@ -713,7 +713,7 @@ def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
713713
success_hook_source = resolve_and_validate_hook(on_success)
714714

715715
if merge and secrets and self._configuration.secrets:
716-
secrets = self._configuration.secrets + secrets
716+
secrets = self._configuration.secrets + list(secrets)
717717

718718
values = dict_utils.remove_none_values(
719719
{

‎src/zenml/steps/step_decorator.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
Union,
2828
overload,
2929
)
30+
from uuid import UUID
3031

3132
from zenml.logger import get_logger
3233

@@ -69,7 +70,7 @@ def step(
6970
step_operator: Optional[str] = None,
7071
output_materializers: Optional["OutputMaterializersSpecification"] = None,
7172
environment: Optional[Dict[str, Any]] = None,
72-
secrets: Optional[List[str]] = None,
73+
secrets: Optional[List[Union[UUID, str]]] = None,
7374
settings: Optional[Dict[str, "SettingsOrDict"]] = None,
7475
extra: Optional[Dict[str, Any]] = None,
7576
on_failure: Optional["HookSpecification"] = None,
@@ -92,7 +93,7 @@ def step(
9293
step_operator: Optional[str] = None,
9394
output_materializers: Optional["OutputMaterializersSpecification"] = None,
9495
environment: Optional[Dict[str, Any]] = None,
95-
secrets: Optional[List[str]] = None,
96+
secrets: Optional[List[Union[UUID, str]]] = None,
9697
settings: Optional[Dict[str, "SettingsOrDict"]] = None,
9798
extra: Optional[Dict[str, Any]] = None,
9899
on_failure: Optional["HookSpecification"] = None,

‎src/zenml/utils/secret_utils.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -218,13 +218,15 @@ def resolve_and_verify_secrets(
218218
resolved_secrets.append(secret.id)
219219
else:
220220
filter_model = SecretFilter(name=secret_name_or_id)
221-
secrets = zen_store.list_secrets(filter_model=filter_model)
222-
if not secrets:
221+
secret_page = zen_store.list_secrets(
222+
secret_filter_model=filter_model
223+
)
224+
if not secret_page.items:
223225
raise KeyError(
224226
f"Secret with name {secret_name_or_id} not found."
225227
)
226228

227-
resolved_secrets.append(secrets[0].id)
229+
resolved_secrets.append(secret_page.items[0].id)
228230

229231
return resolved_secrets
230232
else:

‎src/zenml/zen_server/routers/stack_components_endpoints.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from uuid import UUID
1818

1919
from fastapi import APIRouter, Depends, Security
20+
from pydantic import BaseModel
2021

2122
from zenml.constants import API, COMPONENT_TYPES, STACK_COMPONENTS, VERSION_1
2223
from zenml.enums import StackComponentType
@@ -85,7 +86,7 @@ def create_stack_component(
8586
Returns:
8687
The created stack component.
8788
"""
88-
rbac_read_checks = []
89+
rbac_read_checks: List[BaseModel] = []
8990
if component.connector:
9091
service_connector = zen_store().get_service_connector(
9192
component.connector
@@ -216,16 +217,16 @@ def update_stack_component(
216217
validate_custom_flavors=False,
217218
)
218219

219-
rbac_read_checks = []
220+
rbac_read_checks: List[BaseModel] = []
220221
if component_update.connector:
221222
service_connector = zen_store().get_service_connector(
222223
component_update.connector
223224
)
224225
rbac_read_checks.append(service_connector)
225226

226-
if component_update.secrets:
227+
if component_update.add_secrets:
227228
rbac_read_checks.extend(
228-
[zen_store().get_secret(id) for id in component_update.secrets]
229+
[zen_store().get_secret(id) for id in component_update.add_secrets]
229230
)
230231

231232
batch_verify_permissions_for_models(rbac_read_checks, action=Action.READ)

‎src/zenml/zen_server/routers/stacks_endpoints.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,11 @@
1313
# permissions and limitations under the License.
1414
"""Endpoint definitions for stacks."""
1515

16-
from typing import Optional, Union
16+
from typing import List, Optional, Union
1717
from uuid import UUID
1818

1919
from fastapi import APIRouter, Depends, Security
20+
from pydantic import BaseModel
2021

2122
from zenml.constants import API, STACKS, VERSION_1
2223
from zenml.models import (
@@ -82,7 +83,7 @@ def create_stack(
8283
Returns:
8384
The created stack.
8485
"""
85-
rbac_read_checks = []
86+
rbac_read_checks: List[BaseModel] = []
8687

8788
# Check the service connector creation
8889
is_connector_create_needed = False
@@ -212,7 +213,7 @@ def update_stack(
212213
Returns:
213214
The updated stack.
214215
"""
215-
rbac_read_checks = []
216+
rbac_read_checks: List[BaseModel] = []
216217
if stack_update.components:
217218
rbac_read_checks.extend(
218219
[

‎src/zenml/zen_stores/schemas/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@
9797
"RunMetadataSchema",
9898
"ScheduleSchema",
9999
"SecretSchema",
100+
"SecretResourceSchema",
100101
"ServerSettingsSchema",
101102
"ServiceConnectorSchema",
102103
"ServiceSchema",

‎src/zenml/zen_stores/schemas/stack_schemas.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
import base64
1717
import json
18-
from typing import TYPE_CHECKING, Any, List, Optional
18+
from typing import TYPE_CHECKING, Any, List, Optional, Sequence
1919
from uuid import UUID
2020

2121
from sqlalchemy import UniqueConstraint
@@ -123,7 +123,7 @@ class StackSchema(NamedSchema, table=True):
123123
def from_request(
124124
cls,
125125
request: "StackRequest",
126-
components: List["StackComponentSchema"],
126+
components: Sequence["StackComponentSchema"],
127127
) -> "StackSchema":
128128
"""Create a stack schema from a request.
129129

0 commit comments

Comments
 (0)