kacpermuda commented on code in PR #65788:
URL: https://github.com/apache/airflow/pull/65788#discussion_r3160470051
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -164,6 +173,144 @@ def get_job_name(task: TaskInstance |
RuntimeTaskInstance) -> str:
return f"{task.dag_id}.{task.task_id}"
+def get_task_instance_from_context():
+ """Return the task instance from the currently executing Airflow task
context."""
+ try:
+ from airflow.sdk import get_current_context
+ except (ImportError, AttributeError):
+ from airflow.operators.python import get_current_context # type:
ignore[no-redef]
+
+ task_instance = get_current_context().get("task_instance")
+ if task_instance is None:
+ raise RuntimeError(
+ "No Airflow task instance found in the current context. "
+ "Call this function from within a running Airflow task, "
+ "or pass `task_instance=` explicitly."
+ )
+ return task_instance
+
+
+def get_dag_run_dag_and_task_from_ti(task_instance):
+ """
+ Return a ``(dag_run, dag, task)`` tuple for the given task instance.
+
+ On Airflow 3+, ``task_instance.get_template_context()`` exposes all three
objects.
+ On Airflow 2, we fall back to the Ti attributes.
+ """
+ if AIRFLOW_V_3_0_PLUS:
+ context = task_instance.get_template_context()
+ return context["dag_run"], context["dag"], context["task"]
+ task = task_instance.task
+ return task_instance.dag_run, task.dag, task
+
+
+def build_task_event_run_facets(
Review Comment:
Yes, definitely, I have a separate ticket to clean up the facet generation
logic in the other parts of the code as well, as it's scattered everywhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
