CODEWITHSUNDEEP
python-2.7airflow
NNN
NNN

Reputation: 11

how to use airflow connection as environment variables in python code

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.

Upvotes: 1

Views: 8458

Answers (1)

Daniel Huang
Daniel Huang

Reputation: 6548

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!

Upvotes: 1

Related Questions