Problem updating the connections in Airflow programatically

前端 未结 1 1403
隐瞒了意图╮
隐瞒了意图╮ 2020-12-20 07:55

I am trying to update the Airflow connections using python. I have created a python function that takes an authentication token from API and updates the extra field of conne

1条回答
  •  生来不讨喜
    2020-12-20 08:43

    I doubt that you are fetching connection from Airflow's meta-db in the right way.

    • All things apart, if you are fetching Variable via Variable.get() method, shouldn't Connection be receiving the same treatment (although Connection class doesn't have a get() function, there must be a workaround)?
    • Here you are merely instantiating Connection object with a given conn_id argument (and not really fetching Connection for that conn_id from db)

    Whenever I have to exploit the underlying SQLAlchemy models, I look at cli.py. Taking cues from connections() function, here's what I think should work

    from airflow.models import Connection
    from airflow.settings import Session
    from airflow.utils.db import provide_session
    from typing import List, Dict, Any, Optional
    from sqlalchemy.orm import exc
    
    @provide_session
    def update_conn_extra(conn_id: str, new_extra: Any, session: Optional[Session] = None) -> Optional[Connection]:
        try:
            my_conn: Optional[Connection] = (session
                                             .query(Connection)
                                             .filter(Connection.conn_id == conn_id)
                                             .one())
        except exc.NoResultFound:
            my_conn: Optional[Connection] = None
        except exc.MultipleResultsFound:
            my_conn: Optional[Connection] = None
        if my_conn:
            my_conn.extra: Any = new_extra
            session.add(my_conn)
            session.commit()
    

    Note that here we are simple overwriting the Connection with updated fields (without first deleting the existing one), which I've found to work. If you face some problems, you can delete the existing connection before writing updated one using session.delete(my_conn)

    0 讨论(0)
提交回复
热议问题