CODEWITHSUNDEEP
airflow
Nick
Nick

Reputation: 91

How do I add a new connection type to Airflow?

Looking at the plugin documentation (https://airflow.incubator.apache.org/plugins.html), it is clear how to add to new hooks and operators, but I am adding a new hook that requires connection information. This information seems to be hard-coded in airflow/models.py. Is there a way to add my own connection type to the list without altering Airflow's source code?

Upvotes: 9

Views: 7155

Answers (1)

Han-Jong Ko
Han-Jong Ko

Reputation: 582

airflow Connection's conn_type field allows null value. so if you don't care about giving unique type name to your custom hook, then you can give a default connection value in your hook implementation. 

from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook
from airflow.utils.db import provide_session 

class MyHook(BaseHook):

    # ... impl whatever you want.        

    @classmethod
    @provide_session
    def get_hook(cls, conn_id='myhook_default', session=None):
        try:
            conn = cls.get_connection(conn_id)
        except AirflowException:  
            # create default connection. run only once.
            conn = Connection(
                conn_id=conn_id,
                # conn_type='string500',  # You can give new type string here. But no UI component's for you. Just leave it.
                host='default.example.com',
                port=80,
                login='default_login',
                password='default_pw',
                extra=json.dumps({...extra_defult_you_need...}),
            )
            session.add(conn)
            session.commit()
    return MyHook(conn=conn)

Upvotes: 3

Related Questions