diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index 009f6a48466c2..b15981aab98a9 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.6.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): """