zenml-nightly 0.83.1.dev20250624__py3-none-any.whl → 0.83.1.dev20250626__py3-none-any.whl

zenml/zen_stores/rest_zen_store.py

@@ -18,6 +18,7 @@ import re
 import time
 from datetime import datetime
 from pathlib import Path
+from threading import RLock
 from typing import (
     Any,
     ClassVar,
@@ -39,6 +40,7 @@ from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    PrivateAttr,
     ValidationError,
     field_validator,
     model_validator,
@@ -272,6 +274,7 @@ from zenml.utils.networking_utils import (
     replace_localhost_with_internal_hostname,
 )
 from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.exceptions import exception_from_response
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
@@ -440,6 +443,8 @@ class RestZenStore(BaseZenStore):
     _api_token: Optional[APIToken] = None
     _session: Optional[requests.Session] = None
     _server_info: Optional[ServerModel] = None
+    _session_lock: RLock = PrivateAttr(default_factory=RLock)
+    _last_authenticated: Optional[datetime] = None
 
     # ====================================
     # ZenML Store interface implementation
@@ -2408,10 +2413,21 @@ class RestZenStore(BaseZenStore):
             response_model=ServiceConnectorResponse,
         )
         self._populate_connector_type(connector_model)
+        # Call this to properly split the secrets from the configuration
+        try:
+            connector_model.validate_configuration()
+        except ValueError as e:
+            logger.error(
+                f"Error validating connector configuration for "
+                f"{connector_model.name}: {e}"
+            )
         return connector_model
 
     def get_service_connector(
-        self, service_connector_id: UUID, hydrate: bool = True
+        self,
+        service_connector_id: UUID,
+        hydrate: bool = True,
+        expand_secrets: bool = False,
     ) -> ServiceConnectorResponse:
         """Gets a specific service connector.
 
@@ -2419,6 +2435,8 @@ class RestZenStore(BaseZenStore):
             service_connector_id: The ID of the service connector to get.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            expand_secrets: Flag deciding whether to include the secrets
+                associated with the service connector.
 
         Returns:
             The requested service connector, if it was found.
@@ -2427,15 +2445,25 @@ class RestZenStore(BaseZenStore):
             resource_id=service_connector_id,
             route=SERVICE_CONNECTORS,
             response_model=ServiceConnectorResponse,
-            params={"expand_secrets": False, "hydrate": hydrate},
+            params={"hydrate": hydrate, "expand_secrets": expand_secrets},
         )
         self._populate_connector_type(connector_model)
+        if expand_secrets:
+            try:
+                # Call this to properly split the secrets from the configuration
+                connector_model.validate_configuration()
+            except ValueError as e:
+                logger.error(
+                    f"Error validating connector configuration for "
+                    f"{connector_model.name}: {e}"
+                )
         return connector_model
 
     def list_service_connectors(
         self,
         filter_model: ServiceConnectorFilter,
         hydrate: bool = False,
+        expand_secrets: bool = False,
     ) -> Page[ServiceConnectorResponse]:
         """List all service connectors.
 
@@ -2444,6 +2472,8 @@ class RestZenStore(BaseZenStore):
                 params.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
+            expand_secrets: Flag deciding whether to include the secrets
+                associated with the service connector.
 
         Returns:
             A page of all service connectors.
@@ -2452,9 +2482,19 @@ class RestZenStore(BaseZenStore):
             route=SERVICE_CONNECTORS,
             response_model=ServiceConnectorResponse,
             filter_model=filter_model,
-            params={"expand_secrets": False, "hydrate": hydrate},
+            params={"hydrate": hydrate, "expand_secrets": expand_secrets},
         )
         self._populate_connector_type(*connector_models.items)
+        if expand_secrets:
+            # Call this to properly split the secrets from the configuration
+            for connector_model in connector_models.items:
+                try:
+                    connector_model.validate_configuration()
+                except ValueError as e:
+                    logger.error(
+                        f"Error validating connector configuration for "
+                        f"{connector_model.name}: {e}"
+                    )
         return connector_models
 
     def update_service_connector(
@@ -2494,6 +2534,14 @@ class RestZenStore(BaseZenStore):
             route=SERVICE_CONNECTORS,
         )
         self._populate_connector_type(connector_model)
+        # Call this to properly split the secrets from the configuration
+        try:
+            connector_model.validate_configuration()
+        except ValueError as e:
+            logger.error(
+                f"Error validating connector configuration for "
+                f"{connector_model.name}: {e}"
+            )
         return connector_model
 
     def delete_service_connector(self, service_connector_id: UUID) -> None:
@@ -2660,6 +2708,14 @@ class RestZenStore(BaseZenStore):
 
         connector = ServiceConnectorResponse.model_validate(response_body)
         self._populate_connector_type(connector)
+        # Call this to properly split the secrets from the configuration
+        try:
+            connector.validate_configuration()
+        except ValueError as e:
+            logger.error(
+                f"Error validating connector configuration for connector client "
+                f"{connector.name}: {e}"
+            )
         return connector
 
     def list_service_connector_resources(
@@ -2706,7 +2762,9 @@ class RestZenStore(BaseZenStore):
 
             # Retrieve the resource list locally
             assert resources.id is not None
-            connector = self.get_service_connector(resources.id)
+            connector = self.get_service_connector(
+                resources.id, expand_secrets=True
+            )
             connector_instance = (
                 service_connector_registry.instantiate_connector(
                     model=connector
@@ -4157,7 +4215,7 @@ class RestZenStore(BaseZenStore):
                 elif token.expired:
                     raise CredentialsNotValid(
                         "Your authentication to the current server has expired. "
-                        "Please log in again using 'zenml login --url "
+                        "Please log in again using 'zenml login "
                         f"{self.url}'."
                     )
 
@@ -4203,74 +4261,79 @@ class RestZenStore(BaseZenStore):
         Returns:
             A requests session.
         """
-        if self._session is None:
-            # We only need to initialize the session once over the lifetime
-            # of the client. We can swap the token out when it expires.
-            if self.config.verify_ssl is False:
-                urllib3.disable_warnings(
-                    urllib3.exceptions.InsecureRequestWarning
-                )
+        with self._session_lock:
+            if self._session is None:
+                # We only need to initialize the session once over the lifetime
+                # of the client. We can swap the token out when it expires.
+                if self.config.verify_ssl is False:
+                    urllib3.disable_warnings(
+                        urllib3.exceptions.InsecureRequestWarning
+                    )
 
-            self._session = requests.Session()
-            # Retries are triggered for all HTTP methods (GET, HEAD, POST, PUT,
-            # PATCH, OPTIONS and DELETE) on specific HTTP status codes:
-            #
-            #     408: Request Timeout.
-            #     429: Too Many Requests.
-            #     502: Bad Gateway.
-            #     503: Service Unavailable.
-            #     504: Gateway Timeout
-            #
-            # This also handles connection level errors, if a connection attempt
-            # fails due to transient issues like:
-            #
-            #     DNS resolution errors.
-            #     Connection timeouts.
-            #     Network disruptions.
-            #
-            # Additional errors retried:
-            #
-            #     Read Timeouts: If the server does not send a response within
-            #     the timeout period.
-            #     Connection Refused: If the server refuses the connection.
-            #
-            retries = Retry(
-                connect=5,
-                read=8,
-                redirect=3,
-                status=10,
-                allowed_methods=[
-                    "HEAD",
-                    "GET",
-                    "POST",
-                    "PUT",
-                    "PATCH",
-                    "DELETE",
-                    "OPTIONS",
-                ],
-                status_forcelist=[
-                    408,  # Request Timeout
-                    429,  # Too Many Requests
-                    502,  # Bad Gateway
-                    503,  # Service Unavailable
-                    504,  # Gateway Timeout
-                ],
-                other=3,
-                backoff_factor=1,
-            )
-            self._session.mount("https://", HTTPAdapter(max_retries=retries))
-            self._session.mount("http://", HTTPAdapter(max_retries=retries))
-            self._session.verify = self.config.verify_ssl
-            # Use a custom user agent to identify the ZenML client in the server
-            # logs.
-            self._session.headers.update(
-                {"User-Agent": "zenml/" + zenml.__version__}
-            )
+                self._session = requests.Session()
+                # Retries are triggered for all HTTP methods (GET, HEAD, POST, PUT,
+                # PATCH, OPTIONS and DELETE) on specific HTTP status codes:
+                #
+                #     408: Request Timeout.
+                #     429: Too Many Requests.
+                #     502: Bad Gateway.
+                #     503: Service Unavailable.
+                #     504: Gateway Timeout
+                #
+                # This also handles connection level errors, if a connection attempt
+                # fails due to transient issues like:
+                #
+                #     DNS resolution errors.
+                #     Connection timeouts.
+                #     Network disruptions.
+                #
+                # Additional errors retried:
+                #
+                #     Read Timeouts: If the server does not send a response within
+                #     the timeout period.
+                #     Connection Refused: If the server refuses the connection.
+                #
+                retries = Retry(
+                    connect=5,
+                    read=8,
+                    redirect=3,
+                    status=10,
+                    allowed_methods=[
+                        "HEAD",
+                        "GET",
+                        "POST",
+                        "PUT",
+                        "PATCH",
+                        "DELETE",
+                        "OPTIONS",
+                    ],
+                    status_forcelist=[
+                        408,  # Request Timeout
+                        429,  # Too Many Requests
+                        502,  # Bad Gateway
+                        503,  # Service Unavailable
+                        504,  # Gateway Timeout
+                    ],
+                    other=3,
+                    backoff_factor=1,
+                )
+                self._session.mount(
+                    "https://", HTTPAdapter(max_retries=retries)
+                )
+                self._session.mount(
+                    "http://", HTTPAdapter(max_retries=retries)
+                )
+                self._session.verify = self.config.verify_ssl
+                # Use a custom user agent to identify the ZenML client in the server
+                # logs.
+                self._session.headers.update(
+                    {"User-Agent": "zenml/" + zenml.__version__}
+                )
 
-        # Note that we return an unauthenticated session here. An API token
-        # is only fetched and set in the authorization header when and if it is
-        # needed.
-        return self._session
+            # Note that we return an unauthenticated session here. An API token
+            # is only fetched and set in the authorization header when and if it is
+            # needed.
+            return self._session
 
     def authenticate(self, force: bool = False) -> None:
         """Authenticate or re-authenticate to the ZenML server.
@@ -4325,6 +4388,7 @@ class RestZenStore(BaseZenStore):
             {"Authorization": "Bearer " + new_api_token}
         )
         logger.debug(f"Authenticated to {self.url}")
+        self._last_authenticated = utc_now()
 
     @staticmethod
     def _handle_response(response: requests.Response) -> Json:
@@ -4391,9 +4455,9 @@ class RestZenStore(BaseZenStore):
             CredentialsNotValid: if the request fails due to invalid
                 client credentials.
         """
-        self.session.headers.update(
-            {source_context.name: source_context.get().value}
-        )
+        request_headers = {
+            source_context.name: source_context.get().value,
+        }
 
         # If the server replies with a credentials validation (401 Unauthorized)
         # error, we (re-)authenticate and retry the request here in the
@@ -4414,19 +4478,21 @@ class RestZenStore(BaseZenStore):
         while True:
             # Add a request ID to the request headers
             request_id = str(uuid4())[:8]
-            self.session.headers.update({"X-Request-ID": request_id})
+            request_headers["X-Request-ID"] = request_id
             # Add an idempotency key to the request headers to ensure that
             # requests are idempotent.
-            self.session.headers.update({"Idempotency-Key": str(uuid4())})
+            request_headers["Idempotency-Key"] = str(uuid4())
 
             start_time = time.time()
             logger.debug(f"[{request_id}] {method} {path} started...")
             status_code = "failed"
+            last_authenticated = self._last_authenticated
 
             try:
                 response = self.session.request(
                     method,
                     url,
+                    headers=request_headers,
                     params=params if params else {},
                     verify=self.config.verify_ssl,
                     timeout=timeout or self.config.http_timeout,
@@ -4442,62 +4508,68 @@ class RestZenStore(BaseZenStore):
                 # authenticated at all.
                 credentials_store = get_credentials_store()
 
-                if self._api_token is None:
-                    # The last request was not authenticated with an API
-                    # token at all. We authenticate here and then try the
-                    # request again, this time with a valid API token in the
-                    # header.
-                    logger.debug(
-                        f"[{request_id}] The last request was not "
-                        f"authenticated: {e}\n"
-                        "Re-authenticating and retrying..."
-                    )
-                    self.authenticate()
-                elif not credentials_store.can_login(self.url):
-                    # The request failed either because we're not
-                    # authenticated or our current credentials are not valid
-                    # anymore.
-                    logger.error(
-                        "The current token is no longer valid, and "
-                        "it is not possible to generate a new token using the "
-                        "configured credentials. Please run "
-                        f"`zenml login {self.url}` to "
-                        "re-authenticate to the server or authenticate using "
-                        "an API key. See "
-                        "https://docs.zenml.io/deploying-zenml/connecting-to-zenml/connect-with-a-service-account "
-                        "for more information."
-                    )
-                    # Clear the current token from the credentials store to
-                    # force a new authentication flow next time.
-                    get_credentials_store().clear_token(self.url)
-                    raise e
-                elif not re_authenticated:
-                    # The last request was authenticated with an API token
-                    # that was rejected by the server. We attempt a
-                    # re-authentication here and then retry the request.
-                    logger.debug(
-                        f"[{request_id}] The last request was authenticated "
-                        "with an API token that was rejected by the server: "
-                        f"{e}\n"
-                        "Re-authenticating and retrying..."
-                    )
-                    re_authenticated = True
-                    self.authenticate(
-                        # Ignore the current token and force a re-authentication
-                        force=True
-                    )
-                else:
-                    # The last request was made after re-authenticating but
-                    # still failed. Bailing out.
-                    logger.debug(
-                        f"[{request_id}] The last request failed after "
-                        "re-authenticating: {e}\n"
-                        "Bailing out..."
-                    )
-                    raise CredentialsNotValid(
-                        "The current credentials are no longer valid. Please "
-                        "log in again using 'zenml login'."
-                    ) from e
+                with self._session_lock:
+                    if self._last_authenticated != last_authenticated:
+                        # Another thread has re-authenticated since the last
+                        # request. We simply retry the request.
+                        continue
+
+                    if self._api_token is None:
+                        # The last request was not authenticated with an API
+                        # token at all. We authenticate here and then try the
+                        # request again, this time with a valid API token in the
+                        # header.
+                        logger.debug(
+                            f"[{request_id}] The last request was not "
+                            f"authenticated: {e}\n"
+                            "Re-authenticating and retrying..."
+                        )
+                        self.authenticate()
+                    elif not credentials_store.can_login(self.url):
+                        # The request failed either because we're not
+                        # authenticated or our current credentials are not valid
+                        # anymore.
+                        logger.error(
+                            "The current token is no longer valid, and "
+                            "it is not possible to generate a new token using the "
+                            "configured credentials. Please run "
+                            f"`zenml login {self.url}` to "
+                            "re-authenticate to the server or authenticate using "
+                            "an API key. See "
+                            "https://docs.zenml.io/deploying-zenml/connecting-to-zenml/connect-with-a-service-account "
+                            "for more information."
+                        )
+                        # Clear the current token from the credentials store to
+                        # force a new authentication flow next time.
+                        get_credentials_store().clear_token(self.url)
+                        raise e
+                    elif not re_authenticated:
+                        # The last request was authenticated with an API token
+                        # that was rejected by the server. We attempt a
+                        # re-authentication here and then retry the request.
+                        logger.debug(
+                            f"[{request_id}] The last request was authenticated "
+                            "with an API token that was rejected by the server: "
+                            f"{e}\n"
+                            "Re-authenticating and retrying..."
+                        )
+                        re_authenticated = True
+                        self.authenticate(
+                            # Ignore the current token and force a re-authentication
+                            force=True
+                        )
+                    else:
+                        # The last request was made after re-authenticating but
+                        # still failed. Bailing out.
+                        logger.debug(
+                            f"[{request_id}] The last request failed after "
+                            "re-authenticating: {e}\n"
+                            "Bailing out..."
+                        )
+                        raise CredentialsNotValid(
+                            "The current credentials are no longer valid. Please "
+                            "log in again using 'zenml login'."
+                        ) from e
             finally:
                 end_time = time.time()
                 duration = (end_time - start_time) * 1000

zenml/zen_stores/schemas/secret_schemas.py

@@ -67,6 +67,8 @@ class SecretSchema(NamedSchema, table=True):
 
     private: bool
 
+    internal: bool = Field(default=False)
+
     values: Optional[bytes] = Field(sa_column=Column(TEXT, nullable=True))
 
     user_id: UUID = build_foreign_key_field(
@@ -191,11 +193,13 @@ class SecretSchema(NamedSchema, table=True):
     def from_request(
         cls,
         secret: SecretRequest,
+        internal: bool = False,
     ) -> "SecretSchema":
         """Create a `SecretSchema` from a `SecretRequest`.
 
         Args:
             secret: The `SecretRequest` from which to create the schema.
+            internal: Whether the secret is internal.
 
         Returns:
             The created `SecretSchema`.
@@ -209,6 +213,7 @@ class SecretSchema(NamedSchema, table=True):
             # SQL secret store will call `store_secret_values` to store the
             # values separately if SQL is used as the secrets store.
             values=None,
+            internal=internal,
         )
 
     def update(

zenml/zen_stores/schemas/service_connector_schemas.py

@@ -25,6 +25,7 @@ from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
+    ServiceConnectorConfiguration,
     ServiceConnectorRequest,
     ServiceConnectorResponse,
     ServiceConnectorResponseBody,
@@ -168,6 +169,7 @@ class ServiceConnectorSchema(NamedSchema, table=True):
             The created `ServiceConnectorSchema`.
         """
         assert connector_request.user is not None, "User must be set."
+        configuration = connector_request.configuration.non_secrets
         return cls(
             user_id=connector_request.user,
             name=connector_request.name,
@@ -180,9 +182,9 @@ class ServiceConnectorSchema(NamedSchema, table=True):
             resource_id=connector_request.resource_id,
             supports_instances=connector_request.supports_instances,
             configuration=base64.b64encode(
-                json.dumps(connector_request.configuration).encode("utf-8")
+                json.dumps(configuration).encode("utf-8")
             )
-            if connector_request.configuration
+            if configuration
             else None,
             secret_id=secret_id,
             expires_at=connector_request.expires_at,
@@ -226,15 +228,16 @@ class ServiceConnectorSchema(NamedSchema, table=True):
                     self.expiration_seconds = None
                 continue
             if field == "configuration":
-                self.configuration = (
-                    base64.b64encode(
-                        json.dumps(connector_update.configuration).encode(
-                            "utf-8"
+                if connector_update.configuration is not None:
+                    configuration = connector_update.configuration.non_secrets
+                    if configuration is not None:
+                        self.configuration = (
+                            base64.b64encode(
+                                json.dumps(configuration).encode("utf-8")
+                            )
+                            if configuration
+                            else None
                         )
-                    )
-                    if connector_update.configuration
-                    else None
-                )
             elif field == "resource_types":
                 self.resource_types = base64.b64encode(
                     json.dumps(connector_update.resource_types).encode("utf-8")
@@ -286,12 +289,11 @@ class ServiceConnectorSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = ServiceConnectorResponseMetadata(
-                configuration=json.loads(
-                    base64.b64decode(self.configuration).decode()
+                configuration=ServiceConnectorConfiguration(
+                    **json.loads(base64.b64decode(self.configuration).decode())
                 )
                 if self.configuration
-                else {},
-                secret_id=self.secret_id,
+                else ServiceConnectorConfiguration(),
                 expiration_seconds=self.expiration_seconds,
                 labels=self.labels_dict,
             )

zenml/zen_stores/secrets_stores/service_connector_secrets_store.py

@@ -29,6 +29,7 @@ from pydantic import Field, model_validator
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.logger import get_logger
 from zenml.models import (
+    ServiceConnectorConfiguration,
     ServiceConnectorRequest,
 )
 from zenml.service_connectors.service_connector import ServiceConnector
@@ -133,7 +134,9 @@ class ServiceConnectorSecretsStore(BaseSecretsStore):
                 connector_type=self.SERVICE_CONNECTOR_TYPE,
                 resource_types=[self.SERVICE_CONNECTOR_RESOURCE_TYPE],
                 auth_method=self.config.auth_method,
-                configuration=self.config.auth_config,
+                configuration=ServiceConnectorConfiguration(
+                    **self.config.auth_config
+                ),
             )
             base_connector = service_connector_registry.instantiate_connector(
                 model=request