From 81d7f8194e8bb27df13b350db7f2ec7150b0d159 Mon Sep 17 00:00:00 2001
From: Konstantinos Koukopoulos <koukopoulos@gmail.com>
Date: Fri, 20 May 2022 19:48:51 +0300
Subject: [PATCH 1/2] Support restricted index patterns in Elasticsearch log
 handler

Sometimes Airflow doesn't have the ability to search across all indices
in an Elasticsearch server. This might be due to security settings in
the server. In these cases fetching the remote logs fails. To fix this
we create a index_patterns configuration setting that can be set to a
more restrictive pattern.
---
 airflow/config_templates/config.yml           |  7 +++
 airflow/config_templates/default_airflow.cfg  |  4 ++
 .../elasticsearch/log/es_task_handler.py      |  8 ++-
 .../log/elasticmock/fake_elasticsearch.py     | 25 ++++++---
 .../elasticsearch/log/test_es_task_handler.py | 54 +++++++++++++++++++
 5 files changed, 91 insertions(+), 7 deletions(-)

diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 009f6a48466c2..01ed3a3c5f0a9 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -2317,6 +2317,13 @@
       type: string
       example: ~
       default: "offset"
+    - name: index_patterns
+      description: |
+        Comma separated list of index patterns to use when searching for logs (default: `_all`).
+      version_added: 2.5.0
+      type: string
+      example: something-*
+      default: "_all"
 - name: elasticsearch_configs
   description: ~
   options:
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index 1cff5f8dbe75f..146d0077a9060 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -1163,6 +1163,10 @@ host_field = host
 # The field where offset is stored (normally either `offset` or `log.offset`)
 offset_field = offset
 
+# Comma separated list of index patterns to use when searching for logs (default: `_all`).
+# Example: index_patterns = something-*
+index_patterns = _all
+
 [elasticsearch_configs]
 use_ssl = False
 verify_certs = True
diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
index ecd964a1472f8..c0fecd7f2a4be 100644
--- a/airflow/providers/elasticsearch/log/es_task_handler.py
+++ b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -83,6 +83,7 @@ def __init__(
         offset_field: str = "offset",
         host: str = "localhost:9200",
         frontend: str = "localhost:5601",
+        index_patterns: str | None = conf.get("elasticsearch", "index_patterns", fallback="_all"),
         es_kwargs: dict | None = conf.getsection("elasticsearch_configs"),
         *,
         filename_template: str | None = None,
@@ -114,6 +115,7 @@ def __init__(
         self.json_fields = [label.strip() for label in json_fields.split(",")]
         self.host_field = host_field
         self.offset_field = offset_field
+        self.index_patterns = index_patterns
         self.context_set = False
 
         self.formatter: logging.Formatter
@@ -282,7 +284,11 @@ def es_read(self, log_id: str, offset: str, metadata: dict) -> list:
         :param metadata: log metadata, used for steaming log download.
         """
         # Offset is the unique key for sorting logs given log_id.
-        search = Search(using=self.client).query("match_phrase", log_id=log_id).sort(self.offset_field)
+        search = (
+            Search(index=self.index_patterns, using=self.client)
+            .query("match_phrase", log_id=log_id)
+            .sort(self.offset_field)
+        )
 
         search = search.filter("range", **{self.offset_field: {"gt": int(offset)}})
         max_log_line = search.count()
diff --git a/tests/providers/elasticsearch/log/elasticmock/fake_elasticsearch.py b/tests/providers/elasticsearch/log/elasticmock/fake_elasticsearch.py
index b325abb1ea704..c7887b864fa3e 100644
--- a/tests/providers/elasticsearch/log/elasticmock/fake_elasticsearch.py
+++ b/tests/providers/elasticsearch/log/elasticmock/fake_elasticsearch.py
@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import fnmatch
 import json
 
 from elasticsearch import Elasticsearch
@@ -363,6 +364,19 @@ def match_must_phrase(document, matches, must):
                     # use in as a proxy for match_phrase
                     matches.append(document)
 
+    # Check index(es) exists.
+    def _validate_search_targets(self, targets):
+        # TODO: support allow_no_indices query parameter
+        matches = set()
+        for target in targets:
+            if target == "_all" or target == "":
+                matches.update(self.__documents_dict)
+            elif "*" in target:
+                matches.update(fnmatch.filter(self.__documents_dict, target))
+            elif target not in self.__documents_dict:
+                raise NotFoundError(404, f"IndexMissingException[[{target}] missing]")
+        return matches
+
     def _normalize_index_to_list(self, index):
         # Ensure to have a list of index
         if index is None:
@@ -375,12 +389,11 @@ def _normalize_index_to_list(self, index):
             # Is it the correct exception to use ?
             raise ValueError("Invalid param 'index'")
 
-        # Check index(es) exists
-        for searchable_index in searchable_indexes:
-            if searchable_index not in self.__documents_dict:
-                raise NotFoundError(404, f"IndexMissingException[[{searchable_index}] missing]")
-
-        return searchable_indexes
+        return list(
+            self._validate_search_targets(
+                target for index in searchable_indexes for target in index.split(",")
+            )
+        )
 
     @staticmethod
     def _normalize_doc_type_to_list(doc_type):
diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py b/tests/providers/elasticsearch/log/test_es_task_handler.py
index bdf274732e555..0d0e4e9cfa963 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -105,6 +105,7 @@ def teardown_method(self):
 
     def test_client(self):
         assert isinstance(self.es_task_handler.client, elasticsearch.Elasticsearch)
+        assert self.es_task_handler.index_patterns == "_all"
 
     def test_client_with_config(self):
         es_conf = dict(conf.getsection("elasticsearch_configs"))
@@ -125,6 +126,21 @@ def test_client_with_config(self):
             es_kwargs=es_conf,
         )
 
+    def test_client_with_patterns(self):
+        # ensure creating with index patterns does not fail
+        patterns = "test_*,other_*"
+        handler = ElasticsearchTaskHandler(
+            base_log_folder=self.local_log_location,
+            end_of_log_mark=self.end_of_log_mark,
+            write_stdout=self.write_stdout,
+            json_format=self.json_format,
+            json_fields=self.json_fields,
+            host_field=self.host_field,
+            offset_field=self.offset_field,
+            index_patterns=patterns,
+        )
+        assert handler.index_patterns == patterns
+
     def test_read(self, ti):
         ts = pendulum.now()
         logs, metadatas = self.es_task_handler.read(
@@ -139,6 +155,44 @@ def test_read(self, ti):
         assert "1" == metadatas[0]["offset"]
         assert timezone.parse(metadatas[0]["last_log_timestamp"]) > ts
 
+    def test_read_with_patterns(self, ti):
+        ts = pendulum.now()
+        with mock.patch.object(self.es_task_handler, "index_patterns", new="test_*,other_*"):
+            logs, metadatas = self.es_task_handler.read(
+                ti, 1, {"offset": 0, "last_log_timestamp": str(ts), "end_of_log": False}
+            )
+
+        assert 1 == len(logs)
+        assert len(logs) == len(metadatas)
+        assert len(logs[0]) == 1
+        assert self.test_message == logs[0][0][-1]
+        assert not metadatas[0]["end_of_log"]
+        assert "1" == metadatas[0]["offset"]
+        assert timezone.parse(metadatas[0]["last_log_timestamp"]) > ts
+
+    def test_read_with_patterns_no_match(self, ti):
+        ts = pendulum.now()
+        with mock.patch.object(self.es_task_handler, "index_patterns", new="test_other_*,test_another_*"):
+            logs, metadatas = self.es_task_handler.read(
+                ti, 1, {"offset": 0, "last_log_timestamp": str(ts), "end_of_log": False}
+            )
+
+        assert 1 == len(logs)
+        assert len(logs) == len(metadatas)
+        assert [[]] == logs
+        assert not metadatas[0]["end_of_log"]
+        assert "0" == metadatas[0]["offset"]
+        # last_log_timestamp won't change if no log lines read.
+        assert timezone.parse(metadatas[0]["last_log_timestamp"]) == ts
+
+    def test_read_with_missing_index(self, ti):
+        ts = pendulum.now()
+        with mock.patch.object(self.es_task_handler, "index_patterns", new="nonexistent,test_*"):
+            with pytest.raises(elasticsearch.exceptions.NotFoundError, match=r".*nonexistent.*"):
+                self.es_task_handler.read(
+                    ti, 1, {"offset": 0, "last_log_timestamp": str(ts), "end_of_log": False}
+                )
+
     @pytest.mark.parametrize("seconds", [3, 6])
     def test_read_missing_logs(self, seconds, create_task_instance):
         """

From 9a7ba7790e6a4821651c0592c927c2df5d33ee50 Mon Sep 17 00:00:00 2001
From: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
Date: Wed, 7 Dec 2022 16:02:22 -0700
Subject: [PATCH 2/2] Update airflow/config_templates/config.yml

---
 airflow/config_templates/config.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 01ed3a3c5f0a9..b15981aab98a9 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -2320,7 +2320,7 @@
     - name: index_patterns
       description: |
         Comma separated list of index patterns to use when searching for logs (default: `_all`).
-      version_added: 2.5.0
+      version_added: 2.6.0
       type: string
       example: something-*
       default: "_all"
