SQLAlchemy relationship with task_instance table in airflow

user_3804955 0
Viewed 1338

I'm using airflow and I want to be able to track all files generated by given task instances in a table airflow.file_list, which is part of the same database used by airflow (running on postgres). Using SQLAlchemy I have the following mapper for my file_list table:

from airflow.models import Base

class MySourceFile(Base):
    """ SQLAlchemy mapper class for the file_list table entries."""
    __table__ = Table('file_list', Base.metadata,
        Column('UID', Integer, primary_key=True),
        Column('task_id', String(_ID_LEN), nullable=False),
        Column('dag_id', String(_ID_LEN), nullable=False),
        Column('execution_date', DateTime, nullable=False),
        Column('file_path', String(_ID_LEN), nullable=False),
        Column('file_sha256', String(_ID_LEN), nullable=False),
        ForeignKeyConstraint(
            ['task_id', 'dag_id', 'execution_date'],
            ['task_instance.task_id', 'task_instance.dag_id', 'task_instance.execution_date']
        ),
        extend_existing=True,
    )

    instance_task = relationship(
    TaskInstance,
    primaryjoin=and_(
        TaskInstance.task_id == __table__.c.task_id,
        TaskInstance.dag_id == __table__.c.dag_id,
        TaskInstance.execution_date == __table__.c.execution_date
    ),
    viewonly=True,
    foreign_keys=[__table__.c.task_id, __table__.c.dag_id, __table__.c.execution_date]
)

I'm importing the declarative base from airflow.modles because I've read that it's necessary for interacting mappers to share the same instance of base. In the above codes-snippet I want instance_task to reference the task_instance that created the file. The table columns task_id, dag_id and execution_date in my table airflow.file_list mirror the primary keys in airflow.task_instance. Unfortunately, when I run the airflow server, I'm getting the following error:

sqlalchemy.exc.InvalidRequestError: One or more mappers failed to initialize - can't proceed with initialization of other mappers. Triggering mapper: 'Mapper|MySourceFile|file_list'. Original exception was: Can't determine relationship direction for relationship 'MySourceFile.instance_task' - foreign key columns are present in neither the parent nor the child's mapped tables

I'd prefer not to modify the airflow source if possible. Thank you in advance for any help.

airflow sqlalchemy python
1 Answers

First off specifying __table__ like that is odd (I've never seen it before), and as a general rule anything starting with and underscore is private and should be avoided.

You need to specify that task_id is a foreign key. Something like this:

task_id = Column(String, ForeignKey('task.id'))

The docs on SQL Alchemy relationships are useful if you haven't seem them. And the backref pattern might be useful to create a relationship from task instance down to your custom file without having to alter the TaskInstance model.

user_4391890
Related