def setup_environment_variables()

in images/airflow/2.9.2/python/mwaa/config/setup_environment.py [0:0]

• 34 lines of code
• 9 McCabe index (conditional complexity)

def setup_environment_variables(command: str, executor_type: str) -> Dict:
    """Set up and return environment variables for Airflow execution.

            Configures the necessary environment variables based on the provided command
            and executor type for running Airflow tasks.

            Args:
                command (str): The Airflow command to be executed.
                executor_type (str): The type of executor to be used (e.g., 'Local', 'Celery').

            Returns:
                Dict: A dictionary containing all configured environment variables
                      required for the Airflow execution.
    """
    # Add the necessary environment variables.
    mwaa_essential_airflow_config = get_essential_airflow_config(executor_type)
    mwaa_opinionated_airflow_config = get_opinionated_airflow_config()
    mwaa_essential_airflow_environ = get_essential_environ(command)
    mwaa_opinionated_airflow_environ = get_opinionated_environ()
    user_airflow_config = get_user_airflow_config()


    startup_script_environ = _execute_startup_script(
        command,
        {
            **os.environ,
            **mwaa_opinionated_airflow_config,
            **mwaa_opinionated_airflow_environ,
            **user_airflow_config,
            **mwaa_essential_airflow_environ,
            **mwaa_essential_airflow_config,
        },
    )
    environ = {
        **os.environ,
        # Custom configuration and environment variables that we think are good, but
        # allow the user to override.
        **mwaa_opinionated_airflow_config,
        **mwaa_opinionated_airflow_environ,
        # What the user defined in the startup script.
        **startup_script_environ,
        # What the user passed via Airflow config secrets (specified by the
        # MWAA__CORE__CUSTOM_AIRFLOW_CONFIGS environment variable.)
        **user_airflow_config,
        # The MWAA__x__y environment variables that are passed to the container are
        # considered protected environment variables that cannot be overridden at
        # runtime to avoid breaking the functionality of the container.
        **{
            key: value
            for (key, value) in os.environ.items()
            if _is_protected_os_environ(key)
        },
        # Essential variables that our setup will not function properly without, hence
        # it always has the highest priority.
        **mwaa_essential_airflow_config,
        **mwaa_essential_airflow_environ,
    }

    # IMPORTANT NOTE: The level for this should stay "DEBUG" to avoid logging customer
    # custom environment variables, which potentially contains sensitive credentials,
    # to stdout which, in this case of Fargate hosting (like in Amazon MWAA), ends up
    # being captured and sent to the service hosting.
    logger.debug(f"Environment variables: %s", environ)

    # Export the environment variables to .bashrc and .bash_profile to enable
    # users to run a shell on the container and have the necessary environment
    # variables set for using airflow CLI.
    _export_env_variables(environ)

    return environ