36

Going through Admin -> Connections, we have the ability to create/modify a connection's params, but I'm wondering if I can do the same through API so I can programmatically set the connections

airflow.models.Connection seems like it only deals with actually connecting to the instance instead of saving it to the list. It seems like a function that should have been implemented, but I'm not sure where I can find the docs for this specific function.

JChao
  • 2,178
  • 5
  • 35
  • 65

5 Answers5

52

Connection is actually a model which you can use to query and insert a new connection 

from airflow import settings
from airflow.models import Connection
conn = Connection(
        conn_id=conn_id,
        conn_type=conn_type,
        host=host,
        login=login,
        password=password,
        port=port
) #create a connection object
session = settings.Session() # get the session
session.add(conn)
session.commit() # it will insert the connection object programmatically.
mad_
  • 8,121
  • 2
  • 25
  • 40
  • 1
    Thanks for the quick response. This is exactly what I'm looking for. I'll accept this answer once stackoverflow allows me – JChao Aug 15 '18 at 17:57
  • 11
    is there anyway to list the connection first and then check if connection already exists or not? – Gagan Jun 07 '19 at 00:15
  • @mad_ I thought may be using Settings above, we can access connection list and check from it. Yes we can directly use bashoperator and then airflow connections -l but I was not very much into creating another tasks for it – Gagan Jun 07 '19 at 19:37
  • 1
    How to delete a connection from this session – Anoop Mar 11 '21 at 11:04
  • Use the same model `Connection ` to delete as well – mad_ Mar 12 '21 at 18:37
  • What about closing the session here? Don't we need session.close()? – Vivek Kumar Jul 26 '23 at 09:14
16

You can also add, delete, and list connections from the Airflow CLI if you need to do it outside of Python/Airflow code, via bash, in a Dockerfile, etc.

airflow connections --add ...

Usage:

airflow connections [-h] [-l] [-a] [-d] [--conn_id CONN_ID]
                    [--conn_uri CONN_URI] [--conn_extra CONN_EXTRA]
                    [--conn_type CONN_TYPE] [--conn_host CONN_HOST]
                    [--conn_login CONN_LOGIN] [--conn_password CONN_PASSWORD]
                    [--conn_schema CONN_SCHEMA] [--conn_port CONN_PORT]

https://airflow.apache.org/cli.html#connections

It doesn't look like the CLI currently supports modifying an existing connection, but there is a Jira issue for it with an active open PR on GitHub.

Taylor D. Edmiston
  • 12,088
  • 6
  • 56
  • 76
12

First check if connection exists, after create new Connection using from airflow.models import Connection :

import logging
from airflow import settings
from airflow.models import Connection

def create_conn(conn_id, conn_type, host, login, pwd, port, desc):
    conn = Connection(conn_id=conn_id,
                      conn_type=conn_type,
                      host=host,
                      login=login,
                      password=pwd,
                      port=port,
                      description=desc)
    session = settings.Session()
    conn_name = session.query(Connection).filter(Connection.conn_id == conn.conn_id).first()

        if str(conn_name) == str(conn.conn_id):
            logging.warning(f"Connection {conn.conn_id} already exists")
            return None

    session.add(conn)
    session.commit()
    logging.info(Connection.log_info(conn))
    logging.info(f'Connection {conn_id} is created')
    return conn
Bruno Campos
  • 595
  • 1
  • 7
  • 18
4

You can populate connections using environment variables using the connection URI format.

The environment variable naming convention is AIRFLOW_CONN_<conn_id>, all uppercase.

So if your connection id is my_prod_db then the variable name should be AIRFLOW_CONN_MY_PROD_DB.

In general, Airflow’s URI format is like so:

my-conn-type://my-login:my-password@my-host:5432/my-schema?param1=val1&param2=val2

Note that connections registered in this way do not show up in the Airflow UI.

Tamlyn
  • 22,122
  • 12
  • 111
  • 127
1

To use session = settings.Session(), it assumes the airflow database backend has been initiated. For those who haven't set it up for your development environment, a hybrid method using both Connection class and environment variables will be a workaround.

Below is the example for setting up a S3Hook

from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.models.connection import Connection
import os
import json

aws_default = Connection(
    conn_id="aws_default",
    conn_type="aws",
    login='YOUR-AWS-KEY-ID',
    password='YOUR-AWS-KEY-SECRET',
    extra=json.dumps({'region_name': 'us-east-1'})
    )

os.environ["AIRFLOW_CONN_AWS_DEFAULT"] = aws_default.get_uri()
s3_hook = S3Hook(aws_conn_id='aws_default')
s3_hook.list_keys(bucket_name='YOUR-BUCKET', prefix='YOUR-FILENAME')
Zeno
  • 331
  • 2
  • 5