how to use airflow connection as environment variables in python code

Question

Does anyone know how to access airflow environment variable using AIRFLOW_CONN_ and use in the python code. I know we can use hook to get the password, but have been trying to use AIRFLOW_CONN in my python to connect to the database. I have saved the connection in Airflow UI and in the docs, they mentioned to use AIRFLOW_CONN_ prefix to the conn_id to use. I used it in my python code using os.environ['AIRFLOW_CONN_REDSHIFT'], but it does not identify the environment variable. Please help.

Do you mean to save connection object as environment variables? — Chengzhi
I have saved the connection object in Airflow UI. Now I want to use them in the code. — NNN

Daniel Huang Daniel Huang · Accepted Answer · 2017-09-15T17:35:19

Saving the connection to database and setting an AIRFLOW_CONN_ environment variable are two different ways to add a connection. You should only choose one way, unless you want them stored under connection ids.

Assuming you are running your python code through an operator like PythonOperator, you should be able to fetch your connection just like the BaseHook does.

Stored in database:

@classmethod
def _get_connections_from_db(cls, conn_id):
    session = settings.Session()
    db = (
        session.query(Connection)
        .filter(Connection.conn_id == conn_id)
        .all()
    )
    session.expunge_all()
    session.close()
    if not db:
        raise AirflowException(
            "The conn_id `{0}` isn't defined".format(conn_id))
    return db

Stored in environment variable:

@classmethod
def _get_connection_from_env(cls, conn_id):
    environment_uri = os.environ.get(CONN_ENV_PREFIX + conn_id.upper())
    conn = None
    if environment_uri:
        conn = Connection(conn_id=conn_id, uri=environment_uri)
    return conn

Although I would recommend fetching it via a hook to avoid duplicating this code!

how to use airflow connection as environment variables in python code

1 Answers