You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by mo...@apache.org on 2024/05/21 13:50:04 UTC
(airflow) 01/03: openlineage: execute extraction and message sending in separate process
This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch openlineage-processexecution
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 80e6d94836f85f6da4a0df98f3a781b8560f8566
Author: Maciej Obuchowski <ob...@gmail.com>
AuthorDate: Wed May 15 16:09:46 2024 +0200
openlineage: execute extraction and message sending in separate process
Signed-off-by: Maciej Obuchowski <ob...@gmail.com>
---
airflow/providers/openlineage/plugins/listener.py | 50 +++++++++++++++++++++--
1 file changed, 46 insertions(+), 4 deletions(-)
diff --git a/airflow/providers/openlineage/plugins/listener.py b/airflow/providers/openlineage/plugins/listener.py
index 73f8c8c79e..ef736308f7 100644
--- a/airflow/providers/openlineage/plugins/listener.py
+++ b/airflow/providers/openlineage/plugins/listener.py
@@ -17,13 +17,16 @@
from __future__ import annotations
import logging
+import os
from concurrent.futures import ProcessPoolExecutor
from datetime import datetime
from typing import TYPE_CHECKING
+import psutil
from openlineage.client.serde import Serde
from airflow import __version__ as airflow_version
+from airflow.api_internal.internal_api_call import InternalApiConfig
from airflow.listeners import hookimpl
from airflow.providers.openlineage import conf
from airflow.providers.openlineage.extractors import ExtractorManager
@@ -71,7 +74,7 @@ class OpenLineageListener:
self,
previous_state,
task_instance: TaskInstance,
- session: Session, # This will always be QUEUED
+ session: Session,
):
if not getattr(task_instance, "task", None) is not None:
self.log.warning(
@@ -153,7 +156,20 @@ class OpenLineageListener:
len(Serde.to_json(redacted_event).encode("utf-8")),
)
- on_running()
+ pid = os.fork()
+ if pid:
+ process = psutil.Process(pid)
+ process.wait(5)
+ else:
+ if not InternalApiConfig.get_use_internal_api():
+ # Force a new SQLAlchemy session. We can't share open DB handles
+ # between process. The cli code will re-create this as part of its
+ # normal startup
+ from airflow import settings
+
+ settings.engine.pool.dispose()
+ settings.engine.dispose()
+ on_running()
@hookimpl
def on_task_instance_success(self, previous_state, task_instance: TaskInstance, session):
@@ -215,7 +231,20 @@ class OpenLineageListener:
len(Serde.to_json(redacted_event).encode("utf-8")),
)
- on_success()
+ pid = os.fork()
+ if pid:
+ process = psutil.Process(pid)
+ process.wait(5)
+ else:
+ if not InternalApiConfig.get_use_internal_api():
+ # Force a new SQLAlchemy session. We can't share open DB handles
+ # between process. The cli code will re-create this as part of its
+ # normal startup
+ from airflow import settings
+
+ settings.engine.pool.dispose()
+ settings.engine.dispose()
+ on_success()
@hookimpl
def on_task_instance_failed(self, previous_state, task_instance: TaskInstance, session):
@@ -277,7 +306,20 @@ class OpenLineageListener:
len(Serde.to_json(redacted_event).encode("utf-8")),
)
- on_failure()
+ pid = os.fork()
+ if pid:
+ process = psutil.Process(pid)
+ process.wait(5)
+ else:
+ if not InternalApiConfig.get_use_internal_api():
+ # Force a new SQLAlchemy session. We can't share open DB handles
+ # between process. The cli code will re-create this as part of its
+ # normal startup
+ from airflow import settings
+
+ settings.engine.pool.dispose()
+ settings.engine.dispose()
+ on_failure()
@property
def executor(self):