conn_deco.py

from datetime import timedelta

from airflow.decorators import dag, task
from airflow.operators.python import PythonOperator
import pendulum

from decors import get_connection, remove, setup

def_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "email_on_failure": False,
    "email_on_retry": False,
    "retries": 1,
    "retry_delay": timedelta(minutes=5),
}


@dag(
    default_args=def_args,
    schedule=None,
    start_date=pendulum.today('UTC'),
    tags=["example"],
)
def conn_decorator():
    @task()
    def doing_nothing(conn_id, **kwargs):
        print(f"Using connection {conn_id}")

        ssh_hook = get_connection(conn_id=conn_id, **kwargs)
        with ssh_hook.get_conn() as ssh_client:
            sftp_client = ssh_client.open_sftp()
            print("Connected")
            lst = sftp_client.listdir(path="/tmp/")
            for f in lst:
                print(f)

        return conn_id

    conn_id = PythonOperator(python_callable=setup, task_id="setup_connection")
    # another way of mixing taskflow and classical api:
    a_id = conn_id.output["return_value"]
    dno = doing_nothing(conn_id=a_id)
    en = PythonOperator(
        python_callable=remove, op_kwargs={"conn_id": dno}, task_id="cleanup"
    )

    conn_id >> dno >> en


dag = conn_decorator()