astro-sdk
astro-sdk copied to clipboard
Remove the XCOM pull and push for Open lineage
Describe the bug
Currently we are pushing the values to xcom to get the values in extract_on_complete for open lineage. This is a temporary hack till the issue is not fixed on airflow side https://github.com/apache/airflow/pull/21157
This work after changing the following line from taskinstance.py:
session.merge(self)
to
session.merge(self).task = self.task
With this, I was able to get it to work
task_instance.task.sql
The only change then I have in python-sdk/src/astro/lineage/extractor.py is :
Before:
self.log.debug(self.operator.sql)
After:
self.log.debug(task_instance.task.sql)
Without that change, it didn’t work because OpenLineage saves the “task” before running (before calling the op.execute method but after TI is set to Running state) the task in memory - this line in on_task_instance_running
Expected behavior Remove xcom push and pull on operator