content: """track secrets in db [7b651bf6822e].
  
  Revision ID: 7b651bf6822e
  Revises: 1ac1b9c04da1
  Create Date: 2023-12-22 20:56:18.131906
  
  """
  
  from abc import abstractmethod
  from datetime import datetime
  from typing import Any, Dict, List, Optional
  from uuid import UUID
  
  import sqlalchemy as sa
  from alembic import op
  from pydantic import BaseModel
  
  from zenml.logger import get_logger
  
  logger = get_logger(__name__)
  
  # revision identifiers, used by Alembic.
  
  
  revision = "7b651bf6822e"
  down_revision = "1ac1b9c04da1"
  branch_labels = None
  depends_on = None
  
  
  ZENML_SECRET_LABEL = "zenml"
  ZENML_SECRET_ID_LABEL = "zenml_secret_id"
  ZENML_SECRET_NAME_LABEL = "zenml_secret_name"
  ZENML_SECRET_SCOPE_LABEL = "zenml_secret_scope"
  ZENML_SECRET_USER_LABEL = "zenml_secret_user"
  ZENML_SECRET_WORKSPACE_LABEL = "zenml_secret_workspace"
  
  
  class ZenMLSecretMetadata(BaseModel):
      """ZenML secret metadata.
  
      This class reflects the structure of the ZenML secret metadata stored in the
      backend secret (e.g. in secret tags/labels) at the time when the secret
      metadata was migrated to the database. It is roughly aligned to the
      structure of the ZenML SecretResponse model as it was at that time.
      """
  
      id: UUID
      name: str
      user_id: UUID
      workspace_id: UUID
      created: datetime
      updated: datetime
      scope: str
  
  
  class BaseSecretsStoreBackend(BaseModel):
      """Base class for accessing and listing ZenML secrets."""
  
      server_id: UUID
      client: Any = None
  
      def _get_secret_metadata(
          self,
          secret_id: Optional[UUID] = None,
      ) -> Dict[str, str]:
          """Get a dictionary with metadata that can be used as tags/labels.
  
          This utility method can be used with Secrets Managers that can
          associate metadata (e.g. tags, labels) with a secret. The metadata can
          be configured alongside each secret.
  
          NOTE: the ZENML_SECRET_LABEL is always included in the metadata to
          distinguish ZenML secrets from other secrets that might be stored in
          the same backend, as well as to distinguish between different ZenML
          deployments using the same backend. Its value is set to the ZenML
          deployment ID.
  
          Args:
              secret_id: Optional secret ID to include in the metadata.
  
          Returns:
              Dictionary with secret metadata information.
          """
          # Always include the main ZenML label to distinguish ZenML secrets
          # from other secrets that might be stored in the same backend and
          # to distinguish between different ZenML deployments using the same
          # backend.
          metadata: Dict[str, str] = {ZENML_SECRET_LABEL: str(self.server_id)}
  
          # Include the secret ID if provided.
          if secret_id is not None:
              metadata[ZENML_SECRET_ID_LABEL] = str(secret_id)
  
          return metadata
  
      def _create_secret_from_metadata(
          self,
          metadata: Dict[str, str],
          created: datetime,
          updated: datetime,
      ) -> ZenMLSecretMetadata:
          """Create a ZenML secret metadata object from metadata stored in the secrets store backend.
  
          Args:
              metadata: ZenML secret metadata collected from the backend secret
                  (e.g. from secret tags/labels).
              created: The secret creation time.
              updated: The secret last updated time.
  
          Returns:
              The ZenML secret.
  
          Raises:
              KeyError: If the secret does not have the required metadata, if it
                  is not managed by this ZenML instance or if it is linked to a
                  user or workspace that no longer exists.
          """
          # Double-check that the secret is managed by this ZenML instance.
          if metadata.get(ZENML_SECRET_LABEL) != str(self.server_id):
              raise KeyError("Secret is not managed by this ZenML instance")
  
          # Recover the ZenML secret fields from the input secret metadata.
          try:
              secret_id = UUID(metadata[ZENML_SECRET_ID_LABEL])
              name = metadata[ZENML_SECRET_NAME_LABEL]
              scope = metadata[ZENML_SECRET_SCOPE_LABEL]
              workspace_id = UUID(metadata[ZENML_SECRET_WORKSPACE_LABEL])
              user_id = UUID(metadata[ZENML_SECRET_USER_LABEL])
          except KeyError as e:
              raise KeyError(
                  f"Secret could not be retrieved: missing required metadata: {e}"
              )
  
          secret_model = ZenMLSecretMetadata(
              id=secret_id,
              name=name,
              workspace_id=workspace_id,
              user_id=user_id,
              created=created,
              updated=updated,
              scope=scope,
          )
  
          return secret_model
  
      @abstractmethod
      def list_secrets(
          self,
      ) -> List[ZenMLSecretMetadata]:
          """List all ZenML secrets in the secrets store backend.
  
          Returns:
              A list of all secrets.
          """
  
  
  class AWSSecretsStoreBackend(BaseSecretsStoreBackend):
      """AWS ZenML secrets store backend."""
  
      @staticmethod
      def _get_aws_secret_filters(
          metadata: Dict[str, str],
      ) -> List[Dict[str, str]]:
          """Convert ZenML secret metadata to AWS secret filters.
  
          Args:
              metadata: The ZenML secret metadata.
  
          Returns:
              The AWS secret filters.
          """
          aws_filters: List[Dict[str, Any]] = []
          for k, v in metadata.items():
              aws_filters.append(
                  {
                      "Key": "tag-key",
                      "Values": [
                          k,
                      ],
                  }
              )
              aws_filters.append(
                  {
                      "Key": "tag-value",
                      "Values": [
                          str(v),
                      ],
                  }
              )
  
          return aws_filters
  
      def list_secrets(
          self,
      ) -> List[ZenMLSecretMetadata]:
          """List all ZenML secrets in the AWS secrets store backend.
  
          Returns:
              A list of all secrets.
  
          Raises:
              RuntimeError: If the AWS Secrets Manager API returns an unexpected
                  error.
          """
          from botocore.exceptions import ClientError
  
          # The metadata will always contain at least the filter criteria
          # required to exclude everything but AWS secrets that belong to the
          # current ZenML deployment.
          metadata = self._get_secret_metadata()
          aws_filters = self._get_aws_secret_filters(metadata)
  
          results: List[ZenMLSecretMetadata] = []
  
          try:
              # AWS Secrets Manager API pagination is wrapped around the
              # `list_secrets` method call. We use it because we need to fetch all
              # secrets matching the (partial) filter that we set up. Note that
              # the pagination used here has nothing to do with the pagination
              # that we do for the method caller.
              paginator = self.client.get_paginator("list_secrets")
              pages = paginator.paginate(
                  Filters=aws_filters,
                  PaginationConfig={
                      "PageSize": 100,
                  },
              )
  
              for page in pages:
                  for secret in page["SecretList"]:
                      try:
                          # Convert the AWS secret tags to a metadata dictionary.
                          unpacked_metadata: Dict[str, str] = {
                              tag["Key"]: tag["Value"] for tag in secret["Tags"]
                          }
  
                          secret_model = self._create_secret_from_metadata(
                              metadata=unpacked_metadata,
                              created=secret["CreatedDate"],
                              updated=secret["LastChangedDate"],
                          )
                      except KeyError:
                          # The _convert_aws_secret method raises a KeyError
                          # if the secret is not well formed. Here we pretend that
                          # the secret does not exist.
                          continue
  
                      results.append(secret_model)
          except ClientError as e:
              raise RuntimeError(f"Error listing AWS secrets: {e}")
  
          return results
  
  
  ZENML_GCP_DATE_FORMAT_STRING = "%Y-%m-%d-%H-%M-%S"
  ZENML_GCP_SECRET_CREATED_KEY = "zenml-secret-created"
  ZENML_GCP_SECRET_UPDATED_KEY = "zenml-secret-updated"
  
  
  class GCPSecretsStoreBackend(BaseSecretsStoreBackend):
      """GCP ZenML secrets store backend."""
  
      project_id: str
  
      @property
      def parent_name(self) -> str:
          """Construct the GCP parent path to the secret manager.
  
          Returns:
              The parent path to the secret manager
          """
          return f"projects/{self.project_id}"
  
      def _convert_gcp_secret(
          self,
          labels: Dict[str, str],
      ) -> ZenMLSecretMetadata:
          """Create a ZenML secret model from data stored in an GCP secret.
  
          If the GCP secret cannot be converted, the method acts as if the
          secret does not exist and raises a KeyError.
  
          Args:
              labels: The GCP secret labels.
  
          Returns:
              The ZenML secret model.
  
          Raises:
              KeyError: if the GCP secret cannot be converted.
          """
          # Recover the ZenML secret metadata from the AWS secret tags.
  
          # The GCP secret labels do not really behave like a dictionary: when
          # a key is not found, it does not raise a KeyError, but instead
          # returns an empty string. That's why we make this conversion.
          label_dict = dict(labels)
  
          try:
              created = datetime.strptime(
                  label_dict[ZENML_GCP_SECRET_CREATED_KEY],
                  ZENML_GCP_DATE_FORMAT_STRING,
              )
              updated = datetime.strptime(
                  label_dict[ZENML_GCP_SECRET_UPDATED_KEY],
                  ZENML_GCP_DATE_FORMAT_STRING,
              )
          except KeyError as e:
              raise KeyError(
                  f"Invalid GCP secret: missing required tag '{e}'"
              ) from e
  
          return self._create_secret_from_metadata(
              metadata=label_dict,
              created=created,
              updated=updated,
          )
  
      def list_secrets(self) -> List[ZenMLSecretMetadata]:
          """List all secrets.
  
          Returns:
              A list of all secrets.
  
          Raises:
              RuntimeError: if the GCP Secrets Manager API returns an unexpected
                  error.
          """
          from google.cloud.secretmanager import SecretManagerServiceClient
  
          assert isinstance(self.client, SecretManagerServiceClient)
  
          try:
              # get all the secrets and their labels from GCP
              secrets = []
              for secret in self.client.list_secrets(
                  request={
                      "parent": self.parent_name,
                      "filter": "",
                  }
              ):
                  try:
                      secrets.append(self._convert_gcp_secret(secret.labels))
                  except KeyError:
                      # keep going / ignore if this secret version doesn't exist
                      # or isn't a ZenML secret
                      continue
          except Exception as e:
              raise RuntimeError(f"Error listing GCP secrets: {e}") from e
  
          return secrets
  
  
  ZENML_AZURE_SECRET_CREATED_KEY = "zenml_secret_created"
  ZENML_AZURE_SECRET_UPDATED_KEY = "zenml_secret_updated"
  
  
  class AzureSecretsStoreBackend(BaseSecretsStoreBackend):
      """Azure ZenML secrets store backend."""
  
      def _convert_azure_secret(
          self,
          tags: Dict[str, str],
      ) -> ZenMLSecretMetadata:
          """Create a ZenML secret model from data stored in an Azure secret.
  
          If the Azure secret cannot be converted, the method acts as if the
          secret does not exist and raises a KeyError.
  
          Args:
              tags: The Azure secret tags.
  
          Returns:
              The ZenML secret.
  
          Raises:
              KeyError: if the Azure secret cannot be converted.
          """
          try:
              created = datetime.fromisoformat(
                  tags[ZENML_AZURE_SECRET_CREATED_KEY],
              )
              updated = datetime.fromisoformat(
                  tags[ZENML_AZURE_SECRET_UPDATED_KEY],
              )
          except KeyError as e:
              raise KeyError(
                  f"Secret could not be retrieved: missing required metadata: {e}"
              )
  
          return self._create_secret_from_metadata(
              metadata=tags,
              created=created,
              updated=updated,
          )
  
      def list_secrets(self) -> List[ZenMLSecretMetadata]:
          """List all secrets.
  
          Returns:
              A list of all secrets.
  
          Raises:
              RuntimeError: If the Azure Key Vault API returns an unexpected
                  error.
          """
          from azure.core.exceptions import HttpResponseError
          from azure.keyvault.secrets import SecretClient
  
          assert isinstance(self.client, SecretClient)
  
          results: List[ZenMLSecretMetadata] = []
  
          try:
              all_secrets = self.client.list_properties_of_secrets()
              for secret_property in all_secrets:
                  try:
                      # NOTE: we do not include the secret values in the
                      # response. We would need a separate API call to fetch
                      # them for each secret, which would be very inefficient
                      # anyway.
                      assert secret_property.tags is not None
                      secret_model = self._convert_azure_secret(
                          tags=secret_property.tags,
                      )
                  except KeyError:
                      # The _convert_azure_secret method raises a KeyError
                      # if the secret is tied to a workspace or user that no
                      # longer exists or if it is otherwise not valid. Here we
                      # pretend that the secret does not exist.
                      continue
  
                  results.append(secret_model)
          except HttpResponseError as e:
              raise RuntimeError(f"Error listing Azure Key Vault secrets: {e}")
  
          return results
  
  
  HVAC_ZENML_SECRET_NAME_PREFIX = "zenml"
  ZENML_VAULT_SECRET_METADATA_KEY = "zenml_secret_metadata"
  ZENML_VAULT_SECRET_CREATED_KEY = "zenml_secret_created"
  ZENML_VAULT_SECRET_UPDATED_KEY = "zenml_secret_updated"
  
  
  class HashiCorpVaultSecretsStoreBackend(BaseSecretsStoreBackend):
      """HashiCorp Vault ZenML secrets store backend."""
  
      def _convert_vault_secret(
          self,
          vault_secret: Dict[str, Any],
      ) -> ZenMLSecretMetadata:
          """Create a ZenML secret model from data stored in an HashiCorp Vault secret.
  
          If the HashiCorp Vault secret cannot be converted, the method acts as if
          the secret does not exist and raises a KeyError.
  
          Args:
              vault_secret: The HashiCorp Vault secret in JSON form.
  
          Returns:
              The ZenML secret.
  
          Raises:
              KeyError: if the HashiCorp Vault secret cannot be converted.
          """
          try:
              metadata = vault_secret[ZENML_VAULT_SECRET_METADATA_KEY]
              created = datetime.fromisoformat(
                  vault_secret[ZENML_VAULT_SECRET_CREATED_KEY],
              )
              updated = datetime.fromisoformat(
                  vault_secret[ZENML_VAULT_SECRET_UPDATED_KEY],
              )
          except (KeyError, ValueError) as e:
              raise KeyError(
                  f"Secret could not be retrieved: missing required metadata: {e}"
              )
  
          return self._create_secret_from_metadata(
              metadata=metadata,
              created=created,
              updated=updated,
          )
  
      def list_secrets(self) -> List[ZenMLSecretMetadata]:
          """List all secrets.
  
          Note that returned secrets do not include any secret values. To fetch
          the secret values, use `get_secret`.
  
          Returns:
              A list of all secrets.
  
          Raises:
              RuntimeError: If the HashiCorp Vault API returns an unexpected
                  error.
          """
          from hvac import Client
          from hvac.exceptions import (
              InvalidPath,
              VaultError,
          )
  
          assert isinstance(self.client, Client)
  
          results: List[ZenMLSecretMetadata] = []
  
          try:
              # List all ZenML secrets in the Vault
              all_secrets = (
                  self.client.secrets.kv.v2.list_secrets(
                      path=HVAC_ZENML_SECRET_NAME_PREFIX
                  )
                  .get("data", {})
                  .get("keys", [])
              )
          except InvalidPath:
              # no secrets created yet
              pass
          except VaultError as e:
              raise RuntimeError(f"Error listing HashiCorp Vault secrets: {e}")
          else:
              # Convert the Vault secrets to ZenML secrets
              for secret_uuid in all_secrets:
                  vault_secret_id = (
                      f"{HVAC_ZENML_SECRET_NAME_PREFIX}/{secret_uuid}"
                  )
                  try:
                      vault_secret = (
                          self.client.secrets.kv.v2.read_secret(
                              path=vault_secret_id
                          )
                          .get("data", {})
                          .get("data", {})
                      )
                  except (InvalidPath, VaultError) as e:
                      logger.warning(
                          f"Error fetching secret with ID {vault_secret_id}: {e}",
                      )
                      continue
  
                  try:
                      secret_model = self._convert_vault_secret(
                          vault_secret,
                      )
                  except KeyError as e:
                      # The _convert_vault_secret method raises a KeyError
                      # if the secret is not valid. Here we
                      # pretend that the secret does not exist.
                      logger.warning(
                          f"Error fetching secret with ID {vault_secret_id}: {e}",
                      )
                      continue
  
                  results.append(secret_model)
  
          return results
  
  
  def upgrade() -> None:
      """Upgrade database schema and/or data, creating a new revision.
  
      Raises:
          NotImplementedError: If the secrets store type is not supported.
      """
      from zenml.config.global_config import GlobalConfiguration
      from zenml.enums import SecretsStoreType
      from zenml.zen_stores.sql_zen_store import SqlZenStore
  
      logger.info("Migrating secrets from the external secrets store to the db.")
  
      store_cfg = GlobalConfiguration().store_configuration
  
      # We create a new SqlZenStore instance with the same configuration as the
      # one used to run migrations, but we skip the default registrations and
      # migrations. This is because we can't use the same store instance that is
      # used to run migrations because it would create an infinite loop.
      store = SqlZenStore(
          config=store_cfg,
          skip_default_registrations=True,
          skip_migrations=True,
      )
  
      secrets_store = store.secrets_store
  
      if secrets_store.TYPE == SecretsStoreType.SQL:
          # If the secrets store is already a SQL secrets store, we don't need
          # to transfer secrets from the external secrets store to the db.
          logger.info(
              "Skipping migration of secrets from the external secrets store "
              "to the db because the db is already configured as a SQL secrets "
              "store."
          )
          return
  
      logger.debug(
          f"Transferring secrets from the {secrets_store.type} secrets store "
          f"to the db: {secrets_store.config}"
      )
  
      # Transfer secrets from the external secrets store to the db
  
      conn = op.get_bind()
      meta = sa.MetaData()
      meta.reflect(
          only=("secret", "user", "workspace", "identity"), bind=op.get_bind()
      )
      secrets = sa.Table("secret", meta)
      users = sa.Table("user", meta)
      workspaces = sa.Table("workspace", meta)
      identity = sa.Table("identity", meta)
  
      # Extract the ZenML deployment ID from the identity table
      server_id = UUID(conn.execute(sa.select(identity.c.id)).scalar_one())
  
      # Initialize the secrets store backend based on the backend client extracted
      # from the secrets store.
      backend: BaseSecretsStoreBackend
      if secrets_store.TYPE == SecretsStoreType.AWS:
          from zenml.zen_stores.secrets_stores.aws_secrets_store import (
              AWSSecretsStore,
          )
  
          assert isinstance(secrets_store, AWSSecretsStore)
          backend = AWSSecretsStoreBackend(
              server_id=server_id,
              client=secrets_store.client,
          )
      elif secrets_store.TYPE == SecretsStoreType.GCP:
          from zenml.zen_stores.secrets_stores.gcp_secrets_store import (
              GCPSecretsStore,
          )
  
          assert isinstance(secrets_store, GCPSecretsStore)
          backend = GCPSecretsStoreBackend(
              server_id=server_id,
              client=secrets_store.client,
              project_id=secrets_store.config.project_id,
          )
      elif secrets_store.TYPE == SecretsStoreType.AZURE:
          from zenml.zen_stores.secrets_stores.azure_secrets_store import (
              AzureSecretsStore,
          )
  
          assert isinstance(secrets_store, AzureSecretsStore)
          backend = AzureSecretsStoreBackend(
              server_id=server_id,
              client=secrets_store.client,
          )
      elif secrets_store.TYPE == SecretsStoreType.HASHICORP:
          from zenml.zen_stores.secrets_stores.hashicorp_secrets_store import (
              HashiCorpVaultSecretsStore,
          )
  
          assert isinstance(secrets_store, HashiCorpVaultSecretsStore)
          backend = HashiCorpVaultSecretsStoreBackend(
              server_id=server_id,
              client=secrets_store.client,
          )
      else:
          raise NotImplementedError(
              f"Secrets store type {secrets_store.TYPE} not supported"
          )
  
      # List all secrets in the backend
      external_secrets = backend.list_secrets()
      if len(external_secrets) == 0:
          logger.debug("No secrets found in the external secrets store.")
  
      for secret in external_secrets:
          logger.info(
              f"Transferring secret '{secret.name}' from the "
              f"{secrets_store.type} secrets store to the db."
          )
          # Check if a user with the given ID exists in the db
          user = conn.execute(
              users.select().where(
                  users.c.id == str(secret.user_id).replace("-", "")
              )
          ).first()
          if user is None:
              logger.warning(
                  f"User with ID {secret.user_id} not found in the db. "
                  f"Skipping secret '{secret.name}'."
              )
              continue
          # Check if a workspace with the given ID exists in the db
          workspace = conn.execute(
              workspaces.select().where(
                  workspaces.c.id == str(secret.workspace_id).replace("-", "")
              )
          ).first()
          if workspace is None:
              logger.warning(
                  f"Workspace with ID {secret.workspace_id} not found in the db. "
                  f"Skipping secret '{secret.name}'."
              )
              continue
  
          # Check if there is already a secret with the same ID in the db
          # and delete it if there is. This can happen for example if the user has
          # previously migrated the secrets from the db to an external secrets
          # store and forgot to delete the secrets from the db.
          conn.execute(
              secrets.delete().where(
                  secrets.c.id == str(secret.id).replace("-", "")
              )
          )
          conn.execute(
              secrets.insert().values(
                  id=str(secret.id).replace("-", ""),
                  created=secret.created,
                  updated=secret.updated,
                  name=secret.name,
                  scope=secret.scope,
                  values=None,
                  workspace_id=str(secret.workspace_id).replace("-", ""),
                  user_id=str(secret.user_id).replace("-", ""),
              )
          )
  
  
  def downgrade() -> None:
      """Downgrade database schema and/or data back to the previous revision."""
      pass
  