Source code for allensdk.core.authentication

import os
from typing import Optional, Dict, Any
import logging
from functools import wraps
from abc import ABC, abstractmethod
from collections import namedtuple
from allensdk.core.auth_config import CREDENTIAL_KEYS


logger = logging.getLogger(__name__)

DbCredentials = namedtuple("DbCredentials",
                           ["dbname", "user", "host", "port", "password"])


class CredentialProvider(ABC):
    METHOD = "custom"

    @abstractmethod
    def provide(self, credential):
        pass


class EnvCredentialProvider(CredentialProvider):
    """
    Provides credentials from environment variables for variables listed
    in CREDENTIAL_KEYS.
    """
    METHOD = "env"

    def __init__(self, environ: Optional[Dict[str, Any]] = None):
        """
        Parameters
        ----------
        environ: dictionary or os.environ
            A dictionary that provides the values for keys in
            CREDENTIAL_KEYS. If not provided, defaults to os.environ to
            provide environment variables.
        """
        if environ is None:
            environ = os.environ
        self.credentials = dict((k[0], environ.get(k[0], k[1]))
                                for k in CREDENTIAL_KEYS)

    def provide(self, credential):
        return self.credentials.get(credential)


CREDENTIAL_PROVIDER = EnvCredentialProvider()


def set_credential_provider(provider):
    logger.info(f"Setting provider to method '{provider.METHOD}.")
    global CREDENTIAL_PROVIDER
    CREDENTIAL_PROVIDER = provider


def get_credential_provider():
    return CREDENTIAL_PROVIDER


def credential_injector(credential_map: Dict[str, Any],
                        provider: Optional[CredentialProvider] = None):
    """
    Decorator used to inject credentials from another source if not
    explicitly provided in the function call. This function will only supply
    values for keyword arguments. All keys defined in `credential_map` must
    correspond to keyword arguments in the function signature.

    PARAMETERS
    ----------
    credential_map: Dict[Str: Any]
        Dictionary where the keys are the keyword of a credential kwarg
        passed to the decorated function, and the values are the name
        of the credential in the credential provider (see CREDENTIAL_KEYS).

        Example of credential_map for PostgresQueryMixin connecting to
        LIMS database:
            {
                "dbname": "LIMS_DBNAME",
                "user": "LIMS_USER",
                "host": "LIMS_HOST",
                "password": "LIMS_PASSWORD",
                "port": "LIMS_PORT"
            }
    provider: Optional[CredentialProvider]
        Subclass of CredentialProvider to provide credentials to the
        wrapped function. If left unspecified, will default to
        EnvCredentialProvider, which provides credentials from environment
        variables.
    """
    if provider is None:
        provider = get_credential_provider()

    def injector_decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for kw, credential in credential_map.items():
                if kw not in kwargs.keys():
                    logger.debug(f"No explicit value provided for {kw}. "
                                 "Searching credential provider.")
                    secret = provider.provide(credential)
                    if secret is not None:
                        logger.debug("Found value in credential provider, "
                                     f"from '{provider.METHOD}' method.")
                        kwargs.update({kw: provider.provide(credential)})
                    else:
                        logger.warning(
                            f"Value for {kw} was neither explicitly provided "
                            "nor found in credential provider.")
            return func(*args, **kwargs)
        return wrapper
    return injector_decorator