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

zenml/utils/package_utils.py

@@ -13,10 +13,26 @@
 #  permissions and limitations under the License.
 """Utility functions for the package."""
 
-from typing import List
+import sys
+from typing import Dict, List, Optional, Union, cast
 
 import requests
 from packaging import version
+from packaging.markers import default_environment
+from packaging.requirements import Requirement
+
+if sys.version_info < (3, 10):
+    from importlib_metadata import (
+        PackageNotFoundError,
+        distribution,
+        distributions,
+    )
+else:
+    from importlib.metadata import (
+        PackageNotFoundError,
+        distribution,
+        distributions,
+    )
 
 
 def is_latest_zenml_version() -> bool:
@@ -87,3 +103,97 @@ def clean_requirements(requirements: List[str]) -> List[str]:
         ):
             cleaned[package] = req
     return sorted(cleaned.values())
+
+
+def requirement_installed(requirement: Union[str, Requirement]) -> bool:
+    """Check if a requirement is installed.
+
+    Args:
+        requirement: A requirement string.
+
+    Returns:
+        True if the requirement is installed, False otherwise.
+    """
+    if isinstance(requirement, str):
+        requirement = Requirement(requirement)
+
+    try:
+        dist = distribution(requirement.name)
+    except PackageNotFoundError:
+        return False
+
+    return requirement.specifier.contains(dist.version)
+
+
+def get_dependencies(
+    requirement: Requirement, recursive: bool = False
+) -> List[Requirement]:
+    """Get the dependencies of a requirement.
+
+    Args:
+        requirement: A requirement string.
+        recursive: Whether to include recursive dependencies.
+
+    Returns:
+        A list of requirements.
+    """
+    dist = distribution(requirement.name)
+    marker_environment = cast(Dict[str, str], default_environment())
+
+    dependencies = []
+
+    for req in dist.requires or []:
+        parsed_req = Requirement(req)
+
+        if parsed_req.marker:
+            should_include = False
+
+            marker_environment["extra"] = ""
+            if parsed_req.marker.evaluate(environment=marker_environment):
+                should_include = True
+
+            if not should_include:
+                # Not required without extras, so check if it's required with
+                # any of the requested extras
+                for extra in requirement.extras:
+                    marker_environment["extra"] = extra
+                    if parsed_req.marker.evaluate(
+                        environment=marker_environment
+                    ):
+                        should_include = True
+                        break
+
+            if should_include:
+                dependencies.append(parsed_req)
+        else:
+            # No marker means always include
+            dependencies.append(parsed_req)
+
+    if recursive:
+        for dependency in dependencies:
+            dependencies.extend(get_dependencies(dependency, recursive=True))
+
+    return dependencies
+
+
+def get_package_information(
+    package_names: Optional[List[str]] = None,
+) -> Dict[str, str]:
+    """Get package information.
+
+    Args:
+        package_names: Filter for specific package names. If no package names
+            are provided, all installed packages are returned.
+
+    Returns:
+        A dictionary of the name:version for the package names passed in or
+            all packages and their respective versions.
+    """
+    if package_names:
+        return {
+            dist.name: dist.version
+            for dist in distributions()
+            if dist.name in package_names
+        }
+
+    return {dist.name: dist.version for dist in distributions()}

zenml/zen_server/routers/service_connectors_endpoints.py

@@ -152,6 +152,7 @@ def list_service_connectors(
         resource_type=ResourceType.SERVICE_CONNECTOR,
         list_method=zen_store().list_service_connectors,
         hydrate=hydrate,
+        expand_secrets=expand_secrets,
     )
 
     if expand_secrets:
@@ -163,9 +164,6 @@ def list_service_connectors(
         )
 
         for connector in connectors.items:
-            if not connector.secret_id:
-                continue
-
             if allowed_ids is None or is_owned_by_authenticated_user(
                 connector
             ):
@@ -174,14 +172,8 @@ def list_service_connectors(
                 pass
             elif connector.id not in allowed_ids:
                 # The user is not allowed to read secret values for this
-                # connector. We don't raise an exception here but don't include
-                # the secret values
-                continue
-
-            secret = zen_store().get_secret(secret_id=connector.secret_id)
-
-            # Update the connector configuration with the secret.
-            connector.configuration.update(secret.secret_values)
+                # connector. We remove the secrets from the connector.
+                connector.remove_secrets()
 
     return connectors
 
@@ -253,21 +245,14 @@ def get_service_connector(
         The requested service connector.
     """
     connector = zen_store().get_service_connector(
-        connector_id, hydrate=hydrate
+        connector_id, hydrate=hydrate, expand_secrets=expand_secrets
     )
     verify_permission_for_model(connector, action=Action.READ)
 
-    if (
-        expand_secrets
-        and connector.secret_id
-        and has_permissions_for_model(
-            connector, action=Action.READ_SECRET_VALUE
-        )
+    if expand_secrets and not has_permissions_for_model(
+        connector, action=Action.READ_SECRET_VALUE
     ):
-        secret = zen_store().get_secret(secret_id=connector.secret_id)
-
-        # Update the connector configuration with the secret.
-        connector.configuration.update(secret.secret_values)
+        connector.remove_secrets()
 
     return dehydrate_response_model(connector)
 

zenml/zen_stores/migrations/versions/5bb25e95849c_add_internal_secrets.py

@@ -0,0 +1,62 @@
+"""add internal secrets [5bb25e95849c].
+
+Revision ID: 5bb25e95849c
+Revises: 0.83.1
+Create Date: 2025-06-23 20:49:44.184630
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "5bb25e95849c"
+down_revision = "0.83.1"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # Step 1: Add internal column as nullable
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("internal", sa.Boolean(), nullable=True))
+
+    # Step 2: Update existing records based on service connector references
+    # If a secret is referenced by a service_connector.secret_id, make it internal=True
+    # Otherwise, set it to internal=False
+    connection = op.get_bind()
+
+    # Update secrets that are referenced by service connectors to be internal
+    connection.execute(
+        sa.text("""
+            UPDATE secret 
+            SET internal = TRUE 
+            WHERE id IN (
+                SELECT DISTINCT secret_id 
+                FROM service_connector 
+                WHERE secret_id IS NOT NULL
+            );
+        """)
+    )
+
+    # Update all other secrets to be not internal
+    connection.execute(
+        sa.text("""
+            UPDATE secret 
+            SET internal = FALSE 
+            WHERE internal IS NULL;
+        """)
+    )
+
+    # Step 3: Make internal column non-nullable
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.alter_column(
+            "internal", existing_type=sa.Boolean(), nullable=False
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.drop_column("internal")

zenml/zen_stores/rest_zen_store.py

@@ -2413,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.
 
@@ -2424,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.
@@ -2432,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.
 
@@ -2449,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.
@@ -2457,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(
@@ -2499,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:
@@ -2665,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(
@@ -2711,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

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -20,7 +20,7 @@ from uuid import UUID
 
 from pydantic import ConfigDict
 from sqlalchemy import UniqueConstraint
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import TEXT, Column, Field, Relationship
 
@@ -259,19 +259,19 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         from zenml.zen_stores.schemas import ModelVersionSchema
 
         options = [
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.pipeline)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.stack)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.build)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.schedule)
             ),
-            joinedload(jl_arg(PipelineRunSchema.deployment)).joinedload(
+            selectinload(jl_arg(PipelineRunSchema.deployment)).joinedload(
                 jl_arg(PipelineDeploymentSchema.code_reference)
             ),
         ]
@@ -286,14 +286,14 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         if include_resources:
             options.extend(
                 [
-                    joinedload(
+                    selectinload(
                         jl_arg(PipelineRunSchema.model_version)
                     ).joinedload(
                         jl_arg(ModelVersionSchema.model), innerjoin=True
                     ),
-                    joinedload(jl_arg(PipelineRunSchema.logs)),
-                    joinedload(jl_arg(PipelineRunSchema.user)),
-                    # joinedload(jl_arg(PipelineRunSchema.tags)),
+                    selectinload(jl_arg(PipelineRunSchema.logs)),
+                    selectinload(jl_arg(PipelineRunSchema.user)),
+                    selectinload(jl_arg(PipelineRunSchema.tags)),
                 ]
             )
 

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/schemas/step_run_schemas.py

@@ -21,7 +21,7 @@ from uuid import UUID
 from pydantic import ConfigDict
 from sqlalchemy import TEXT, Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship, SQLModel
 
@@ -50,6 +50,7 @@ from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
     PipelineDeploymentSchema,
+    StepConfigurationSchema,
 )
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.project_schemas import ProjectSchema
@@ -187,6 +188,14 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     original_step_run: Optional["StepRunSchema"] = Relationship(
         sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
     )
+    step_configuration_schema: Optional["StepConfigurationSchema"] = (
+        Relationship(
+            sa_relationship_kwargs=dict(
+                viewonly=True,
+                primaryjoin="and_(foreign(StepConfigurationSchema.name) == StepRunSchema.name, foreign(StepConfigurationSchema.deployment_id) == StepRunSchema.deployment_id)",
+            ),
+        )
+    )
 
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
@@ -209,17 +218,25 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         Returns:
             A list of query options.
         """
-        from zenml.zen_stores.schemas import ModelVersionSchema
+        from zenml.zen_stores.schemas import (
+            ArtifactVersionSchema,
+            ModelVersionSchema,
+        )
 
         options = [
-            joinedload(jl_arg(StepRunSchema.deployment)),
-            joinedload(jl_arg(StepRunSchema.pipeline_run)),
+            selectinload(jl_arg(StepRunSchema.deployment)).load_only(
+                jl_arg(PipelineDeploymentSchema.pipeline_configuration)
+            ),
+            selectinload(jl_arg(StepRunSchema.pipeline_run)).load_only(
+                jl_arg(PipelineRunSchema.start_time)
+            ),
+            joinedload(jl_arg(StepRunSchema.step_configuration_schema)),
         ]
 
         if include_metadata:
             options.extend(
                 [
-                    joinedload(jl_arg(StepRunSchema.logs)),
+                    selectinload(jl_arg(StepRunSchema.logs)),
                     # joinedload(jl_arg(StepRunSchema.parents)),
                     # joinedload(jl_arg(StepRunSchema.run_metadata)),
                 ]
@@ -228,12 +245,28 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         if include_resources:
             options.extend(
                 [
-                    joinedload(jl_arg(StepRunSchema.model_version)).joinedload(
+                    selectinload(
+                        jl_arg(StepRunSchema.model_version)
+                    ).joinedload(
                         jl_arg(ModelVersionSchema.model), innerjoin=True
                     ),
-                    joinedload(jl_arg(StepRunSchema.user)),
-                    # joinedload(jl_arg(StepRunSchema.input_artifacts)),
-                    # joinedload(jl_arg(StepRunSchema.output_artifacts)),
+                    selectinload(jl_arg(StepRunSchema.user)),
+                    selectinload(jl_arg(StepRunSchema.input_artifacts))
+                    .joinedload(
+                        jl_arg(StepRunInputArtifactSchema.artifact_version),
+                        innerjoin=True,
+                    )
+                    .joinedload(
+                        jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
+                    ),
+                    selectinload(jl_arg(StepRunSchema.output_artifacts))
+                    .joinedload(
+                        jl_arg(StepRunOutputArtifactSchema.artifact_version),
+                        innerjoin=True,
+                    )
+                    .joinedload(
+                        jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
+                    ),
                 ]
             )
 
@@ -290,10 +323,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         """
         step = None
         if self.deployment is not None:
-            step_configurations = self.deployment.get_step_configurations(
-                include=[self.name]
-            )
-            if step_configurations:
+            if self.step_configuration_schema:
                 pipeline_configuration = (
                     PipelineConfiguration.model_validate_json(
                         self.deployment.pipeline_configuration
@@ -304,7 +334,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
                     inplace=True,
                 )
                 step = Step.from_dict(
-                    json.loads(step_configurations[0].config),
+                    json.loads(self.step_configuration_schema.config),
                     pipeline_configuration=pipeline_configuration,
                 )
         if not step and self.step_configuration:

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