zenml-nightly 0.83.1.dev20250625__py3-none-any.whl → 0.83.1.dev20250627__py3-none-any.whl

zenml/zen_stores/sql_zen_store.py

@@ -54,7 +54,6 @@ from packaging import version
 from pydantic import (
     ConfigDict,
     Field,
-    SecretStr,
     SerializeAsAny,
     field_validator,
     model_validator,
@@ -65,7 +64,7 @@ from sqlalchemy.exc import (
     ArgumentError,
     IntegrityError,
 )
-from sqlalchemy.orm import Mapped, joinedload, noload
+from sqlalchemy.orm import Mapped, noload, selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlalchemy.util import immutabledict
 from sqlmodel import Session as SqlModelSession
@@ -264,6 +263,7 @@ from zenml.models import (
     ServiceAccountRequest,
     ServiceAccountResponse,
     ServiceAccountUpdate,
+    ServiceConnectorConfiguration,
     ServiceConnectorFilter,
     ServiceConnectorRequest,
     ServiceConnectorResourcesModel,
@@ -3037,18 +3037,17 @@ class SqlZenStore(BaseZenStore):
 
         if artifact is None:
             try:
-                with session.begin_nested():
-                    artifact_request = ArtifactRequest(
-                        name=name,
-                        project=project_id,
-                        has_custom_name=has_custom_name,
-                    )
-                    self._set_request_user_id(
-                        request_model=artifact_request, session=session
-                    )
-                    artifact = ArtifactSchema.from_request(artifact_request)
-                    session.add(artifact)
-                    session.commit()
+                artifact_request = ArtifactRequest(
+                    name=name,
+                    project=project_id,
+                    has_custom_name=has_custom_name,
+                )
+                self._set_request_user_id(
+                    request_model=artifact_request, session=session
+                )
+                artifact = ArtifactSchema.from_request(artifact_request)
+                session.add(artifact)
+                session.commit()
                 session.refresh(artifact)
             except IntegrityError:
                 # We have to rollback the failed session first in order to
@@ -3218,6 +3217,7 @@ class SqlZenStore(BaseZenStore):
                         artifact_version_id=artifact_version_schema.id,
                     )
                     session.add(vis_schema)
+                session.commit()
 
             # Save tags of the artifact
             self._attach_tags_to_resources(
@@ -5332,13 +5332,17 @@ class SqlZenStore(BaseZenStore):
                 schema_class=PipelineRunSchema,
                 session=session,
                 query_options=[
-                    joinedload(jl_arg(PipelineRunSchema.deployment)),
-                    # joinedload(jl_arg(PipelineRunSchema.step_runs)).sele(
-                    #     jl_arg(StepRunSchema.input_artifacts)
-                    # ),
-                    # joinedload(jl_arg(PipelineRunSchema.step_runs)).joinedload(
-                    #     jl_arg(StepRunSchema.output_artifacts)
-                    # ),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.deployment)
+                    ).load_only(
+                        jl_arg(PipelineDeploymentSchema.pipeline_configuration)
+                    ),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.step_runs)
+                    ).selectinload(jl_arg(StepRunSchema.input_artifacts)),
+                    selectinload(
+                        jl_arg(PipelineRunSchema.step_runs)
+                    ).selectinload(jl_arg(StepRunSchema.output_artifacts)),
                 ],
             )
             assert run.deployment is not None
@@ -6743,6 +6747,40 @@ class SqlZenStore(BaseZenStore):
         if self.backup_secrets_store:
             self.backup_secrets_store.delete_secret_values(secret_id=secret_id)
 
+    def _create_secret_schema(
+        self, secret: SecretRequest, session: Session, internal: bool = False
+    ) -> SecretSchema:
+        """Creates a new secret schema.
+
+        Args:
+            secret: The secret to create.
+            session: The session to use.
+            internal: Whether the secret is internal.
+
+        Returns:
+            The newly created secret schema.
+        """
+        new_secret = SecretSchema.from_request(
+            secret,
+            internal=internal,
+        )
+        session.add(new_secret)
+        session.commit()
+
+        try:
+            # Set the secret values in the configured secrets store
+            self._set_secret_values(
+                secret_id=new_secret.id, values=secret.secret_values
+            )
+        except:
+            # If setting the secret values fails, delete the secret from the
+            # database.
+            session.delete(new_secret)
+            session.commit()
+            raise
+
+        return new_secret
+
     @track_decorator(AnalyticsEvent.CREATED_SECRET)
     def create_secret(self, secret: SecretRequest) -> SecretResponse:
         """Creates a new secret.
@@ -6777,31 +6815,17 @@ class SqlZenStore(BaseZenStore):
             if secret_exists:
                 raise EntityExistsError(msg)
 
-            new_secret = SecretSchema.from_request(
-                secret,
+            new_secret = self._create_secret_schema(
+                secret=secret,
+                session=session,
             )
-            session.add(new_secret)
-            session.commit()
 
             secret_model = new_secret.to_model(
                 include_metadata=True, include_resources=True
             )
 
-        try:
-            # Set the secret values in the configured secrets store
-            self._set_secret_values(
-                secret_id=new_secret.id, values=secret.secret_values
-            )
-        except:
-            # If setting the secret values fails, delete the secret from the
-            # database.
-            with Session(self.engine) as session:
-                session.delete(new_secret)
-                session.commit()
-            raise
-
-        secret_model.set_secrets(secret.secret_values)
-        return secret_model
+            secret_model.set_secrets(secret.secret_values)
+            return secret_model
 
     def get_secret(
         self, secret_id: UUID, hydrate: bool = True
@@ -6821,7 +6845,11 @@ class SqlZenStore(BaseZenStore):
         """
         with Session(self.engine) as session:
             secret_in_db = session.exec(
-                select(SecretSchema).where(SecretSchema.id == secret_id)
+                select(SecretSchema).where(
+                    SecretSchema.id == secret_id,
+                    # Don't return internal secrets
+                    col(SecretSchema.internal).is_(False),
+                )
             ).first()
             if (
                 secret_in_db is None
@@ -6870,6 +6898,10 @@ class SqlZenStore(BaseZenStore):
                 self._get_active_user(session).id
             )
             query = select(SecretSchema)
+            # Don't return internal secrets
+            query = query.where(
+                col(SecretSchema.internal).is_(False),
+            )
             return self.filter_and_paginate(
                 session=session,
                 query=query,
@@ -6910,7 +6942,11 @@ class SqlZenStore(BaseZenStore):
         """
         with Session(self.engine) as session:
             existing_secret = session.exec(
-                select(SecretSchema).where(SecretSchema.id == secret_id)
+                select(SecretSchema).where(
+                    SecretSchema.id == secret_id,
+                    # Don't update internal secrets
+                    col(SecretSchema.internal).is_(False),
+                )
             ).first()
 
             active_user = self._get_active_user(session)
@@ -6978,6 +7014,27 @@ class SqlZenStore(BaseZenStore):
 
         return secret_model
 
+    def _delete_secret_schema(self, secret_id: UUID, session: Session) -> None:
+        """Deletes a secret schema.
+
+        Args:
+            secret_id: The ID of the secret to delete.
+            session: The session to use.
+        """
+        # Delete the secret values in the configured secrets store
+        try:
+            self._delete_secret_values(secret_id=secret_id)
+        except KeyError:
+            # If the secret values don't exist in the secrets store, we don't
+            # need to raise an error.
+            pass
+
+        secret_in_db = session.exec(
+            select(SecretSchema).where(SecretSchema.id == secret_id)
+        ).one()
+        session.delete(secret_in_db)
+        session.commit()
+
     def delete_secret(self, secret_id: UUID) -> None:
         """Delete a secret.
 
@@ -6989,7 +7046,11 @@ class SqlZenStore(BaseZenStore):
         """
         with Session(self.engine) as session:
             existing_secret = session.exec(
-                select(SecretSchema).where(SecretSchema.id == secret_id)
+                select(SecretSchema).where(
+                    SecretSchema.id == secret_id,
+                    # Don't delete internal secrets
+                    col(SecretSchema.internal).is_(False),
+                )
             ).first()
 
             if not existing_secret or (
@@ -7003,19 +7064,7 @@ class SqlZenStore(BaseZenStore):
                     "not owned by the current user."
                 )
 
-            # Delete the secret values in the configured secrets store
-            try:
-                self._delete_secret_values(secret_id=secret_id)
-            except KeyError:
-                # If the secret values don't exist in the secrets store, we don't
-                # need to raise an error.
-                pass
-
-            secret_in_db = session.exec(
-                select(SecretSchema).where(SecretSchema.id == secret_id)
-            ).one()
-            session.delete(secret_in_db)
-            session.commit()
+            self._delete_secret_schema(secret_id=secret_id, session=session)
 
     def backup_secrets(
         self, ignore_errors: bool = True, delete_secrets: bool = False
@@ -7379,7 +7428,6 @@ class SqlZenStore(BaseZenStore):
                 resource_types=service_connector.resource_types,
                 resource_id=service_connector.resource_id,
                 configuration=service_connector.configuration,
-                secrets=service_connector.secrets,
             )
 
         with Session(self.engine) as session:
@@ -7397,7 +7445,8 @@ class SqlZenStore(BaseZenStore):
             # Create the secret
             secret_id = self._create_connector_secret(
                 connector_name=service_connector.name,
-                secrets=service_connector.secrets,
+                secrets=service_connector.configuration.secrets,
+                session=session,
             )
             try:
                 # Create the service connector
@@ -7425,12 +7474,20 @@ class SqlZenStore(BaseZenStore):
             connector = new_service_connector.to_model(
                 include_metadata=True, include_resources=True
             )
+            if new_service_connector.secret_id:
+                secrets = self._get_secret_values(
+                    secret_id=new_service_connector.secret_id
+                )
+                connector.add_secrets(secrets)
             self._populate_connector_type(connector)
 
             return connector
 
     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.
 
@@ -7438,6 +7495,8 @@ class SqlZenStore(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 in the
+                output model.
 
         Returns:
             The requested service connector, if it was found.
@@ -7453,12 +7512,20 @@ class SqlZenStore(BaseZenStore):
                 include_metadata=hydrate, include_resources=True
             )
             self._populate_connector_type(connector)
+
+            if expand_secrets and service_connector.secret_id:
+                secrets = self._get_secret_values(
+                    secret_id=service_connector.secret_id
+                )
+                connector.add_secrets(secrets)
+
             return connector
 
     def list_service_connectors(
         self,
         filter_model: ServiceConnectorFilter,
         hydrate: bool = False,
+        expand_secrets: bool = False,
     ) -> Page[ServiceConnectorResponse]:
         """List all service connectors.
 
@@ -7467,6 +7534,8 @@ class SqlZenStore(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 in the
+                output models.
 
         Returns:
             A page of all service connectors.
@@ -7499,6 +7568,27 @@ class SqlZenStore(BaseZenStore):
 
             return items
 
+        def to_model_and_expand_secrets(
+            schema: ServiceConnectorSchema,
+        ) -> ServiceConnectorResponse:
+            """Convert a service connector schema to a model and expand the secrets.
+
+            Args:
+                schema: The service connector schema to convert.
+
+            Returns:
+                The converted service connector model.
+            """
+            model = schema.to_model(
+                include_metadata=hydrate, include_resources=True
+            )
+
+            if expand_secrets and schema.secret_id:
+                secrets = self._get_secret_values(secret_id=schema.secret_id)
+                model.add_secrets(secrets)
+
+            return model
+
         with Session(self.engine) as session:
             query = select(ServiceConnectorSchema)
             paged_connectors: Page[ServiceConnectorResponse] = (
@@ -7508,11 +7598,13 @@ class SqlZenStore(BaseZenStore):
                     table=ServiceConnectorSchema,
                     filter_model=filter_model,
                     custom_fetch=fetch_connectors,
+                    custom_schema_to_model_conversion=to_model_and_expand_secrets,
                     hydrate=hydrate,
                 )
             )
 
             self._populate_connector_type(*paged_connectors.items)
+
             return paged_connectors
 
     def update_service_connector(
@@ -7524,16 +7616,13 @@ class SqlZenStore(BaseZenStore):
         set to None in the model, the field is not updated, but there are
         special rules concerning some fields:
 
-        * the `configuration` and `secrets` fields together represent a full
-        valid configuration update, not just a partial update. If either is
-        set (i.e. not None) in the update, their values are merged together and
-        will replace the existing configuration and secrets values.
+        * the `configuration` field represents a full valid configuration
+        update, not just a partial update. If it is set (i.e. not None) in the
+        update, its values will replace the existing configuration values.
         * the `resource_id` field value is also a full replacement value: if set
         to `None`, the resource ID is removed from the service connector.
         * the `expiration_seconds` field value is also a full replacement value:
         if set to `None`, the expiration is removed from the service connector.
-        * the `secret_id` field value in the update is ignored, given that
-        secrets are managed internally by the ZenML store.
         * the `labels` field is also a full labels update: if set (i.e. not
         `None`), all existing labels are removed and replaced by the new labels
         in the update.
@@ -7631,24 +7720,25 @@ class SqlZenStore(BaseZenStore):
                 update.auth_method = (
                     update.auth_method or existing_connector_model.auth_method
                 )
-                # Validate the configuration update. If the configuration or
-                # secrets fields are set, together they are merged into a
-                # full configuration that is validated against the connector
-                # type schema and replaces the existing configuration and
-                # secrets values
+                # Validate the configuration update. If the configuration is
+                # set, it is validated against the connector type schema and
+                # replaces the existing configuration values
                 update.validate_and_configure_resources(
                     connector_type=connector_type,
                     resource_types=update.resource_types,
                     resource_id=update.resource_id,
                     configuration=update.configuration,
-                    secrets=update.secrets,
                 )
 
-            # Update secret
-            secret_id = self._update_connector_secret(
-                existing_connector=existing_connector_model,
-                updated_connector=update,
-            )
+            secret_id = existing_connector.secret_id
+            if update.configuration is not None:
+                # Update secret
+                secret_id = self._update_connector_secret(
+                    connector_name=existing_connector.name,
+                    existing_secret_id=existing_connector.secret_id,
+                    secrets=update.configuration.secrets,
+                    session=session,
+                )
 
             existing_connector.update(
                 connector_update=update, secret_id=secret_id
@@ -7659,6 +7749,10 @@ class SqlZenStore(BaseZenStore):
             connector = existing_connector.to_model(
                 include_metadata=True, include_resources=True
             )
+            if secret_id:
+                secrets = self._get_secret_values(secret_id=secret_id)
+                connector.add_secrets(secrets)
+
             self._populate_connector_type(connector)
             return connector
 
@@ -7693,7 +7787,10 @@ class SqlZenStore(BaseZenStore):
 
             if service_connector.secret_id:
                 try:
-                    self.delete_secret(service_connector.secret_id)
+                    self._delete_secret_schema(
+                        secret_id=service_connector.secret_id,
+                        session=session,
+                    )
                 except KeyError:
                     # If the secret doesn't exist anymore, we can ignore
                     # this error
@@ -7704,7 +7801,8 @@ class SqlZenStore(BaseZenStore):
     def _create_connector_secret(
         self,
         connector_name: str,
-        secrets: Optional[Dict[str, Optional[SecretStr]]],
+        secrets: Dict[str, PlainSerializedSecretStr],
+        session: Session,
     ) -> Optional[UUID]:
         """Creates a new secret to store the service connector secret credentials.
 
@@ -7712,6 +7810,7 @@ class SqlZenStore(BaseZenStore):
             connector_name: The name of the service connector for which to
                 create a secret.
             secrets: The secret credentials to store.
+            session: The session to use.
 
         Returns:
             The ID of the newly created secret or None, if the service connector
@@ -7729,19 +7828,23 @@ class SqlZenStore(BaseZenStore):
         # that is not already in use
         while True:
             secret_name = f"connector-{connector_name}-{random_str(4)}".lower()
-            existing_secrets = self.list_secrets(
-                SecretFilter(
-                    name=secret_name,
+            existing_secrets = session.exec(
+                select(SecretSchema).where(
+                    SecretSchema.name == secret_name,
                 )
-            )
-            if not existing_secrets.size:
+            ).all()
+            if not existing_secrets:
                 try:
-                    return self.create_secret(
+                    return self._create_secret_schema(
                         SecretRequest(
+                            user=self._get_active_user(session=session).id,
                             name=secret_name,
                             private=False,
                             values=secrets,
-                        )
+                        ),
+                        session=session,
+                        # Hide service connector secrets from the user
+                        internal=True,
                     ).id
                 except KeyError:
                     # The secret already exists, try again
@@ -7811,48 +7914,51 @@ class SqlZenStore(BaseZenStore):
 
     def _update_connector_secret(
         self,
-        existing_connector: ServiceConnectorResponse,
-        updated_connector: ServiceConnectorUpdate,
+        connector_name: str,
+        existing_secret_id: Optional[UUID],
+        secrets: Dict[str, PlainSerializedSecretStr],
+        session: Session,
     ) -> Optional[UUID]:
         """Updates the secret for a service connector.
 
-        If the secrets field in the service connector update is set (i.e. not
-        None), the existing secret, if any, is replaced. If the secrets field is
-        set to an empty dict, the existing secret is deleted.
-
         Args:
-            existing_connector: Existing service connector for which to update a
-                secret.
-            updated_connector: Updated service connector.
+            connector_name: The name of the service connector for which to
+                update a secret.
+            existing_secret_id: The ID of the existing secret to update, if one
+                exists.
+            secrets: The secrets to store.
+            session: The session to use.
 
         Returns:
             The ID of the updated secret or None, if the new service connector
             does not contain any secret credentials.
         """
-        if updated_connector.secrets is None:
-            # If the connector update does not contain a secrets update, keep
-            # the existing secret (if any)
-            return existing_connector.secret_id
-
-        # Delete the existing secret (if any), to be replaced by the new secret
-        if existing_connector.secret_id:
-            try:
-                self.delete_secret(existing_connector.secret_id)
-            except KeyError:
-                # Ignore if the secret no longer exists
-                pass
-
         # If the new service connector does not contain any secret credentials,
         # return None
-        if not updated_connector.secrets:
+        if not secrets:
+            if existing_secret_id:
+                try:
+                    self.delete_secret(existing_secret_id)
+                except KeyError:
+                    # Ignore if the secret no longer exists
+                    pass
             return None
 
-        assert existing_connector.user is not None
-        # A secret does not exist yet, create a new one
-        return self._create_connector_secret(
-            connector_name=updated_connector.name or existing_connector.name,
-            secrets=updated_connector.secrets,
+        if not existing_secret_id:
+            # A secret does not exist yet, create a new one
+            return self._create_connector_secret(
+                connector_name=connector_name,
+                secrets=secrets,
+                session=session,
+            )
+
+        # Update the existing secret - we only need to update the values
+        self._update_secret_values(
+            secret_id=existing_secret_id,
+            values={k: v.get_secret_value() for k, v in secrets.items()},
+            overwrite=True,
         )
+        return existing_secret_id
 
     def verify_service_connector_config(
         self,
@@ -7896,7 +8002,9 @@ class SqlZenStore(BaseZenStore):
             The list of resources that the service connector has access to,
             scoped to the supplied resource type and ID, if provided.
         """
-        connector = self.get_service_connector(service_connector_id)
+        connector = self.get_service_connector(
+            service_connector_id, expand_secrets=True
+        )
 
         connector_instance = service_connector_registry.instantiate_connector(
             model=connector
@@ -7925,7 +8033,9 @@ class SqlZenStore(BaseZenStore):
             A service connector client that can be used to access the given
             resource.
         """
-        connector = self.get_service_connector(service_connector_id)
+        connector = self.get_service_connector(
+            service_connector_id, expand_secrets=True
+        )
 
         connector_instance = service_connector_registry.instantiate_connector(
             model=connector
@@ -8131,7 +8241,9 @@ class SqlZenStore(BaseZenStore):
                     # Fetch an existing service connector
                     if isinstance(connector_id_or_info, UUID):
                         existing_service_connector = (
-                            self.get_service_connector(connector_id_or_info)
+                            self.get_service_connector(
+                                connector_id_or_info, expand_secrets=True
+                            )
                         )
                         if need_to_generate_permanent_tokens:
                             if (
@@ -8140,20 +8252,22 @@ class SqlZenStore(BaseZenStore):
                                 )
                                 is not False
                             ):
-                                connector_config = (
-                                    existing_service_connector.configuration
-                                )
+                                connector_config = existing_service_connector.configuration.plain
                                 connector_config[
                                     "generate_temporary_tokens"
                                 ] = False
                                 self.update_service_connector(
                                     existing_service_connector.id,
                                     ServiceConnectorUpdate(
-                                        configuration=connector_config
+                                        configuration=ServiceConnectorConfiguration(
+                                            **connector_config
+                                        )
                                     ),
                                 )
                         service_connectors.append(
-                            self.get_service_connector(connector_id_or_info)
+                            self.get_service_connector(
+                                connector_id_or_info, expand_secrets=True
+                            )
                         )
                     # Create a new service connector
                     else:
@@ -8169,7 +8283,9 @@ class SqlZenStore(BaseZenStore):
                                     name=connector_name,
                                     connector_type=connector_id_or_info.type,
                                     auth_method=connector_id_or_info.auth_method,
-                                    configuration=connector_config,
+                                    configuration=ServiceConnectorConfiguration(
+                                        **connector_config
+                                    ),
                                     labels={
                                         k: str(v)
                                         for k, v in stack.labels.items()