CODEWITHSUNDEEP
pythonairflow
JChao
JChao

Reputation: 2309

Is there a way to create/modify connections through Airflow API

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.

Upvotes: 41

Views: 32365

Answers (5)

Bruno Campos
Bruno Campos

Reputation: 666

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

Upvotes: 15

Zeno
Zeno

Reputation: 351

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')

Upvotes: 2

Tamlyn
Tamlyn

Reputation: 23544

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.

Upvotes: 6

Taylor D. Edmiston
Taylor D. Edmiston

Reputation: 13016

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.

Upvotes: 16

mad_
mad_

Reputation: 8273

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.

Upvotes: 57

Related Questions