How do we set OS environment variables in Airflow

user_2856499 0
Viewed 3271

I have an airflow dag and what im trying to do is read my variables stored in the airflow UI (username and password) and pass those variable values as exported values in the OS. the reason is because I am using dbt yml file which requires me to read the environment variable 'dbt_user'. (the only other way is to set the password in the yaml file which is not secure.

default:
  target: dev
  outputs:
    dev:
      type: snowflake
      account: xxxx

      user: "{{ env_var('dbt_user') }}"

I tried to write a dag which does the bashoperator export but it doesnt seem to set the environment variable.

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime
from airflow.models import Variable
import os


default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2020,8,1),
    'retries': 0
}


with DAG('sample', default_args=default_args, schedule_interval='@once') as dag:
    task_1 = BashOperator(
        task_id='get_variables',
        bash_command='export dbt_user={{ var.value.dbt_user }} ',
        env = os.environ.copy(),
        dag=dag
    )

    task_2 = BashOperator(
        task_id='load_seed_data_once',
        bash_command='echo $dbt_user',
        dag=dag
    )

task_1 >> task_2

when I tried to echo we can see nothing is being set. Does anyone know how to set an environment variable using the bashoperator?

[2021-11-04 12:00:34,452] {subprocess.py:63} INFO - Running command: ['bash', '-c', 'echo $dbt_user']
[2021-11-04 12:00:34,463] {subprocess.py:74} INFO - Output:
[2021-11-04 12:00:34,464] {subprocess.py:78} INFO - 
[2021-11-04 12:00:34,465] {subprocess.py:82} INFO - Command exited with return code 0
[2021-11-04 12:00:34,494] {taskinstance.py:1211} INFO - Marking task as SUCCESS. dag_id=sample, task_id=load_seed_data_once, execution_date=20211104T120032, start_date=20211104T120034, end_date=20211104T120034
[2021-11-04 12:00:34,517] {taskinstance.py:1265} INFO - 0 downstream tasks scheduled from follow-on schedule check
[2021-11-04 12:00:34,555] {local_task_job.py:149} INFO - Task exited with return code 0

update:

I also tried doing via the python operator but it didnt work as well. It gave me a raise KeyError(key) from None KeyError: 'variable_1'

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime
from airflow.models import Variable
import os


default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2020,8,1),
    'retries': 0
}

def set_env():
    os.environ["variable_1"] = "value_1"

def print_env_var():
    print(os.environ["variable_1"])


with DAG('sample', default_args=default_args, schedule_interval='@once') as dag:
    set_env_task = PythonOperator(
        task_id='python_task', 
        python_callable=set_env,
        dag=dag
    )



    print_env_task = PythonOperator(
        task_id='load_seed_data_once',
        python_callable=print_env_var,
        dag=dag
    )

set_env_task >> print_env_task
airflow
1 Answers

The BashOperator and PythonOperator - any operator I think - start a new child shell and will only inherit the environment set in the container in the image, at runtime (e.g. compose or k8s deploy) or by a script before starting airflow (e.g. entrypoint).

That's why you have the env parameter in BashOperator to pass whatever dict of env vars you want to set for the script. You can pass the dbt_user and password from the AF variables there too as env is templated.


env={'dbt_user': '{{ var.value.dbt_user }}'}

You can set env in the dag defaults too to make it available to all tasks so you don't need to set individually.

Lastly, if you use the LocalExecutor, you can instead do the following in the first bash :

echo "export dbt_user={{ var.value.dbt_user }} >> ~/.bashrc

It will make the exported vars accessible in any new shell. Note this would not work with KubernetesExecutor as a new container is started - but there are ways around it.

user_65113610
Related