content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """SQL Zen Store implementation."""
  
  from zenml.models.v2.core.step_run import StepHeartbeatResponse
  
  try:
      import sqlalchemy  # noqa
  except ImportError:
      raise ImportError(
          "It seems like you've installed the `zenml` package without the "
          "`local` extra, but are trying to use ZenML with a local database.\n"
          "* If you want to use ZenML in a local setup, please install "
          "`zenml[local]` instead, e.g. using `pip install 'zenml[local]'`\n"
          "* If you want to connect to a server, run `zenml login`"
      ) from None
  
  import base64
  import inspect
  import json
  import logging
  import math
  import os
  import random
  import re
  import sys
  import time
  from collections import defaultdict
  from datetime import datetime, timedelta, timezone
  from functools import lru_cache
  from pathlib import Path
  from typing import (
      TYPE_CHECKING,
      Any,
      Callable,
      ClassVar,
      Dict,
      ForwardRef,
      List,
      Literal,
      NoReturn,
      Optional,
      Sequence,
      Set,
      Tuple,
      Type,
      TypeVar,
      Union,
      cast,
      get_origin,
      overload,
  )
  from uuid import UUID
  
  from packaging import version
  from pydantic import (
      ConfigDict,
      Field,
      SerializeAsAny,
      field_validator,
      model_validator,
  )
  from sqlalchemy import QueuePool, event, func, update
  from sqlalchemy.engine import URL, Engine, make_url
  from sqlalchemy.exc import (
      ArgumentError,
      IntegrityError,
  )
  from sqlalchemy.orm import (
      Mapped,
      load_only,
      noload,
      selectinload,
  )
  from sqlalchemy.sql.base import ExecutableOption
  from sqlalchemy.util import immutabledict
  from sqlmodel import Session as SqlModelSession
  
  # Important to note: The select function of SQLModel works slightly differently
  # from the select function of sqlalchemy. If you input only one entity on the
  # select function of SQLModel, it automatically maps it to a SelectOfScalar.
  # As a result, it will not return a tuple as a result, but the first entity in
  # the tuple. While this is convenient in most cases, in unique cases like using
  # the "add_columns" functionality, one might encounter unexpected results.
  from sqlmodel import (
      SQLModel,
      and_,
      col,
      create_engine,
      delete,
      desc,
      or_,
      select,
  )
  from sqlmodel.sql.expression import Select, SelectOfScalar
  
  from zenml.analytics import track
  from zenml.analytics.context import AnalyticsContext
  from zenml.analytics.enums import AnalyticsEvent
  from zenml.analytics.utils import (
      analytics_disabler,
      track_decorator,
      track_handler,
  )
  from zenml.config.global_config import GlobalConfiguration
  from zenml.config.pipeline_configurations import PipelineConfiguration
  from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
  from zenml.config.secrets_store_config import SecretsStoreConfiguration
  from zenml.config.server_config import ServerConfiguration
  from zenml.config.source import Source
  from zenml.config.step_configurations import Step, StepConfiguration, StepSpec
  from zenml.config.store_config import StoreConfiguration
  from zenml.constants import (
      DEFAULT_PASSWORD,
      DEFAULT_STACK_AND_COMPONENT_NAME,
      DEFAULT_USERNAME,
      ENV_ZENML_DEFAULT_USER_NAME,
      ENV_ZENML_DEFAULT_USER_PASSWORD,
      ENV_ZENML_DISABLE_DATABASE_MIGRATION,
      ENV_ZENML_SERVER,
      FINISHED_ONBOARDING_SURVEY_KEY,
      MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION,
      SQL_STORE_BACKUP_DIRECTORY_NAME,
      TEXT_FIELD_MAX_LENGTH,
      handle_bool_env_var,
      is_false_string_value,
      is_true_string_value,
  )
  from zenml.enums import (
      ArtifactSaveType,
      AuthScheme,
      DatabaseBackupStrategy,
      DeploymentStatus,
      ExecutionMode,
      ExecutionStatus,
      LoggingLevels,
      MetadataResourceTypes,
      ModelStages,
      OnboardingStep,
      SecretResourceTypes,
      SecretsStoreType,
      StackComponentType,
      StackDeploymentProvider,
      StepRunInputArtifactType,
      StoreType,
      TaggableResourceTypes,
      VisualizationResourceTypes,
  )
  from zenml.exceptions import (
      AuthorizationException,
      BackupSecretsStoreNotConfiguredError,
      EntityCreationError,
      EntityExistsError,
      IllegalOperationError,
      SecretsStoreNotConfiguredError,
  )
  from zenml.io import fileio
  from zenml.logger import get_console_handler, get_logger, get_logging_level
  from zenml.metadata.metadata_types import get_metadata_type
  from zenml.models import (
      ActionFilter,
      ActionRequest,
      ActionResponse,
      ActionUpdate,
      APIKeyFilter,
      APIKeyInternalResponse,
      APIKeyInternalUpdate,
      APIKeyRequest,
      APIKeyResponse,
      APIKeyRotateRequest,
      APIKeyUpdate,
      ApiTransactionRequest,
      ApiTransactionResponse,
      ApiTransactionUpdate,
      ArtifactFilter,
      ArtifactRequest,
      ArtifactResponse,
      ArtifactUpdate,
      ArtifactVersionFilter,
      ArtifactVersionRequest,
      ArtifactVersionResponse,
      ArtifactVersionUpdate,
      ArtifactVisualizationResponse,
      BaseFilter,
      BaseIdentifiedResponse,
      BaseRequest,
      BaseResponse,
      BaseUpdate,
      CodeReferenceRequest,
      CodeReferenceResponse,
      CodeRepositoryFilter,
      CodeRepositoryRequest,
      CodeRepositoryResponse,
      CodeRepositoryUpdate,
      ComponentFilter,
      ComponentRequest,
      ComponentResponse,
      ComponentUpdate,
      CuratedVisualizationRequest,
      CuratedVisualizationResponse,
      CuratedVisualizationUpdate,
      DefaultComponentRequest,
      DefaultStackRequest,
      DeployedStack,
      DeploymentFilter,
      DeploymentRequest,
      DeploymentResponse,
      DeploymentUpdate,
      EventSourceFilter,
      EventSourceRequest,
      EventSourceResponse,
      EventSourceUpdate,
      FlavorFilter,
      FlavorRequest,
      FlavorResponse,
      FlavorUpdate,
      LogsResponse,
      ModelFilter,
      ModelRequest,
      ModelResponse,
      ModelUpdate,
      ModelVersionArtifactFilter,
      ModelVersionArtifactRequest,
      ModelVersionArtifactResponse,
      ModelVersionFilter,
      ModelVersionPipelineRunFilter,
      ModelVersionPipelineRunRequest,
      ModelVersionPipelineRunResponse,
      ModelVersionRequest,
      ModelVersionResponse,
      ModelVersionUpdate,
      OAuthDeviceFilter,
      OAuthDeviceInternalRequest,
      OAuthDeviceInternalResponse,
      OAuthDeviceInternalUpdate,
      OAuthDeviceResponse,
      OAuthDeviceUpdate,
      Page,
      PipelineBuildFilter,
      PipelineBuildRequest,
      PipelineBuildResponse,
      PipelineFilter,
      PipelineRequest,
      PipelineResponse,
      PipelineRunDAG,
      PipelineRunFilter,
      PipelineRunRequest,
      PipelineRunResponse,
      PipelineRunUpdate,
      PipelineSnapshotFilter,
      PipelineSnapshotRequest,
      PipelineSnapshotResponse,
      PipelineSnapshotRunRequest,
      PipelineSnapshotUpdate,
      PipelineUpdate,
      ProjectFilter,
      ProjectRequest,
      ProjectResponse,
      ProjectScopedFilter,
      ProjectScopedRequest,
      ProjectUpdate,
      RunMetadataRequest,
      RunMetadataResource,
      RunTemplateFilter,
      RunTemplateRequest,
      RunTemplateResponse,
      RunTemplateUpdate,
      ScheduleFilter,
      ScheduleRequest,
      ScheduleResponse,
      ScheduleUpdate,
      SecretFilter,
      SecretRequest,
      SecretResponse,
      SecretUpdate,
      ServerActivationRequest,
      ServerDatabaseType,
      ServerDeploymentType,
      ServerModel,
      ServerSettingsResponse,
      ServerSettingsUpdate,
      ServiceAccountFilter,
      ServiceAccountInternalRequest,
      ServiceAccountRequest,
      ServiceAccountResponse,
      ServiceAccountUpdate,
      ServiceConnectorConfiguration,
      ServiceConnectorFilter,
      ServiceConnectorRequest,
      ServiceConnectorResourcesModel,
      ServiceConnectorResponse,
      ServiceConnectorTypeModel,
      ServiceConnectorUpdate,
      ServiceFilter,
      ServiceRequest,
      ServiceResponse,
      ServiceUpdate,
      StackDeploymentConfig,
      StackDeploymentInfo,
      StackFilter,
      StackRequest,
      StackResponse,
      StackUpdate,
      StepRunFilter,
      StepRunRequest,
      StepRunResponse,
      StepRunUpdate,
      TagFilter,
      TagRequest,
      TagResourceRequest,
      TagResourceResponse,
      TagResponse,
      TagUpdate,
      TriggerExecutionFilter,
      TriggerExecutionRequest,
      TriggerExecutionResponse,
      TriggerFilter,
      TriggerRequest,
      TriggerResponse,
      TriggerUpdate,
      UserAuthModel,
      UserFilter,
      UserRequest,
      UserResponse,
      UserScopedRequest,
      UserUpdate,
  )
  from zenml.service_connectors.service_connector_registry import (
      service_connector_registry,
  )
  from zenml.stack.flavor_registry import FlavorRegistry
  from zenml.stack_deployments.utils import get_stack_deployment_class
  from zenml.utils import tag_utils, uuid_utils
  from zenml.utils.enum_utils import StrEnum
  from zenml.utils.networking_utils import (
      replace_localhost_with_internal_hostname,
  )
  from zenml.utils.secret_utils import PlainSerializedSecretStr
  from zenml.utils.string_utils import (
      format_name_template,
      random_str,
      validate_name,
  )
  from zenml.utils.time_utils import utc_now
  from zenml.zen_stores import template_utils
  from zenml.zen_stores.base_zen_store import (
      BaseZenStore,
  )
  from zenml.zen_stores.dag_generator import DAGGeneratorHelper
  from zenml.zen_stores.migrations.alembic import (
      Alembic,
  )
  from zenml.zen_stores.migrations.utils import MigrationUtils
  from zenml.zen_stores.schemas import (
      ActionSchema,
      APIKeySchema,
      ApiTransactionSchema,
      ArtifactSchema,
      ArtifactVersionSchema,
      BaseSchema,
      CodeReferenceSchema,
      CodeRepositorySchema,
      CuratedVisualizationSchema,
      DeploymentSchema,
      EventSourceSchema,
      FlavorSchema,
      ModelSchema,
      ModelVersionArtifactSchema,
      ModelVersionPipelineRunSchema,
      ModelVersionSchema,
      NamedSchema,
      OAuthDeviceSchema,
      PipelineBuildSchema,
      PipelineRunSchema,
      PipelineSchema,
      PipelineSnapshotSchema,
      ProjectSchema,
      RunMetadataResourceSchema,
      RunMetadataSchema,
      RunTemplateSchema,
      ScheduleSchema,
      SecretResourceSchema,
      SecretSchema,
      ServerSettingsSchema,
      ServiceConnectorSchema,
      StackComponentSchema,
      StackSchema,
      StepConfigurationSchema,
      StepRunInputArtifactSchema,
      StepRunOutputArtifactSchema,
      StepRunParentsSchema,
      StepRunSchema,
      TagResourceSchema,
      TagSchema,
      TriggerExecutionSchema,
      UserSchema,
  )
  from zenml.zen_stores.schemas.artifact_visualization_schemas import (
      ArtifactVisualizationSchema,
  )
  from zenml.zen_stores.schemas.logs_schemas import LogsSchema
  from zenml.zen_stores.schemas.service_schemas import ServiceSchema
  from zenml.zen_stores.schemas.trigger_schemas import TriggerSchema
  from zenml.zen_stores.schemas.utils import (
      get_resource_type_name,
      jl_arg,
  )
  from zenml.zen_stores.secrets_stores.base_secrets_store import BaseSecretsStore
  from zenml.zen_stores.secrets_stores.sql_secrets_store import (
      SqlSecretsStoreConfiguration,
  )
  
  if TYPE_CHECKING:
      from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
  
  AnyNamedSchema = TypeVar("AnyNamedSchema", bound=NamedSchema)
  AnySchema = TypeVar("AnySchema", bound=BaseSchema)
  
  AnyResponse = TypeVar("AnyResponse", bound=BaseResponse)  # type: ignore[type-arg]  # noqa: F821
  AnyIdentifiedResponse = TypeVar(
      "AnyIdentifiedResponse",
      bound=BaseIdentifiedResponse,  # type: ignore[type-arg]  # noqa: F821
  )
  
  # Enable SQL compilation caching to remove the https://sqlalche.me/e/14/cprf
  # warning
  SelectOfScalar.inherit_cache = True
  Select.inherit_cache = True
  
  logger = get_logger(__name__)
  
  ZENML_SQLITE_DB_FILENAME = "zenml.db"
  
  
  def exponential_backoff_with_jitter(
      attempt: int, base_duration: float = 0.05
  ) -> float:
      """Exponential backoff with jitter.
  
      Implemented the `Full jitter` algorithm described in
      https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
  
      Args:
          attempt: The backoff attempt.
          base_duration: The backoff base duration.
  
      Returns:
          The backoff duration.
      """
      exponential_backoff = base_duration * 1.5**attempt
      return random.uniform(0, exponential_backoff)
  
  
  class Session(SqlModelSession):
      """Session subclass that automatically tracks duration and calling context."""
  
      def _get_metrics(self) -> Dict[str, Any]:
          """Get the metrics for the session.
  
          Returns:
              The metrics for the session.
          """
          # Get SQLAlchemy connection pool info
          assert isinstance(self.bind, Engine)
          assert isinstance(self.bind.pool, QueuePool)
          checked_out_connections = self.bind.pool.checkedout()
          available_connections = self.bind.pool.checkedin()
          overflow = self.bind.pool.overflow()
  
          return {
              "active_connections": checked_out_connections,
              "idle_connections": available_connections,
              "overflow_connections": overflow,
          }
  
      def _get_metrics_log_str(self) -> str:
          """Get the metrics for the session as a string for logging.
  
          Returns:
              The metrics for the session as a string for logging.
          """
          if not logger.isEnabledFor(logging.DEBUG):
              return ""
          metrics = self._get_metrics()
          # Add the server metrics if running in a server
          if handle_bool_env_var(ENV_ZENML_SERVER):
              from zenml.zen_server.utils import get_system_metrics
  
              metrics.update(get_system_metrics())
  
          return (
              " [ "
              + " ".join([f"{key}: {value}" for key, value in metrics.items()])
              + " ]"
          )
  
      def __enter__(self) -> "Session":
          """Enter the context manager.
  
          Returns:
              The SqlModel session.
          """
          if logger.isEnabledFor(logging.DEBUG):
              self.log_request_id = "N/A"
              self.log_request = ""
              if handle_bool_env_var(ENV_ZENML_SERVER):
                  # Running inside server
                  from zenml.zen_server.utils import get_current_request_context
  
                  # If the code is running on the server, use the auth context.
                  try:
                      request_context = get_current_request_context()
                      self.log_request_id = request_context.log_request_id
                      self.log_request = request_context.log_request
                  except RuntimeError:
                      pass
  
              # Look up the stack to find the SQLZenStore method
              for frame in inspect.stack():
                  if "self" in frame.frame.f_locals:
                      instance = frame.frame.f_locals["self"]
                      if isinstance(instance, SqlZenStore):
                          self.caller_method = (
                              f"{instance.__class__.__name__}.{frame.function}"
                          )
                          break
              else:
                  self.caller_method = "unknown"
  
              logger.debug(
                  f"[{self.log_request_id}] SQL STATS - "
                  f"{self.log_request} "
                  f"'{self.caller_method}' STARTED {self._get_metrics_log_str()}"
              )
  
              self.start_time = time.time()
  
          return super().__enter__()
  
      def __exit__(
          self,
          exc_type: Optional[Any],
          exc_val: Optional[Any],
          exc_tb: Optional[Any],
      ) -> None:
          """Exit the context manager.
  
          Args:
              exc_type: The exception type.
              exc_val: The exception value.
              exc_tb: The exception traceback.
          """
          if logger.isEnabledFor(logging.DEBUG):
              duration = (time.time() - self.start_time) * 1000
  
              # Add error information to the log
              error_info = ""
              if exc_type is not None:
                  error_info = " with ERROR"
  
              logger.debug(
                  f"[{self.log_request_id}] SQL STATS - "
                  f"{self.log_request} "
                  f"'{self.caller_method}' COMPLETED in "
                  f"{duration:.2f}ms {error_info} {self._get_metrics_log_str()}"
              )
  
          super().__exit__(exc_type, exc_val, exc_tb)
  
  
  class SQLDatabaseDriver(StrEnum):
      """SQL database drivers supported by the SQL ZenML store."""
  
      MYSQL = "mysql"
      SQLITE = "sqlite"
  
  
  class SqlZenStoreConfiguration(StoreConfiguration):
      """SQL ZenML store configuration.
  
      Attributes:
          type: The type of the store.
          secrets_store: The configuration of the secrets store to use.
              This defaults to a SQL secrets store that extends the SQL ZenML
              store.
          backup_secrets_store: The configuration of a backup secrets store to
              use in addition to the primary one as an intermediate step during
              the migration to a new secrets store.
          driver: The SQL database driver.
          database: database name. If not already present on the server, it will
              be created automatically on first access.
          username: The database username.
          password: The database password.
          ssl: Whether to use SSL.
          ssl_ca: certificate authority certificate. Required for SSL
              enabled authentication if the CA certificate is not part of the
              certificates shipped by the operating system.
          ssl_cert: client certificate. Required for SSL enabled
              authentication if client certificates are used.
          ssl_key: client certificate private key. Required for SSL
              enabled if client certificates are used.
          ssl_verify_server_cert: set to verify the identity of the server
              against the provided server certificate.
          pool_size: The maximum number of connections to keep in the SQLAlchemy
              pool.
          max_overflow: The maximum number of connections to allow in the
              SQLAlchemy pool in addition to the pool_size.
          pool_pre_ping: Enable emitting a test statement on the SQL connection
              at the start of each connection pool checkout, to test that the
              database connection is still viable.
      """
  
      type: StoreType = StoreType.SQL
  
      secrets_store: Optional[SerializeAsAny[SecretsStoreConfiguration]] = None
      backup_secrets_store: Optional[
          SerializeAsAny[SecretsStoreConfiguration]
      ] = None
  
      driver: Optional[SQLDatabaseDriver] = None
      database: Optional[str] = None
      username: Optional[PlainSerializedSecretStr] = None
      password: Optional[PlainSerializedSecretStr] = None
      ssl: bool = False
      ssl_ca: Optional[PlainSerializedSecretStr] = None
      ssl_cert: Optional[PlainSerializedSecretStr] = None
      ssl_key: Optional[PlainSerializedSecretStr] = None
      ssl_verify_server_cert: bool = False
      pool_size: int = 20
      max_overflow: int = 20
      pool_pre_ping: bool = True
  
      backup_strategy: DatabaseBackupStrategy = DatabaseBackupStrategy.IN_MEMORY
      # database backup directory
      backup_directory: str = Field(
          default_factory=lambda: os.path.join(
              GlobalConfiguration().config_directory,
              SQL_STORE_BACKUP_DIRECTORY_NAME,
          )
      )
      backup_database: Optional[str] = None
  
      @field_validator("secrets_store")
      @classmethod
      def validate_secrets_store(
          cls, secrets_store: Optional[SecretsStoreConfiguration]
      ) -> SecretsStoreConfiguration:
          """Ensures that the secrets store is initialized with a default SQL secrets store.
  
          Args:
              secrets_store: The secrets store config to be validated.
  
          Returns:
              The validated secrets store config.
          """
          if secrets_store is None:
              secrets_store = SqlSecretsStoreConfiguration()
  
          return secrets_store
  
      @model_validator(mode="after")
      def _validate_backup_strategy(self) -> "SqlZenStoreConfiguration":
          """Validate the backup strategy.
  
          Returns:
              The model attribute values.
  
          Raises:
              ValueError: If the backup database name is not set when the backup
                  database is requested.
          """
          if (
              self.backup_strategy == DatabaseBackupStrategy.DATABASE
              and not self.backup_database
          ):
              raise ValueError(
                  "The `backup_database` attribute must also be set if the "
                  "backup strategy is set to use a backup database."
              )
  
          return self
  
      @model_validator(mode="after")
      def _validate_url(self) -> "SqlZenStoreConfiguration":
          """Validate the SQL URL.
  
          The validator also moves the MySQL username, password and database
          parameters from the URL into the other configuration arguments, if they
          are present in the URL.
  
          Returns:
              The validated values.
  
          Raises:
              ValueError: If the URL is invalid or the SQL driver is not
                  supported.
          """
          if self.url is None:
              return self
  
          # When running inside a container, if the URL uses localhost, the
          # target service will not be available. We try to replace localhost
          # with one of the special Docker or K3D internal hostnames.
          url = replace_localhost_with_internal_hostname(self.url)
  
          try:
              sql_url = make_url(url)
          except ArgumentError as e:
              raise ValueError(
                  "Invalid SQL URL `%s`: %s. The URL must be in the format "
                  "`driver://[[username:password@]hostname:port]/database["
                  "?<extra-args>]`.",
                  url,
                  str(e),
              )
  
          if sql_url.drivername not in SQLDatabaseDriver.values():
              raise ValueError(
                  "Invalid SQL driver value `%s`: The driver must be one of: %s.",
                  url,
                  ", ".join(SQLDatabaseDriver.values()),
              )
          self.driver = SQLDatabaseDriver(sql_url.drivername)
          if sql_url.drivername == SQLDatabaseDriver.SQLITE:
              if (
                  sql_url.username
                  or sql_url.password
                  or sql_url.query
                  or sql_url.database is None
              ):
                  raise ValueError(
                      "Invalid SQLite URL `%s`: The URL must be in the "
                      "format `sqlite:///path/to/database.db`.",
                      url,
                  )
              if self.username or self.password:
                  raise ValueError(
                      "Invalid SQLite configuration: The username and password "
                      "must not be set",
                      url,
                  )
              self.database = sql_url.database
          elif sql_url.drivername == SQLDatabaseDriver.MYSQL:
              if sql_url.username:
                  self.username = PlainSerializedSecretStr(sql_url.username)
                  sql_url = sql_url._replace(username=None)
              if sql_url.password:
                  self.password = PlainSerializedSecretStr(sql_url.password)
                  sql_url = sql_url._replace(password=None)
              if sql_url.database:
                  self.database = sql_url.database
                  sql_url = sql_url._replace(database=None)
              if sql_url.query:
  
                  def _get_query_result(
                      result: Union[str, Tuple[str, ...]],
                  ) -> Optional[str]:
                      """Returns the only or the first result of a query.
  
                      Args:
                          result: The result of the query.
  
                      Returns:
                          The only or the first result, None otherwise.
                      """
                      if isinstance(result, str):
                          return result
                      elif isinstance(result, tuple) and len(result) > 0:
                          return result[0]
                      else:
                          return None
  
                  for k, v in sql_url.query.items():
                      if k == "ssl":
                          if r := _get_query_result(v):
                              self.ssl = is_true_string_value(r)
                      elif k == "ssl_ca":
                          if r := _get_query_result(v):
                              self.ssl_ca = PlainSerializedSecretStr(r)
                              self.ssl = True
                      elif k == "ssl_cert":
                          if r := _get_query_result(v):
                              self.ssl_cert = PlainSerializedSecretStr(r)
                              self.ssl = True
                      elif k == "ssl_key":
                          if r := _get_query_result(v):
                              self.ssl_key = PlainSerializedSecretStr(r)
                              self.ssl = True
                      elif k == "ssl_verify_server_cert":
                          if r := _get_query_result(v):
                              if is_true_string_value(r):
                                  self.ssl_verify_server_cert = True
                              elif is_false_string_value(r):
                                  self.ssl_verify_server_cert = False
                      else:
                          raise ValueError(
                              "Invalid MySQL URL query parameter `%s`: The "
                              "parameter must be one of: ssl, ssl_ca, ssl_cert, "
                              "ssl_key, or ssl_verify_server_cert.",
                              k,
                          )
                  sql_url = sql_url._replace(query=immutabledict())
  
              database = self.database
              if not self.username or not self.password or not database:
                  raise ValueError(
                      "Invalid MySQL configuration: The username, password and "
                      "database must be set in the URL or as configuration "
                      "attributes",
                  )
  
              regexp = r"^[^\\/?%*:|\"<>.-]{1,64}$"
              match = re.match(regexp, database)
              if not match:
                  raise ValueError(
                      f"The database name does not conform to the required "
                      f"format "
                      f"rules ({regexp}): {database}"
                  )
  
              # Save the certificates in a secure location on disk
              secret_folder = Path(
                  GlobalConfiguration().local_stores_path,
                  "certificates",
              )
              for key in ["ssl_key", "ssl_ca", "ssl_cert"]:
                  content = getattr(self, key)
                  if content and not os.path.isfile(content.get_secret_value()):
                      fileio.makedirs(str(secret_folder))
                      file_path = Path(secret_folder, f"{key}.pem")
                      with os.fdopen(
                          os.open(
                              file_path, flags=os.O_RDWR | os.O_CREAT, mode=0o600
                          ),
                          "w",
                      ) as f:
                          f.write(content.get_secret_value())
                      setattr(self, key, str(file_path))
  
          self.url = str(sql_url)
          return self
  
      @staticmethod
      def get_local_url(path: str) -> str:
          """Get a local SQL url for a given local path.
  
          Args:
              path: The path to the local sqlite file.
  
          Returns:
              The local SQL url for the given path.
          """
          return f"sqlite:///{path}/{ZENML_SQLITE_DB_FILENAME}"
  
      @classmethod
      def supports_url_scheme(cls, url: str) -> bool:
          """Check if a URL scheme is supported by this store.
  
          Args:
              url: The URL to check.
  
          Returns:
              True if the URL scheme is supported, False otherwise.
          """
          return make_url(url).drivername in SQLDatabaseDriver.values()
  
      def get_sqlalchemy_config(
          self,
          database: Optional[str] = None,
      ) -> Tuple[URL, Dict[str, Any], Dict[str, Any]]:
          """Get the SQLAlchemy engine configuration for the SQL ZenML store.
  
          Args:
              database: Custom database name to use. If not set, the database name
                  from the configuration will be used.
  
          Returns:
              The URL and connection arguments for the SQLAlchemy engine.
  
          Raises:
              NotImplementedError: If the SQL driver is not supported.
          """
          sql_url = make_url(self.url)
          sqlalchemy_connect_args: Dict[str, Any] = {}
          engine_args = {}
          if sql_url.drivername == SQLDatabaseDriver.SQLITE:
              assert self.database is not None
              # The following default value is needed for sqlite to avoid the
              # Error:
              #   sqlite3.ProgrammingError: SQLite objects created in a thread can
              #   only be used in that same thread.
              sqlalchemy_connect_args = {"check_same_thread": False}
          elif sql_url.drivername == SQLDatabaseDriver.MYSQL:
              # all these are guaranteed by our root validator
              assert self.database is not None
              assert self.username is not None
              assert self.password is not None
              assert sql_url.host is not None
  
              if not database:
                  database = self.database
  
              engine_args = {
                  "pool_size": self.pool_size,
                  "max_overflow": self.max_overflow,
                  "pool_pre_ping": self.pool_pre_ping,
              }
  
              sql_url = sql_url._replace(
                  drivername="mysql+pymysql",
                  username=self.username.get_secret_value(),
                  password=self.password.get_secret_value(),
                  database=database,
              )
  
              sqlalchemy_ssl_args: Dict[str, Any] = {}
  
              # Handle SSL params
              if self.ssl:
                  sqlalchemy_ssl_args["ssl"] = True
                  for key in ["ssl_key", "ssl_ca", "ssl_cert"]:
                      ssl_setting = getattr(self, key)
                      if not ssl_setting:
                          continue
                      if not os.path.isfile(ssl_setting.get_secret_value()):
                          logger.warning(
                              f"Database SSL setting `{key}` is not a file. "
                          )
                      sqlalchemy_ssl_args[key.removeprefix("ssl_")] = (
                          ssl_setting.get_secret_value()
                      )
                  sqlalchemy_ssl_args["check_hostname"] = (
                      self.ssl_verify_server_cert
                  )
                  sqlalchemy_connect_args["ssl"] = sqlalchemy_ssl_args
          else:
              raise NotImplementedError(
                  f"SQL driver `{sql_url.drivername}` is not supported."
              )
  
          return sql_url, sqlalchemy_connect_args, engine_args
  
      model_config = ConfigDict(
          # Don't validate attributes when assigning them. This is necessary
          # because the certificate attributes can be expanded to the contents
          # of the certificate files.
          validate_assignment=False,
          # Forbid extra attributes set in the class.
          extra="ignore",
      )
  
  
  class SqlZenStore(BaseZenStore):
      """Store Implementation that uses SQL database backend.
  
      Attributes:
          config: The configuration of the SQL ZenML store.
          skip_migrations: Whether to skip migrations when initializing the store.
          TYPE: The type of the store.
          CONFIG_TYPE: The type of the store configuration.
          _engine: The SQLAlchemy engine.
      """
  
      config: SqlZenStoreConfiguration
      skip_migrations: bool = False
      TYPE: ClassVar[StoreType] = StoreType.SQL
      CONFIG_TYPE: ClassVar[Type[StoreConfiguration]] = SqlZenStoreConfiguration
  
      _engine: Optional[Engine] = None
      _migration_utils: Optional[MigrationUtils] = None
      _alembic: Optional[Alembic] = None
      _secrets_store: Optional[BaseSecretsStore] = None
      _backup_secrets_store: Optional[BaseSecretsStore] = None
      _should_send_user_enriched_events: bool = False
      _cached_onboarding_state: Optional[Set[str]] = None
      _default_user: Optional[UserResponse] = None
  
      @property
      def secrets_store(self) -> "BaseSecretsStore":
          """The secrets store associated with this store.
  
          Returns:
  