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.
  """Stack component mixin for authentication."""
  
  from typing import Optional, Type, TypeVar, cast
  
  from pydantic import BaseModel, Field
  
  from zenml.client import Client
  from zenml.models import SecretResponse
  from zenml.stack.stack_component import StackComponent, StackComponentConfig
  
  T = TypeVar("T", bound=BaseModel)
  
  
  class AuthenticationConfigMixin(StackComponentConfig):
      """Base config for authentication mixins.
  
      Any stack component that implements `AuthenticationMixin` should have a
      config that inherits from this class.
  
      Field descriptions are defined inline using Field() descriptors.
      """
  
      authentication_secret: Optional[str] = Field(
          default=None,
          description="Name of the ZenML secret containing authentication credentials.",
      )
  
  
  class AuthenticationMixin(StackComponent):
      """Stack component mixin for authentication.
  
      Any stack component that implements this mixin should have a config that
      inherits from `AuthenticationConfigMixin`.
      """
  
      @property
      def config(self) -> AuthenticationConfigMixin:
          """Returns the `AuthenticationConfigMixin` config.
  
          Returns:
              The configuration.
          """
          return cast(AuthenticationConfigMixin, self._config)
  
      def get_authentication_secret(
          self,
      ) -> Optional[SecretResponse]:
          """Gets the secret referred to by the authentication secret attribute.
  
          Returns:
              The secret if the `authentication_secret` attribute is set,
              `None` otherwise.
  
          Raises:
              KeyError: If the secret does not exist.
          """
          if not self.config.authentication_secret:
              return None
  
          # Try to resolve the secret using the secret store
          try:
              return Client().get_secret_by_name_and_private_status(
                  name=self.config.authentication_secret,
              )
          except (KeyError, NotImplementedError):
              raise KeyError(
                  f"The authentication secret {self.config.authentication_secret} "
                  f"referenced by the `{self.name}` `{self.type}` stack "
                  "component does not exist."
              )
  
      def get_typed_authentication_secret(
          self, expected_schema_type: Type[T]
      ) -> Optional[T]:
          """Gets a typed secret referred to by the authentication secret attribute.
  
          Args:
              expected_schema_type: A Pydantic model class that represents the
                  expected schema type of the secret.
  
          Returns:
              The secret values extracted from the secret and converted into the
              indicated Pydantic type, if the `authentication_secret` attribute is
              set, `None` otherwise.
  
          Raises:
              TypeError: If the secret cannot be converted into the indicated
                  Pydantic type.
          """
          secret = self.get_authentication_secret()
  
          if not secret:
              return None
  
          try:
              typed_secret = expected_schema_type(
                  **secret.secret_values,
              )
          except (TypeError, ValueError) as e:
              raise TypeError(
                  f"Authentication secret `{self.config.authentication_secret}` "
                  f"referenced by the `{self.name}` `{self.type}` stack component"
                  f"could not be converted to {expected_schema_type}: {e}"
              )
  
          return typed_secret
  