From 81d7f8194e8bb27df13b350db7f2ec7150b0d159 Mon Sep 17 00:00:00 2001 From: Konstantinos Koukopoulos 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"