From 42c2e011a0e74d6dfd898c15e539350e1295e367 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Tue, 1 Feb 2022 14:21:16 -0800
Subject: [PATCH 1/7] Emit "logs not found" message when ES logs appear to be
 missing

Current ES log handler will wait up to 5 minutes for logs to appear (or for _more_ logs to appear since last log message was emitted).  This produces undesirable behavior when the log message has been deleted from the elasticsearch cluster.  A user may wait a long time thinking that the logs are coming when they are not.

To resolve this, if no logs whatsoever have been retrieved after 5 seconds of trying, we give up and emit a "logs not  found" message.

If the task has only just started, this may be a "false negative", and we guide the user to refresh if they think that might be the case.
---
 .../elasticsearch/log/es_task_handler.py      | 21 ++++++---
 .../elasticsearch/log/test_es_task_handler.py | 44 ++++++++++++++++---
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index bb6b1f648baf6..5f07a2ea276c6 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -187,15 +187,24 @@ def _read(
         metadata['end_of_log'] = False if not logs else len(loading_hosts) == 0
 
         cur_ts = pendulum.now()
-        # Assume end of log after not receiving new log for 5 min,
-        # as executor heartbeat is 1 min and there might be some
-        # delay before Elasticsearch makes the log available.
         if 'last_log_timestamp' in metadata:
             last_log_ts = timezone.parse(metadata['last_log_timestamp'])
-            if (
+
+            # if we are not getting any logs at all after more than N seconds of trying,
+            # assume logs do not exist
+            if int(next_offset) == 0 and cur_ts.diff(last_log_ts).in_seconds() > 5:
+                metadata['end_of_log'] = True
+                message = (
+                    f"*** Log {log_id} not found in elasticsearch. "
+                    f"If your task started recently, please wait a moment and reload this page. "
+                    f"Otherwise, the logs for this task instance may have been removed."
+                )
+                return [('', message)], metadata
+            elif (
+                # Assume end of log after not receiving new log for N min,
                 cur_ts.diff(last_log_ts).in_minutes() >= 5
-                or 'max_offset' in metadata
-                and int(offset) >= int(metadata['max_offset'])
+                # if max_offset specified, respect it
+                or ('max_offset' in metadata and int(offset) >= int(metadata['max_offset']))
             ):
                 metadata['end_of_log'] = True
 
diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py b/tests/providers/elasticsearch/log/test_es_task_handler.py
index 05e53c21ba6c5..10802ec7f8d38 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -19,6 +19,7 @@
 import json
 import logging
 import os
+import re
 import shutil
 from unittest import mock
 from urllib.parse import quote
@@ -38,6 +39,19 @@
 from .elasticmock import elasticmock
 
 
+def get_ti(dag_id, task_id, execution_date, create_task_instance):
+    ti = create_task_instance(
+        dag_id=dag_id,
+        task_id=task_id,
+        execution_date=execution_date,
+        dagrun_state=DagRunState.RUNNING,
+        state=TaskInstanceState.RUNNING,
+    )
+    ti.try_number = 1
+    ti.raw = False
+    return ti
+
+
 class TestElasticsearchTaskHandler:
     DAG_ID = 'dag_for_testing_es_task_handler'
     TASK_ID = 'task_for_testing_es_log_handler'
@@ -47,16 +61,12 @@ class TestElasticsearchTaskHandler:
 
     @pytest.fixture()
     def ti(self, create_task_instance):
-        ti = create_task_instance(
+        yield get_ti(
             dag_id=self.DAG_ID,
             task_id=self.TASK_ID,
             execution_date=self.EXECUTION_DATE,
-            dagrun_state=DagRunState.RUNNING,
-            state=TaskInstanceState.RUNNING,
+            create_task_instance=create_task_instance,
         )
-        ti.try_number = 1
-        ti.raw = False
-        yield ti
         clear_db_runs()
         clear_db_dags()
 
@@ -131,6 +141,28 @@ def test_read(self, ti):
         assert '1' == metadatas[0]['offset']
         assert timezone.parse(metadatas[0]['last_log_timestamp']) > ts
 
+    def test_read_missing_logs(self, create_task_instance):
+        """
+        When the log actually isn't there to be found, we only want to wait for 5 seconds.
+        In this case we expect to receive a message of the form 'Log {log_id} not found in elasticsearch ...'
+        """
+        ti = get_ti(
+            self.DAG_ID,
+            self.TASK_ID,
+            pendulum.instance(self.EXECUTION_DATE).add(days=1),  # so logs are not found
+            create_task_instance=create_task_instance,
+        )
+        ts = pendulum.now().add(seconds=-6)
+        logs, metadatas = self.es_task_handler.read(ti, 1, {'offset': 0, 'last_log_timestamp': str(ts)})
+
+        assert 1 == len(logs)
+        assert re.match(r'.*Log .* not found in elasticsearch.*', logs[0][0][1]) is not None
+        assert len(logs) == len(metadatas)
+        assert len(logs[0]) == 1
+        assert metadatas[0]['end_of_log'] is True
+        assert '0' == metadatas[0]['offset']
+        assert timezone.parse(metadatas[0]['last_log_timestamp']) == ts
+
     def test_read_with_match_phrase_query(self, ti):
         similar_log_id = (
             f'{TestElasticsearchTaskHandler.TASK_ID}-'

From 17a3137150facf81f495c7063a56ec39c907b748 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Tue, 1 Feb 2022 15:48:24 -0800
Subject: [PATCH 2/7] Apply suggestions from code review

Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
---
 airflow/providers/elasticsearch/log/es_task_handler.py    | 7 ++++---
 tests/providers/elasticsearch/log/test_es_task_handler.py | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index 5f07a2ea276c6..30e773551f5d7 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -196,11 +196,12 @@ def _read(
                 metadata['end_of_log'] = True
                 message = (
                     f"*** Log {log_id} not found in elasticsearch. "
-                    f"If your task started recently, please wait a moment and reload this page. "
-                    f"Otherwise, the logs for this task instance may have been removed."
+                    "If your task started recently, please wait a moment and reload this page. "
+                    "Otherwise, the logs for this task instance may have been removed."
                 )
                 return [('', message)], metadata
-            elif (
+
+            if (
                 # Assume end of log after not receiving new log for N min,
                 cur_ts.diff(last_log_ts).in_minutes() >= 5
                 # if max_offset specified, respect it
diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py b/tests/providers/elasticsearch/log/test_es_task_handler.py
index 10802ec7f8d38..7d8d35cf22410 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -156,7 +156,7 @@ def test_read_missing_logs(self, create_task_instance):
         logs, metadatas = self.es_task_handler.read(ti, 1, {'offset': 0, 'last_log_timestamp': str(ts)})
 
         assert 1 == len(logs)
-        assert re.match(r'.*Log .* not found in elasticsearch.*', logs[0][0][1]) is not None
+        assert re.match(r'^\*\*\* Log .* not found in elasticsearch.*', logs[0][0][1]) is not None
         assert len(logs) == len(metadatas)
         assert len(logs[0]) == 1
         assert metadatas[0]['end_of_log'] is True

From bc1285edbd67c089f08d14453725f7b6157585e9 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Wed, 2 Feb 2022 15:22:21 -0800
Subject: [PATCH 3/7] fixup! Apply suggestions from code review

---
 .../elasticsearch/log/es_task_handler.py         | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index 30e773551f5d7..3ca90dfd603f6 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -194,12 +194,16 @@ def _read(
             # assume logs do not exist
             if int(next_offset) == 0 and cur_ts.diff(last_log_ts).in_seconds() > 5:
                 metadata['end_of_log'] = True
-                message = (
-                    f"*** Log {log_id} not found in elasticsearch. "
-                    "If your task started recently, please wait a moment and reload this page. "
-                    "Otherwise, the logs for this task instance may have been removed."
-                )
-                return [('', message)], metadata
+                return [
+                    (
+                        '',
+                        (
+                            f"*** Log {log_id} not found in elasticsearch. "
+                            "If your task started recently, please wait a moment and reload this page. "
+                            "Otherwise, the logs for this task instance may have been removed."
+                        ),
+                    )
+                ], metadata
 
             if (
                 # Assume end of log after not receiving new log for N min,

From c90c537bddf901739a986e356b124bb990cf89b5 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Wed, 2 Feb 2022 16:16:13 -0800
Subject: [PATCH 4/7] fixup! Apply suggestions from code review

---
 .../elasticsearch/log/test_es_task_handler.py | 35 ++++++++++++++-----
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py b/tests/providers/elasticsearch/log/test_es_task_handler.py
index 7d8d35cf22410..3357cd64a608a 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -141,7 +141,8 @@ def test_read(self, ti):
         assert '1' == metadatas[0]['offset']
         assert timezone.parse(metadatas[0]['last_log_timestamp']) > ts
 
-    def test_read_missing_logs(self, create_task_instance):
+    @pytest.mark.parametrize('seconds', [3, 6])
+    def test_read_missing_logs(self, seconds, create_task_instance):
         """
         When the log actually isn't there to be found, we only want to wait for 5 seconds.
         In this case we expect to receive a message of the form 'Log {log_id} not found in elasticsearch ...'
@@ -152,14 +153,23 @@ def test_read_missing_logs(self, create_task_instance):
             pendulum.instance(self.EXECUTION_DATE).add(days=1),  # so logs are not found
             create_task_instance=create_task_instance,
         )
-        ts = pendulum.now().add(seconds=-6)
+        ts = pendulum.now().add(seconds=-seconds)
         logs, metadatas = self.es_task_handler.read(ti, 1, {'offset': 0, 'last_log_timestamp': str(ts)})
 
         assert 1 == len(logs)
-        assert re.match(r'^\*\*\* Log .* not found in elasticsearch.*', logs[0][0][1]) is not None
+        if seconds > 5:
+            # we expect a log not found message when checking began more than 5 seconds ago
+            assert len(logs[0]) == 1
+            actual_message = logs[0][0][1]
+            expected_pattern = r'^\*\*\* Log .* not found in elasticsearch.*'
+            assert re.match(expected_pattern, actual_message) is not None
+            assert metadatas[0]['end_of_log'] is True
+        else:
+            # we've "waited" less than 5 seconds so it should not be "end of log" and should be no log message
+            assert len(logs[0]) == 0
+            assert logs == [[]]
+            assert metadatas[0]['end_of_log'] is False
         assert len(logs) == len(metadatas)
-        assert len(logs[0]) == 1
-        assert metadatas[0]['end_of_log'] is True
         assert '0' == metadatas[0]['offset']
         assert timezone.parse(metadatas[0]['last_log_timestamp']) == ts
 
@@ -242,15 +252,24 @@ def test_read_timeout(self, ti):
         ts = pendulum.now().subtract(minutes=5)
 
         self.es.delete(index=self.index_name, doc_type=self.doc_type, id=1)
+        # in the below call, offset=1 implies that we have already retrieved something
+        # if we had never retrieved any logs at all (offset=0), then we would have gotten
+        # a "logs not found" message after 5 seconds of trying
+        offset = 1
         logs, metadatas = self.es_task_handler.read(
-            ti, 1, {'offset': 0, 'last_log_timestamp': str(ts), 'end_of_log': False}
+            task_instance=ti,
+            try_number=1,
+            metadata={
+                'offset': offset,
+                'last_log_timestamp': str(ts),
+                'end_of_log': False,
+            },
         )
         assert 1 == len(logs)
         assert len(logs) == len(metadatas)
         assert [[]] == logs
         assert metadatas[0]['end_of_log']
-        # offset should be initialized to 0 if not provided.
-        assert '0' == metadatas[0]['offset']
+        assert str(offset) == metadatas[0]['offset']
         assert timezone.parse(metadatas[0]['last_log_timestamp']) == ts
 
     def test_read_as_download_logs(self, ti):

From 6aadc0fb3943427141efda06a72867251b3d6304 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Wed, 2 Feb 2022 21:39:02 -0800
Subject: [PATCH 5/7] Update es_task_handler.py

---
 airflow/providers/elasticsearch/log/es_task_handler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index 3ca90dfd603f6..7666af86b5a61 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -198,7 +198,7 @@ def _read(
                     (
                         '',
                         (
-                            f"*** Log {log_id} not found in elasticsearch. "
+                            f"*** Log {log_id} not found in Elasticsearch. "
                             "If your task started recently, please wait a moment and reload this page. "
                             "Otherwise, the logs for this task instance may have been removed."
                         ),

From 075f220b260487df911a62fe6d0ba5ae196fc20b Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Wed, 2 Feb 2022 23:39:29 -0800
Subject: [PATCH 6/7] fix readabilitya

---
 .../elasticsearch/log/es_task_handler.py        | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index 7666af86b5a61..7b043b126e978 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -194,17 +194,12 @@ def _read(
             # assume logs do not exist
             if int(next_offset) == 0 and cur_ts.diff(last_log_ts).in_seconds() > 5:
                 metadata['end_of_log'] = True
-                return [
-                    (
-                        '',
-                        (
-                            f"*** Log {log_id} not found in Elasticsearch. "
-                            "If your task started recently, please wait a moment and reload this page. "
-                            "Otherwise, the logs for this task instance may have been removed."
-                        ),
-                    )
-                ], metadata
-
+                missing_log_message = (
+                    f"*** Log {log_id} not found in Elasticsearch. "
+                    "If your task started recently, please wait a moment and reload this page. "
+                    "Otherwise, the logs for this task instance may have been removed."
+                )
+                return [('', missing_log_message)], metadata
             if (
                 # Assume end of log after not receiving new log for N min,
                 cur_ts.diff(last_log_ts).in_minutes() >= 5

From cf8d5df9633b0bc036807d4841056e21ef9e10ba Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Thu, 3 Feb 2022 11:35:55 -0800
Subject: [PATCH 7/7] fixup! fix readabilitya

---
 tests/providers/elasticsearch/log/test_es_task_handler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py b/tests/providers/elasticsearch/log/test_es_task_handler.py
index 3357cd64a608a..03eab3dbb78a6 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -161,7 +161,7 @@ def test_read_missing_logs(self, seconds, create_task_instance):
             # we expect a log not found message when checking began more than 5 seconds ago
             assert len(logs[0]) == 1
             actual_message = logs[0][0][1]
-            expected_pattern = r'^\*\*\* Log .* not found in elasticsearch.*'
+            expected_pattern = r'^\*\*\* Log .* not found in Elasticsearch.*'
             assert re.match(expected_pattern, actual_message) is not None
             assert metadatas[0]['end_of_log'] is True
         else:
