This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 69f98974344 Look up task logs by task try UUID (#74204)
69f98974344 is described below
commit 69f98974344dec7de5a70c91a9395bc72163687c
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Mon Oct 5 23:44:34 2026 +0100
Look up task logs by task try UUID (#74204)
These handlers previously identified a task's logs by coordinates (dag,
task, run, map index, try number). Once a task can run in several loop
passes that share those coordinates, the identity has to come from
something that is unique per attempt, and with the ownership change
every try already has its own UUID.
Elasticsearch and OpenSearch entries now carry the task instance's id as
``ti_id`` and Cloud Logging entries carry it as a label of the same
name, next to the existing log_id or labels, so the human-readable
identity and any Kibana or Logs Explorer, etc. queries built on them
continue to work. Readers match on that id alone, or fall back to the
log_id or labels for entries that have no ``ti_id``, so logs written
before this change stay readable without reindexing.
``ti_id`` is where the task process has bound the UUID in to its own log
lines since 3.3, so entries from 3.3 onward already match. Nothing
earlier wrote it, and most lines the supervisor logs itself do not
carry it either, so the Elasticsearch and OpenSearch uploaders add it to
every line they ship.
Readers only match on the id from Airflow 3.4. Before then clearing a
task instance gives the live row a new id, so it may no longer be the id
its logs were written under, and earlier versions keep matching on
log_id or labels alone. Downloading every try on Google Cloud Logging
drops the id condition just as it already drops the try number.
---
.../kubernetes/executors/kubernetes_executor.py | 3 +
.../executors/kubernetes_executor_types.py | 1 +
.../executors/kubernetes_executor_utils.py | 2 +
.../executors/test_kubernetes_executor.py | 29 ++++++
providers/elasticsearch/docs/logging/index.rst | 5 +
.../providers/elasticsearch/log/es_task_handler.py | 43 +++++++-
.../providers/elasticsearch/version_compat.py | 2 +
.../elasticsearch/log/test_es_remote_log_io.py | 32 ++++++
.../unit/elasticsearch/log/test_es_task_handler.py | 44 +++++++-
.../google/cloud/log/cloud_logging_task_handler.py | 26 ++++-
.../src/airflow/providers/google/version_compat.py | 9 +-
.../cloud/log/test_cloud_logging_task_handler.py | 112 +++++++++++++++++----
.../providers/opensearch/log/os_task_handler.py | 43 ++++++--
.../airflow/providers/opensearch/version_compat.py | 2 +
.../opensearch/log/test_os_remote_log_io.py | 32 ++++++
.../unit/opensearch/log/test_os_task_handler.py | 48 ++++++++-
16 files changed, 397 insertions(+), 36 deletions(-)
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
index 819c404da15..f19a64b08dd 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
@@ -52,6 +52,7 @@ from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types impor
ADOPTED,
POD_EXECUTOR_DONE_KEY,
TASK_INSTANCE_ID_ANNOTATION,
+ TASK_INSTANCE_ID_LABEL,
FailureDetails,
KubernetesJob,
KubernetesResults,
@@ -1042,6 +1043,8 @@ class KubernetesExecutor(BaseExecutor):
run_id=ti.run_id,
airflow_worker=ti.queued_by_job_id,
)
+ if self.supports_task_instance_uuid:
+ selector += f",{TASK_INSTANCE_ID_LABEL}={ti.id}"
namespace = self._get_pod_namespace(ti)
pod_list = client.list_namespaced_pod(
namespace=namespace,
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
index 3f41f398228..2eea8f07125 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py
@@ -102,6 +102,7 @@ So we want events on a revoked pod to be ignored.
TASK_INSTANCE_ID_ANNOTATION = "task_instance_id"
+TASK_INSTANCE_ID_LABEL = "ti_id"
def task_instance_id_from_pod(pod: k8s.V1Pod) -> UUID | None:
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
index ce574867191..576a6fe796f 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
@@ -38,6 +38,7 @@ from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types impor
POD_EXECUTOR_DONE_KEY,
POD_REVOKED_KEY,
TASK_INSTANCE_ID_ANNOTATION,
+ TASK_INSTANCE_ID_LABEL,
FailureDetails,
KubernetesJob,
KubernetesResults,
@@ -658,6 +659,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
)
if self.supports_task_instance_uuid and isinstance(key,
TaskInstanceUuid):
pod.metadata.annotations[TASK_INSTANCE_ID_ANNOTATION] = str(key)
+ pod.metadata.labels[TASK_INSTANCE_ID_LABEL] = str(key)
# Reconcile the pod generated by the Operator and the Pod
# generated by the .cfg file
self.log.info(
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
index 7baff0322c5..84e03ce35fb 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
@@ -47,6 +47,7 @@ from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor import (
)
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types
import (
ADOPTED,
+ TASK_INSTANCE_ID_LABEL,
KubernetesJob,
KubernetesResults,
KubernetesWatch,
@@ -4366,6 +4367,7 @@ class TestKubernetesExecutorUuid:
)
assert pod.metadata.annotations["task_instance_id"] ==
str(uuid_workload.ti.id)
+ assert pod.metadata.labels[TASK_INSTANCE_ID_LABEL] ==
str(uuid_workload.ti.id)
assert pod.metadata.annotations["dag_id"] == "uuid_dag"
assert pod.metadata.annotations["try_number"] == "1"
@@ -4396,6 +4398,33 @@ class TestKubernetesExecutorUuid:
assert statement.compile().params == {"id_1": uuid_workload.ti.id}
assert state == State.RESTARTING
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_get_streaming_task_log_selects_pod_by_attempt_uuid_label(self,
mock_get_kube_client):
+ ti_id = uuid4()
+ mock_kube_client = mock_get_kube_client.return_value
+ mock_kube_client.list_namespaced_pod.return_value.items = [
+ k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="current"))
+ ]
+ mock_kube_client.read_namespaced_pod_log.return_value = [b"a_"]
+ ti = mock.MagicMock(
+ id=ti_id,
+ try_number=2,
+ dag_id="test_k8s_log_dag",
+ task_id="test_task",
+ map_index=-1,
+ run_id="test_run",
+ queued_by_job_id=None,
+ hostname="",
+ executor_config={},
+ )
+
+ messages, _ = KubernetesExecutor().get_streaming_task_log(ti=ti,
try_number=2)
+
+ assert messages[-1] == "Found logs through kube API"
+ label_selector =
mock_kube_client.list_namespaced_pod.call_args.kwargs["label_selector"]
+ assert f"{TASK_INSTANCE_ID_LABEL}={ti_id}" in label_selector.split(",")
+ assert
mock_kube_client.read_namespaced_pod_log.call_args.kwargs["name"] == "current"
+
def test_revoke_selects_uuid_among_pods_with_reused_coordinates(self,
uuid_workload, mocker):
executor = KubernetesExecutor()
executor.kube_scheduler =
mocker.create_autospec(AirflowKubernetesScheduler, instance=True)
diff --git a/providers/elasticsearch/docs/logging/index.rst
b/providers/elasticsearch/docs/logging/index.rst
index 2e8d732ea3f..8fd3883f705 100644
--- a/providers/elasticsearch/docs/logging/index.rst
+++ b/providers/elasticsearch/docs/logging/index.rst
@@ -294,6 +294,11 @@ The following fields are recognized and displayed by the
Airflow UI when present
* - ``host``
- The hostname of the worker that produced the log. Used to group log
lines by source.
- Field name can be customized with the ``host_field`` handler parameter.
+ * - ``ti_id``
+ - UUID of the task try that produced the log line.
+ - Written by the Airflow 3 task process, so a log shipper forwards it
from the task's JSON output
+ without extra configuration. From Airflow 3.4 logs are read by this
field, falling back to
+ ``log_id`` for documents that lack it.
**Field mappings**
diff --git
a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
index 2fd81a541c5..c70e0800137 100644
---
a/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
+++
b/providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -49,7 +49,11 @@ from airflow.providers.common.compat.sdk import conf,
timezone
from airflow.providers.elasticsearch._compat import apply_compat_with
from airflow.providers.elasticsearch.log.es_json_formatter import
ElasticsearchJSONFormatter
from airflow.providers.elasticsearch.log.es_response import
ElasticSearchResponse, Hit, resolve_nested
-from airflow.providers.elasticsearch.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
+from airflow.providers.elasticsearch.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
from airflow.utils.log.file_task_handler import FileTaskHandler
from airflow.utils.log.logging_mixin import ExternalLoggingMixin, LoggingMixin
from airflow.utils.state import TaskInstanceState
@@ -226,6 +230,31 @@ def _render_log_id(log_id_template: str, ti: TaskInstance
| TaskInstanceKey, try
)
+def _get_ti_id_fields(ti: TaskInstance | TaskInstanceKey) -> dict[str, str]:
+ # Before 3.4 a try can reuse the previous try's id, so only log_id
identifies it.
+ if not AIRFLOW_V_3_4_PLUS:
+ return {}
+ return {"ti_id": str(ti_id)} if (ti_id := getattr(ti, "id", None)) else {}
+
+
+def _build_log_query(log_id: str, ti: RuntimeTI) -> list[dict[str, Any]]:
+ log_id_match = {"match_phrase": {"log_id": log_id}}
+ # Before 3.4 a cleared task instance gets a new id, which can differ from
the id its logs were written under.
+ if not AIRFLOW_V_3_4_PLUS:
+ return [log_id_match]
+ return [
+ {
+ "bool": {
+ "should": [
+ {"match_phrase": {"ti_id": str(ti.id)}},
+ {"bool": {"must": [log_id_match], "must_not": {"exists":
{"field": "ti_id"}}}},
+ ],
+ "minimum_should_match": 1,
+ }
+ }
+ ]
+
+
def _clean_date(value: datetime | None) -> str:
"""
Clean up a date value so that it is safe to query in elasticsearch by
removing reserved characters.
@@ -761,21 +790,24 @@ class ElasticsearchRemoteLogIO(LoggingMixin): # noqa:
D101
local_loc = self.base_log_folder.joinpath(path)
log_id = _render_log_id(self.log_id_template, ti, ti.try_number) #
type: ignore[arg-type]
+ ti_id_fields = _get_ti_id_fields(ti) # type: ignore[arg-type]
if local_loc.is_file() and self.write_stdout:
# Intentionally construct the log_id and offset field
- log_lines = self._parse_raw_log(local_loc.read_text(), log_id)
+ log_lines = self._parse_raw_log(local_loc.read_text(), log_id,
ti_id_fields)
for line in log_lines:
sys.stdout.write(json.dumps(line) + "\n")
sys.stdout.flush()
if local_loc.is_file() and self.write_to_es:
- log_lines = self._parse_raw_log(local_loc.read_text(), log_id)
+ log_lines = self._parse_raw_log(local_loc.read_text(), log_id,
ti_id_fields)
success = self._write_to_es(log_lines)
if success and self.delete_local_copy:
shutil.rmtree(os.path.dirname(local_loc))
- def _parse_raw_log(self, log: str, log_id: str) -> list[dict[str, Any]]:
+ def _parse_raw_log(
+ self, log: str, log_id: str, extra_fields: dict[str, str] | None = None
+ ) -> list[dict[str, Any]]:
logs = log.split("\n")
parsed_logs = []
offset = 1
@@ -801,6 +833,7 @@ class ElasticsearchRemoteLogIO(LoggingMixin): # noqa: D101
log_dict.update(
{
"log_id": log_id,
+ **(extra_fields or {}),
self.offset_field: offset,
}
)
@@ -883,7 +916,7 @@ class ElasticsearchRemoteLogIO(LoggingMixin): # noqa: D101
query: dict[Any, Any] = {
"bool": {
"filter": [{"range": {self.offset_field: {"gt":
int(offset)}}}],
- "must": [{"match_phrase": {"log_id": log_id}}],
+ "must": _build_log_query(log_id, ti),
}
}
diff --git
a/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py
b/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py
index 806dbac18d2..8c63b0494d7 100644
---
a/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py
+++
b/providers/elasticsearch/src/airflow/providers/elasticsearch/version_compat.py
@@ -36,10 +36,12 @@ AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >=
(3, 0, 0)
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
diff --git
a/providers/elasticsearch/tests/integration/elasticsearch/log/test_es_remote_log_io.py
b/providers/elasticsearch/tests/integration/elasticsearch/log/test_es_remote_log_io.py
index e6da3dc7514..8ba26091deb 100644
---
a/providers/elasticsearch/tests/integration/elasticsearch/log/test_es_remote_log_io.py
+++
b/providers/elasticsearch/tests/integration/elasticsearch/log/test_es_remote_log_io.py
@@ -39,6 +39,7 @@ class _MockTI:
run_id: str = "integration_test_run"
try_number: int = 1
map_index: int = -1
+ id: uuid.UUID = dataclasses.field(default_factory=uuid.uuid4)
@pytest.mark.integration("elasticsearch")
@@ -145,3 +146,34 @@ class TestElasticsearchRemoteLogIOIntegration:
log_entry = json.loads(log_messages[0])
assert "error_detail" in log_entry
assert log_entry["error_detail"] == error_detail
+
+ @pytest.mark.parametrize(
+ ("written", "reader", "expected"),
+ [
+ pytest.param(["legacy"], "x", ["legacy"],
id="pre-upgrade-logs-stay-readable"),
+ pytest.param(["x", "y"], "x", ["x"],
id="post-upgrade-reader-sees-only-its-try"),
+ pytest.param(["x", "y"], "y", ["y"], id="post-upgrade-other-try"),
+ pytest.param(["legacy", "x", "y"], "x", ["legacy", "x"],
id="mixed-index"),
+ ],
+ )
+ @patch(
+ "airflow.providers.elasticsearch.log.es_task_handler.TASK_LOG_FIELDS",
+ ["message"],
+ )
+ def test_read_scopes_to_try_and_keeps_pre_upgrade_logs(self, tmp_path,
written, reader, expected):
+ """Tries sharing one log_id stay separate, and entries without ti_id
stay readable."""
+ try_ids = {"x": uuid.uuid4(), "y": uuid.uuid4()}
+ log_id = _render_log_id(self.elasticsearch_io.log_id_template,
_MockTI(), 1)
+ for name in written:
+ if name == "legacy":
+ legacy_entry = {"message": name, "log_id": log_id, "offset": 1}
+ self.elasticsearch_io.client.index(index=self.target_index,
document=legacy_entry)
+ continue
+ log_file = tmp_path / f"{name}.log"
+ log_file.write_text(json.dumps({"message": name}) + "\n")
+ self.elasticsearch_io.upload(log_file, _MockTI(id=try_ids[name]))
+ self.elasticsearch_io.client.indices.refresh(index=self.target_index)
+
+ _, log_messages = self.elasticsearch_io.read("",
_MockTI(id=try_ids[reader]))
+
+ assert sorted(json.loads(message)["event"] for message in
log_messages) == sorted(expected)
diff --git
a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
index 5fa6f1ad6ef..2517975349c 100644
---
a/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
+++
b/providers/elasticsearch/tests/unit/elasticsearch/log/test_es_task_handler.py
@@ -24,9 +24,11 @@ import os
import re
from io import StringIO
from pathlib import Path
+from types import SimpleNamespace
from unittest import mock
from unittest.mock import Mock, patch
from urllib.parse import quote
+from uuid import UUID, uuid4
import elasticsearch
import pendulum
@@ -41,8 +43,10 @@ from airflow.providers.elasticsearch.log.es_task_handler
import (
ElasticsearchRemoteLogIO,
ElasticsearchTaskHandler,
_build_log_fields,
+ _build_log_query,
_clean_date,
_format_error_detail,
+ _get_ti_id_fields,
_render_log_id,
_safe_build_structured_log_message,
_strip_userinfo,
@@ -57,6 +61,41 @@ from tests_common.test_utils.db import clear_db_dags,
clear_db_runs
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
[email protected](
+ ("is_airflow_3_4_plus", "expected"),
+ [(False, {}), (True, {"ti_id": "some-ti-id"})],
+)
+def test_ti_id_is_only_written_from_airflow_3_4(is_airflow_3_4_plus, expected):
+ ti = SimpleNamespace(id="some-ti-id")
+
+ with
patch("airflow.providers.elasticsearch.log.es_task_handler.AIRFLOW_V_3_4_PLUS",
is_airflow_3_4_plus):
+ assert _get_ti_id_fields(ti) == expected
+
+
[email protected]("is_airflow_3_4_plus", [False, True])
+def test_log_query_matches_ti_id_or_documents_without_it(is_airflow_3_4_plus):
+ ti = SimpleNamespace(id=uuid4())
+ log_id_match = {"match_phrase": {"log_id": "some-log-id"}}
+
+ with
patch("airflow.providers.elasticsearch.log.es_task_handler.AIRFLOW_V_3_4_PLUS",
is_airflow_3_4_plus):
+ must = _build_log_query("some-log-id", ti)
+
+ if not is_airflow_3_4_plus:
+ assert must == [log_id_match]
+ return
+ assert must == [
+ {
+ "bool": {
+ "should": [
+ {"match_phrase": {"ti_id": str(ti.id)}},
+ {"bool": {"must": [log_id_match], "must_not": {"exists":
{"field": "ti_id"}}}},
+ ],
+ "minimum_should_match": 1,
+ }
+ }
+ ]
+
+
@dataclasses.dataclass
class _MockTI:
dag_id: str = "dag_for_testing_es_log_handler"
@@ -64,6 +103,7 @@ class _MockTI:
run_id: str = "run_for_testing_es_log_handler"
try_number: int = 1
map_index: int = -1
+ id: UUID = dataclasses.field(default_factory=uuid4)
def get_ti(dag_id, task_id, run_id, logical_date, create_task_instance):
@@ -706,6 +746,7 @@ class TestElasticsearchRemoteLogIO:
file_path.write_text("\n".join(json.dumps(log) for log in sample_logs)
+ "\n")
return file_path
+
@patch("airflow.providers.elasticsearch.log.es_task_handler.AIRFLOW_V_3_4_PLUS",
True)
def test_write_to_stdout(self, tmp_json_file, ti, capsys):
self.elasticsearch_io.write_to_es = False
self.elasticsearch_io.upload(tmp_json_file, ti)
@@ -714,6 +755,7 @@ class TestElasticsearchRemoteLogIO:
stdout_lines = captured.out.strip().splitlines()
log_entries = [json.loads(line) for line in stdout_lines]
assert [entry["message"] for entry in log_entries] == ["start",
"processing", "end"]
+ assert {entry["ti_id"] for entry in log_entries} == {str(ti.id)}
def test_invalid_task_log_file_path(self, ti):
with (
@@ -816,7 +858,7 @@ class TestElasticsearchRemoteLogIO:
query = {
"bool": {
"filter": [{"range": {self.elasticsearch_io.offset_field:
{"gt": 2}}}],
- "must": [{"match_phrase": {"log_id": log_id}}],
+ "must": _build_log_query(log_id, ti),
}
}
diff --git
a/providers/google/src/airflow/providers/google/cloud/log/cloud_logging_task_handler.py
b/providers/google/src/airflow/providers/google/cloud/log/cloud_logging_task_handler.py
index 3bc9be0274b..1a9a2b296c7 100644
---
a/providers/google/src/airflow/providers/google/cloud/log/cloud_logging_task_handler.py
+++
b/providers/google/src/airflow/providers/google/cloud/log/cloud_logging_task_handler.py
@@ -46,7 +46,7 @@ from airflow.providers.common.compat.sdk import conf
from airflow.providers.google.cloud.utils.credentials_provider import
get_credentials_and_project_id
from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_4_PLUS
from airflow.utils.log.file_task_handler import FileTaskHandler
from airflow.utils.log.logging_mixin import LoggingMixin
@@ -82,6 +82,7 @@ LABEL_TASK_ID = "task_id"
LABEL_DAG_ID = "dag_id"
LABEL_LOGICAL_DATE = "logical_date" if AIRFLOW_V_3_0_PLUS else "execution_date"
LABEL_TRY_NUMBER = "try_number"
+LABEL_TI_ID = "ti_id"
@attrs.define(kw_only=True)
@@ -209,6 +210,8 @@ class CloudLoggingRemoteLogIO(LoggingMixin):
if ti:
labels.update(_task_instance_to_labels(ti))
else:
+ if ti_id := event.get("ti_id"):
+ labels[LABEL_TI_ID] = str(ti_id)
if dag_id := event.get("dag_id"):
labels[LABEL_DAG_ID] = str(dag_id)
if task_id := event.get("task_id"):
@@ -269,8 +272,18 @@ class CloudLoggingRemoteLogIO(LoggingMixin):
for key, value in self.resource.labels.items():
log_filters.append(f"resource.labels.{escape_label_key(key)}={escape_label_value(value)}")
- for key, value in ti_labels.items():
-
log_filters.append(f"labels.{escape_label_key(key)}={escape_label_value(value)}")
+ label_conditions = {
+ key: f"labels.{escape_label_key(key)}={escape_label_value(value)}"
+ for key, value in ti_labels.items()
+ }
+ ti_id_condition = label_conditions.pop(LABEL_TI_ID, None)
+ # Before 3.4 a cleared task instance gets a new id, which can differ
from the id its logs were written under.
+ if ti_id_condition and AIRFLOW_V_3_4_PLUS:
+ # Entries written before the ``ti_id`` label existed are matched
by the other labels.
+ legacy_conditions = [*label_conditions.values(), f"NOT
labels.{LABEL_TI_ID}:*"]
+ log_filters.append(f"({ti_id_condition} OR ({' AND
'.join(legacy_conditions)}))")
+ else:
+ log_filters.extend(label_conditions.values())
return "\n".join(log_filters)
def read_logs(
@@ -317,7 +330,7 @@ class CloudLoggingRemoteLogIO(LoggingMixin):
def _task_instance_to_labels(ti) -> dict[str, str]:
"""Convert a task instance to Cloud Logging labels."""
- return {
+ labels = {
LABEL_TASK_ID: ti.task_id,
LABEL_DAG_ID: ti.dag_id,
LABEL_LOGICAL_DATE: str(ti.logical_date.isoformat())
@@ -325,6 +338,10 @@ def _task_instance_to_labels(ti) -> dict[str, str]:
else str(ti.execution_date.isoformat()),
LABEL_TRY_NUMBER: str(ti.try_number),
}
+ # Before 3.4 a try can reuse the previous try's id, so only the other
labels identify it.
+ if AIRFLOW_V_3_4_PLUS and (ti_id := getattr(ti, "id", None)):
+ labels[LABEL_TI_ID] = str(ti_id)
+ return labels
class CloudLoggingTaskHandler(logging.Handler):
@@ -472,6 +489,7 @@ class CloudLoggingTaskHandler(logging.Handler):
ti_labels[LABEL_TRY_NUMBER] = str(try_number)
else:
del ti_labels[LABEL_TRY_NUMBER]
+ ti_labels.pop(LABEL_TI_ID, None)
log_filter = self.io.prepare_log_filter(ti_labels)
next_page_token = metadata.get("next_page_token", None)
diff --git a/providers/google/src/airflow/providers/google/version_compat.py
b/providers/google/src/airflow/providers/google/version_compat.py
index 57d8840c9ef..bb32aa84ff2 100644
--- a/providers/google/src/airflow/providers/google/version_compat.py
+++ b/providers/google/src/airflow/providers/google/version_compat.py
@@ -36,6 +36,7 @@ AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3,
0, 0)
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
# BaseTrigger.on_kill() was introduced in 3.3.0; use this flag to guard
trigger migrations
AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
# Version-compatible imports
# BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator
3.0.x
@@ -49,4 +50,10 @@ else:
# Explicitly export these imports to protect them from being removed by linters
-__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_3_PLUS",
"BaseOperator"]
+__all__ = [
+ "AIRFLOW_V_3_0_PLUS",
+ "AIRFLOW_V_3_1_PLUS",
+ "AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
+ "BaseOperator",
+]
diff --git
a/providers/google/tests/unit/google/cloud/log/test_cloud_logging_task_handler.py
b/providers/google/tests/unit/google/cloud/log/test_cloud_logging_task_handler.py
index b360f9c2525..bdbb9889f77 100644
---
a/providers/google/tests/unit/google/cloud/log/test_cloud_logging_task_handler.py
+++
b/providers/google/tests/unit/google/cloud/log/test_cloud_logging_task_handler.py
@@ -20,6 +20,7 @@ import logging
import warnings
from contextlib import nullcontext
from pathlib import Path
+from types import SimpleNamespace
from unittest import mock
from unittest.mock import PropertyMock
from urllib.parse import parse_qs, urlsplit
@@ -33,12 +34,32 @@ from airflow.providers.common.compat.sdk import timezone
from airflow.providers.google.cloud.log.cloud_logging_task_handler import (
CloudLoggingRemoteLogIO,
CloudLoggingTaskHandler,
+ _task_instance_to_labels,
)
from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import clear_db_dags, clear_db_runs
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_4_PLUS
+
+
[email protected](
+ ("is_airflow_3_4_plus", "expect_ti_id"),
+ [(False, False), (True, True)],
+)
+def test_ti_id_label_is_only_written_from_airflow_3_4(is_airflow_3_4_plus,
expect_ti_id):
+ now = timezone.utcnow()
+ ti = SimpleNamespace(
+ id="some-ti-id", task_id="t", dag_id="d", try_number=1,
logical_date=now, execution_date=now
+ )
+
+ with mock.patch(
+
"airflow.providers.google.cloud.log.cloud_logging_task_handler.AIRFLOW_V_3_4_PLUS",
+ is_airflow_3_4_plus,
+ ):
+ labels = _task_instance_to_labels(ti)
+
+ assert ("ti_id" in labels) is expect_ti_id
def _create_list_log_entries_response_mock(messages, token):
@@ -289,6 +310,7 @@ class TestCloudLoggingRemoteLogIO:
"run_id": "test_run_id",
"try_number": 2,
"map_index": -1,
+ "ti_id": "2a4ab344-e2b8-41c9-808f-2411659ba64f",
}
result = proc(logger, "info", event)
@@ -300,6 +322,7 @@ class TestCloudLoggingRemoteLogIO:
labels = kwargs.get("labels", {})
assert labels == {
+ "ti_id": event["ti_id"],
"dag_id": "test_dag_id",
"task_id": "test_task_id",
"run_id": "test_run_id",
@@ -325,6 +348,42 @@ class TestCloudLoggingRemoteLogIO:
assert 'labels.task_id="test_task"' in log_filter
assert 'labels.dag_id="test_dag"' in log_filter
+ @mock.patch(
+
"airflow.providers.google.cloud.log.cloud_logging_task_handler.get_credentials_and_project_id"
+ )
+ @pytest.mark.parametrize(
+ ("is_airflow_3_4_plus", "expected_conditions"),
+ [
+ pytest.param(
+ False,
+ ['labels.task_id="test_task"', 'labels.dag_id="test_dag"',
'labels.try_number="1"'],
+ id="before-3.4",
+ ),
+ pytest.param(
+ True,
+ [
+ '(labels.ti_id="abc" OR (labels.task_id="test_task" AND
labels.dag_id="test_dag" '
+ 'AND labels.try_number="1" AND NOT labels.ti_id:*))'
+ ],
+ id="3.4-or-later",
+ ),
+ ],
+ )
+ def test_prepare_log_filter_matches_ti_id_or_legacy_labels(
+ self, mock_get_creds_and_project_id, is_airflow_3_4_plus,
expected_conditions
+ ):
+ mock_get_creds_and_project_id.return_value = ("creds", "project_id")
+
+ with mock.patch(
+
"airflow.providers.google.cloud.log.cloud_logging_task_handler.AIRFLOW_V_3_4_PLUS",
+ is_airflow_3_4_plus,
+ ):
+ log_filter = self.io.prepare_log_filter(
+ {"task_id": "test_task", "dag_id": "test_dag", "try_number":
"1", "ti_id": "abc"}
+ )
+
+ assert log_filter.splitlines()[2:] == expected_conditions
+
@mock.patch(
"airflow.providers.google.cloud.log.cloud_logging_task_handler.get_credentials_and_project_id"
)
@@ -491,6 +550,15 @@ class TestCloudLoggingHandlerTask:
clear_db_runs()
clear_db_dags()
+ def _ti_id_labels(self) -> dict[str, str]:
+ return {"ti_id": str(self.ti.id)} if AIRFLOW_V_3_4_PLUS else {}
+
+ def _expect_label_filter(self, *label_conditions: str) -> str:
+ if not AIRFLOW_V_3_4_PLUS:
+ return "\n".join(label_conditions)
+ legacy_conditions = " AND ".join([*label_conditions, "NOT
labels.ti_id:*"])
+ return f'(labels.ti_id="{self.ti.id}" OR ({legacy_conditions}))'
+
def _setup_handler(self, **handler_kwargs):
self.transport_mock = mock.MagicMock()
handler_kwargs = {"transport": self.transport_mock, **handler_kwargs}
@@ -518,6 +586,7 @@ class TestCloudLoggingHandlerTask:
"dag_id": self.DAG_ID,
date_key: "2016-01-01T00:00:00+00:00",
"try_number": "1",
+ **self._ti_id_labels(),
}
resource = Resource(type="global", labels={})
self.transport_mock.return_value.send.assert_called_once_with(
@@ -545,6 +614,7 @@ class TestCloudLoggingHandlerTask:
"dag_id": self.DAG_ID,
date_key: "2016-01-01T00:00:00+00:00",
"try_number": "1",
+ **self._ti_id_labels(),
"product.googleapis.com/task_id": "test-value",
}
resource = Resource(type="global", labels={})
@@ -636,10 +706,12 @@ class TestCloudLoggingHandlerTask:
filter_str = (
'resource.type="global"\n'
'logName="projects/project_id/logs/airflow"\n'
- 'labels.task_id="task_for_testing_cloud_logging_task_handler"\n'
- 'labels.dag_id="dag_for_testing_cloud_logging_task_handler"\n'
- f'labels.{date_label}="2016-01-01T00:00:00+00:00"\n'
- 'labels.try_number="3"'
+ + self._expect_label_filter(
+ 'labels.task_id="task_for_testing_cloud_logging_task_handler"',
+ 'labels.dag_id="dag_for_testing_cloud_logging_task_handler"',
+ f'labels.{date_label}="2016-01-01T00:00:00+00:00"',
+ 'labels.try_number="3"',
+ )
)
mock_client.return_value.list_log_entries.assert_called_once_with(
request=ListLogEntriesRequest(
@@ -670,10 +742,12 @@ class TestCloudLoggingHandlerTask:
filter_str = (
'resource.type="global"\n'
'logName="projects/project_id/logs/airflow"\n'
- 'labels.task_id="task_for_testing_cloud_logging_task_handler"\n'
- 'labels.dag_id="dag_for_testing_cloud_logging_task_handler"\n'
- f'labels.{date_label}="2016-01-01T00:00:00+00:00"\n'
- 'labels.try_number="3"'
+ + self._expect_label_filter(
+ 'labels.task_id="task_for_testing_cloud_logging_task_handler"',
+ 'labels.dag_id="dag_for_testing_cloud_logging_task_handler"',
+ f'labels.{date_label}="2016-01-01T00:00:00+00:00"',
+ 'labels.try_number="3"',
+ )
)
mock_client.return_value.list_log_entries.assert_called_once_with(
request=ListLogEntriesRequest(
@@ -696,10 +770,12 @@ class TestCloudLoggingHandlerTask:
filter=(
'resource.type="global"\n'
'logName="projects/project_id/logs/airflow"\n'
-
'labels.task_id="task_for_testing_cloud_logging_task_handler"\n'
-
'labels.dag_id="dag_for_testing_cloud_logging_task_handler"\n'
- f'labels.{date_label}="2016-01-01T00:00:00+00:00"\n'
- 'labels.try_number="3"'
+ + self._expect_label_filter(
+
'labels.task_id="task_for_testing_cloud_logging_task_handler"',
+
'labels.dag_id="dag_for_testing_cloud_logging_task_handler"',
+ f'labels.{date_label}="2016-01-01T00:00:00+00:00"',
+ 'labels.try_number="3"',
+ )
),
order_by="timestamp asc",
page_size=1000,
@@ -817,10 +893,12 @@ class TestCloudLoggingHandlerTask:
expected_filter = [
'resource.type="global"',
'logName="projects/project_id/logs/airflow"',
- f'labels.task_id="{self.ti.task_id}"',
- f'labels.dag_id="{self.DAG_ID}"',
- f'labels.{date_label}="{self.ti.logical_date.isoformat() if
AIRFLOW_V_3_0_PLUS else self.ti.execution_date.isoformat()}"',
- f'labels.try_number="{self.ti.try_number}"',
+ *self._expect_label_filter(
+ f'labels.task_id="{self.ti.task_id}"',
+ f'labels.dag_id="{self.DAG_ID}"',
+ f'labels.{date_label}="{self.ti.logical_date.isoformat() if
AIRFLOW_V_3_0_PLUS else self.ti.execution_date.isoformat()}"',
+ f'labels.try_number="{self.ti.try_number}"',
+ ).splitlines(),
]
assert set(expected_filter) == set(filter_params)
diff --git
a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
index 1f164cf56c8..6b61b7a6b75 100644
---
a/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
+++
b/providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py
@@ -46,7 +46,11 @@ from airflow.providers.common.compat.module_loading import
import_string
from airflow.providers.common.compat.sdk import AirflowException,
TaskInstanceState, conf
from airflow.providers.opensearch.log.os_json_formatter import
OpensearchJSONFormatter
from airflow.providers.opensearch.log.os_response import Hit,
OpensearchResponse
-from airflow.providers.opensearch.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
+from airflow.providers.opensearch.version_compat import (
+ AIRFLOW_V_3_0_PLUS,
+ AIRFLOW_V_3_2_PLUS,
+ AIRFLOW_V_3_4_PLUS,
+)
from airflow.utils.log.file_task_handler import FileTaskHandler
from airflow.utils.log.logging_mixin import ExternalLoggingMixin, LoggingMixin
from airflow.utils.session import create_session
@@ -265,6 +269,31 @@ def _render_log_id(
)
+def _get_ti_id_fields(ti: TaskInstance | RuntimeTI) -> dict[str, str]:
+ # Before 3.4 a try can reuse the previous try's id, so only log_id
identifies it.
+ if not AIRFLOW_V_3_4_PLUS:
+ return {}
+ return {"ti_id": str(ti.id)}
+
+
+def _build_log_query(log_id: str, ti: TaskInstance | RuntimeTI) ->
list[dict[str, Any]]:
+ log_id_match = {"match_phrase": {"log_id": log_id}}
+ # Before 3.4 a cleared task instance gets a new id, which can differ from
the id its logs were written under.
+ if not AIRFLOW_V_3_4_PLUS:
+ return [log_id_match]
+ return [
+ {
+ "bool": {
+ "should": [
+ {"match_phrase": {"ti_id": str(ti.id)}},
+ {"bool": {"must": [log_id_match], "must_not": {"exists":
{"field": "ti_id"}}}},
+ ],
+ "minimum_should_match": 1,
+ }
+ }
+ ]
+
+
def _resolve_nested(hit: dict[Any, Any], parent_class=None) -> type[Hit]:
"""
Resolve nested hits from OpenSearch by iteratively navigating the
`_nested` field.
@@ -692,7 +721,7 @@ class OpensearchTaskHandler(FileTaskHandler,
ExternalLoggingMixin, LoggingMixin)
"query": {
"bool": {
"filter": [{"range": {self.offset_field: {"gt":
int(offset)}}}],
- "must": [{"match_phrase": {"log_id": log_id}}],
+ "must": _build_log_query(log_id, ti),
}
}
}
@@ -945,7 +974,7 @@ class OpensearchRemoteLogIO(LoggingMixin): # noqa: D101
log_id = _render_log_id(self.log_id_template, ti, ti.try_number) #
type: ignore[arg-type]
if self.write_stdout or self.write_to_opensearch:
- log_lines = self._parse_raw_log(local_loc.read_text(), log_id)
+ log_lines = self._parse_raw_log(local_loc.read_text(), log_id,
_get_ti_id_fields(ti)) # type: ignore[arg-type]
if self.write_stdout:
for line in log_lines:
@@ -965,7 +994,9 @@ class OpensearchRemoteLogIO(LoggingMixin): # noqa: D101
parent.rmdir()
parent = parent.parent
- def _parse_raw_log(self, log: str, log_id: str) -> list[dict[str, Any]]:
+ def _parse_raw_log(
+ self, log: str, log_id: str, extra_fields: dict[str, str] | None = None
+ ) -> list[dict[str, Any]]:
parsed_logs = []
offset = 1
for line in log.split("\n"):
@@ -976,7 +1007,7 @@ class OpensearchRemoteLogIO(LoggingMixin): # noqa: D101
except json.JSONDecodeError:
self.log.warning("Skipping non-JSON log line: %r", line)
log_dict = {"event": line}
- log_dict.update({"log_id": log_id, self.offset_field: offset})
+ log_dict.update({"log_id": log_id, **(extra_fields or {}),
self.offset_field: offset})
offset += 1
parsed_logs.append(log_dict)
return parsed_logs
@@ -1048,7 +1079,7 @@ class OpensearchRemoteLogIO(LoggingMixin): # noqa: D101
"query": {
"bool": {
"filter": [{"range": {self.offset_field: {"gt":
int(offset)}}}],
- "must": [{"match_phrase": {"log_id": log_id}}],
+ "must": _build_log_query(log_id, ti),
}
}
}
diff --git
a/providers/opensearch/src/airflow/providers/opensearch/version_compat.py
b/providers/opensearch/src/airflow/providers/opensearch/version_compat.py
index eb02fb81909..0909a24cc19 100644
--- a/providers/opensearch/src/airflow/providers/opensearch/version_compat.py
+++ b/providers/opensearch/src/airflow/providers/opensearch/version_compat.py
@@ -36,10 +36,12 @@ AIRFLOW_V_3_0_PLUS: bool = get_base_airflow_version_tuple()
>= (3, 0, 0)
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
diff --git
a/providers/opensearch/tests/integration/opensearch/log/test_os_remote_log_io.py
b/providers/opensearch/tests/integration/opensearch/log/test_os_remote_log_io.py
index a81c3eda1c5..024ea66e1a3 100644
---
a/providers/opensearch/tests/integration/opensearch/log/test_os_remote_log_io.py
+++
b/providers/opensearch/tests/integration/opensearch/log/test_os_remote_log_io.py
@@ -40,6 +40,7 @@ class _MockTI:
run_id: str = "integration_test_run"
try_number: int = 1
map_index: int = -1
+ id: uuid.UUID = dataclasses.field(default_factory=uuid.uuid4)
@pytest.mark.integration("opensearch")
@@ -156,3 +157,34 @@ class TestOpensearchRemoteLogIOIntegration:
log_entry = json.loads(log_messages[0])
assert "error_detail" in log_entry
assert log_entry["error_detail"] == error_detail
+
+ @pytest.mark.parametrize(
+ ("written", "reader", "expected"),
+ [
+ pytest.param(["legacy"], "x", ["legacy"],
id="pre-upgrade-logs-stay-readable"),
+ pytest.param(["x", "y"], "x", ["x"],
id="post-upgrade-reader-sees-only-its-try"),
+ pytest.param(["x", "y"], "y", ["y"], id="post-upgrade-other-try"),
+ pytest.param(["legacy", "x", "y"], "x", ["legacy", "x"],
id="mixed-index"),
+ ],
+ )
+ @patch(
+ "airflow.providers.opensearch.log.os_task_handler.TASK_LOG_FIELDS",
+ ["message"],
+ )
+ def test_read_scopes_to_try_and_keeps_pre_upgrade_logs(self, tmp_path,
written, reader, expected):
+ """Tries sharing one log_id stay separate, and entries without ti_id
stay readable."""
+ try_ids = {"x": uuid.uuid4(), "y": uuid.uuid4()}
+ log_id = _render_log_id(self.opensearch_io.log_id_template, _MockTI(),
1)
+ for name in written:
+ if name == "legacy":
+ legacy_entry = {"message": name, "log_id": log_id, "offset": 1}
+ self.opensearch_io.client.index(index=self.target_index,
body=legacy_entry)
+ continue
+ log_file = tmp_path / f"{name}.log"
+ log_file.write_text(json.dumps({"message": name}) + "\n")
+ self.opensearch_io.upload(log_file, _MockTI(id=try_ids[name]))
+ self.opensearch_io.client.indices.refresh(index=self.target_index)
+
+ _, log_messages = self.opensearch_io.read("",
_MockTI(id=try_ids[reader]))
+
+ assert sorted(json.loads(message)["event"] for message in
log_messages) == sorted(expected)
diff --git
a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
index b9e6582c663..db3f80e9c4b 100644
--- a/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
+++ b/providers/opensearch/tests/unit/opensearch/log/test_os_task_handler.py
@@ -24,7 +24,9 @@ import os
import re
from io import StringIO
from pathlib import Path
+from types import SimpleNamespace
from unittest.mock import Mock, call, patch
+from uuid import UUID, uuid4
import pendulum
import pytest
@@ -37,7 +39,9 @@ from airflow.providers.opensearch.log.os_task_handler import (
OpensearchRemoteLogIO,
OpensearchTaskHandler,
_build_log_fields,
+ _build_log_query,
_format_error_detail,
+ _get_ti_id_fields,
_render_log_id,
_safe_build_structured_log_message,
_strip_userinfo,
@@ -54,6 +58,41 @@ from tests_common.test_utils.version_compat import
AIRFLOW_V_3_0_PLUS
opensearchpy = pytest.importorskip("opensearchpy")
[email protected](
+ ("is_airflow_3_4_plus", "expected"),
+ [(False, {}), (True, {"ti_id": "some-ti-id"})],
+)
+def test_ti_id_is_only_written_from_airflow_3_4(is_airflow_3_4_plus, expected):
+ ti = SimpleNamespace(id="some-ti-id")
+
+ with
patch("airflow.providers.opensearch.log.os_task_handler.AIRFLOW_V_3_4_PLUS",
is_airflow_3_4_plus):
+ assert _get_ti_id_fields(ti) == expected
+
+
[email protected]("is_airflow_3_4_plus", [False, True])
+def test_log_query_matches_ti_id_or_documents_without_it(is_airflow_3_4_plus):
+ ti = SimpleNamespace(id=uuid4())
+ log_id_match = {"match_phrase": {"log_id": "some-log-id"}}
+
+ with
patch("airflow.providers.opensearch.log.os_task_handler.AIRFLOW_V_3_4_PLUS",
is_airflow_3_4_plus):
+ must = _build_log_query("some-log-id", ti)
+
+ if not is_airflow_3_4_plus:
+ assert must == [log_id_match]
+ return
+ assert must == [
+ {
+ "bool": {
+ "should": [
+ {"match_phrase": {"ti_id": str(ti.id)}},
+ {"bool": {"must": [log_id_match], "must_not": {"exists":
{"field": "ti_id"}}}},
+ ],
+ "minimum_should_match": 1,
+ }
+ }
+ ]
+
+
@dataclasses.dataclass
class _MockTI:
dag_id: str = "dag_for_testing_os_log_handler"
@@ -61,6 +100,7 @@ class _MockTI:
run_id: str = "run_for_testing_os_log_handler"
try_number: int = 1
map_index: int = -1
+ id: UUID = dataclasses.field(default_factory=uuid4)
def get_ti(dag_id, task_id, logical_date, create_task_instance):
@@ -729,15 +769,19 @@ class TestOpensearchRemoteLogIO:
mock_parse.assert_not_called()
mock_write.assert_not_called()
+
@patch("airflow.providers.opensearch.log.os_task_handler.AIRFLOW_V_3_4_PLUS",
True)
def test_write_to_opensearch(self, tmp_json_file, ti):
self.opensearch_io.write_stdout = False
log_id = _render_log_id(self.opensearch_io.log_id_template, ti,
ti.try_number)
- expected_log_lines =
self.opensearch_io._parse_raw_log(tmp_json_file.read_text(), log_id)
+ expected_log_lines = self.opensearch_io._parse_raw_log(
+ tmp_json_file.read_text(), log_id, {"ti_id": str(ti.id)}
+ )
with patch.object(self.opensearch_io, "_write_to_opensearch",
return_value=True) as mock_write:
self.opensearch_io.upload(tmp_json_file, ti)
mock_write.assert_called_once_with(expected_log_lines)
+ assert all(line["ti_id"] == str(ti.id) for line in expected_log_lines)
def test_raw_log_contains_log_id_and_offset(self, tmp_json_file, ti):
raw_log = tmp_json_file.read_text()
@@ -764,7 +808,7 @@ class TestOpensearchRemoteLogIO:
"query": {
"bool": {
"filter": [{"range": {self.opensearch_io.offset_field:
{"gt": 2}}}],
- "must": [{"match_phrase": {"log_id": log_id}}],
+ "must": _build_log_query(log_id, ti),
}
}
}
