Setting custom environment variables in managed apache airflow

Question

We are planning to switch from managing airflow ourselves to Managed Apache Airflow services of AWS. Our original dags use some custom environment variables that need to be set in Managed airflow as well. So far I was not able to find a way to set custom environment variables while setting up airflow environment in MWAA. Please let me know if anyone knows how to set them.

Bimonthly

Is docs.aws.amazon.com/mwaa/latest/userguide/… what you need? – Bony

Yes. But here we can only choose from the available configurations. There is a dropdown that shows up listing configuration options. We can configure only those options. But if I want to define a configuration option by myself as we can do in normal airflow configurations, Ex. AIRFLOW__CORE__MYCONFIG = 'something' I am not able to see how we can do that here. I require this for 2 purposes, one to directly use them in DAGs and the second, I need to set up ENV variables for my non-DAG scripts in plugins. – Bimonthly

Answer 1

An approach for setting environment variables is to use Airflow Variables.

In MWAA, you can store Airflow Variables in AWS Secrets Manager. This approach is documented in MWAA's official documentation. Note that this approach requires specific configuration for the MWAA environment. Once the MWAA environment is set up and the variables are stored in AWS Secrets Manager, the variables become accessible through the Airflow Variable APIs.

from airflow.models import Variable

# Normal call style
foo = Variable.get("foo")

See Step two: Create the Secrets Manager backend as an Apache Airflow configuration option
See Step four: Add the variables in Secrets Manager

Elinaelinor

Answer 2

you should be able to set core.myconfig env variable. mwaa accepts it even though its not on the list. mwaa will create AIRFLOW__CORE__MYCONFIG env variable.

Lammergeier

Thanks. I was trying to avoid doing that as the Environment variables necessarily have AIRFLOW_SECTION__ as a prefix to them if we do that. Our common files already have specific environment variable names without the prefix of AIRFLOW__SECTION__. But it seems like there is no direct way to set them as we like it. So finally we have decided to go ahead with this approach and change our common files. – Bimonthly

perhaps this one is better: docs.aws.amazon.com/mwaa/latest/userguide/… – Lammergeier

Already tried this. The ENV variables created here work for the dag files but not for non-dag files. – Bimonthly

Answer 3

You could also create a custom plugin that generates runtime environment variables.

Link Ref: https://docs.aws.amazon.com/mwaa/latest/userguide/samples-env-variables.html

Highstrung

Answer 4

I fixed it by setting them in the configuration options as such:

  airflow_configuration_options = {
    "var.my_var" = "my_val"
  }

then you can get them by:

my_var = Variable.get('MY_VAR')

Gutshall

Answer 5

To summarize, there are a few ways to set environment variables in MWAA:

1. In the startup.sh script.

Write a script named startup.sh and add the necessary environment variables using the export command:

#!/bin/bash
export MY_VAR1=value1
export MY_VAR2=value2
export FOO_USER=bar
export FOO_BAR=startup.sh

Upload the script to S3 and configure MWAA to use it. This approach is ideal for setting environment variables as it is simple to implement and manage.

2. Using MWAA configuration options in MWAA Console

In the MWAA console, you can set custom configuration options, e.g. foo.user -> bar, which will then be exposed as the AIRFLOW__FOO__USER environment variable. This is the easiest method, but you can’t customize the name of the variable; it will always have the AIRFLOW prefix and underscores.

3. Creating a Custom Airflow plugin. This approach offers the greatest flexibility and customization, but it is also the most complex to implement and maintain. For instance, you can fetch values from AWS Secrets Manager or AWS Parameter Store. By creating a custom Airflow plugin, you can programmatically set environment variables or perform other initialization tasks when Airflow starts.

---

By the way, you can use an Airflow plugin that I’ve developed called mwaa-env-var-plugin. It’s designed to give a little boost to the second approach I mentioned earlier. This plugin converts MWAA configuration options into friendly named environment variables. For example:

env.foo_bar -> $FOO_BAR

Let me know if you need any further adjustments!

Backpedal