[AIRFLOW-4084] fix ElasticSearch log download (#5177)

cong-zhu · KevinYang21 · commit 56a038aae77b · 2019-05-13T17:49:07.000-07:00
diff --git a/airflow/utils/log/es_task_handler.py b/airflow/utils/log/es_task_handler.py
@@ -97,7 +97,7 @@ def _read(self, ti, try_number, metadata=None):
         offset = metadata['offset']
         log_id = self._render_log_id(ti, try_number)
 
-        logs = self.es_read(log_id, offset)
+        logs = self.es_read(log_id, offset, metadata)
 
         next_offset = offset if not logs else logs[-1].offset
 
@@ -113,7 +113,8 @@ def _read(self, ti, try_number, metadata=None):
         # delay before Elasticsearch makes the log available.
         if 'last_log_timestamp' in metadata:
             last_log_ts = timezone.parse(metadata['last_log_timestamp'])
-            if cur_ts.diff(last_log_ts).in_minutes() >= 5:
+            if cur_ts.diff(last_log_ts).in_minutes() >= 5 or 'max_offset' in metadata \
+                    and offset >= metadata['max_offset']:
                 metadata['end_of_log'] = True
 
         if offset != next_offset or 'last_log_timestamp' not in metadata:
@@ -123,14 +124,16 @@ def _read(self, ti, try_number, metadata=None):
 
         return message, metadata
 
-    def es_read(self, log_id, offset):
+    def es_read(self, log_id, offset, metadata):
         """
         Returns the logs matching log_id in Elasticsearch and next offset.
         Returns '' if no log is found or there was an error.
         :param log_id: the log_id of the log to read.
         :type log_id: str
         :param offset: the offset start to read log from.
         :type offset: str
+        :param metadata: log metadata, used for steaming log download.
+        :type metadata: dict
         """
 
         # Offset is the unique key for sorting logs given log_id.
@@ -139,9 +142,15 @@ def es_read(self, log_id, offset):
             .sort('offset')
 
         s = s.filter('range', offset={'gt': offset})
+        max_log_line = s.count()
+        if 'download_logs' in metadata and metadata['download_logs'] and 'max_offset' not in metadata:
+            try:
+                metadata['max_offset'] = s[max_log_line - 1].execute()[-1].offset if max_log_line > 0 else 0
+            except Exception:
+                self.log.exception('Could not get current log size with log_id: {}'.format(log_id))
 
         logs = []
-        if s.count() != 0:
+        if max_log_line != 0:
             try:
 
                 logs = s[self.MAX_LINE_PER_PAGE * self.PAGE:self.MAX_LINE_PER_PAGE] \
diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -19,7 +19,6 @@
 #
 
 import copy
-from io import BytesIO
 import itertools
 import json
 import logging
@@ -36,7 +35,7 @@
 import sqlalchemy as sqla
 from flask import (
     redirect, request, Markup, Response, render_template,
-    make_response, flash, jsonify, send_file, url_for)
+    make_response, flash, jsonify, url_for)
 from flask_appbuilder import BaseView, ModelView, expose, has_access
 from flask_appbuilder.actions import action
 from flask_appbuilder.models.sqla.filters import BaseFilter
@@ -545,30 +544,49 @@ def get_logs_with_metadata(self, session=None):
             models.TaskInstance.dag_id == dag_id,
             models.TaskInstance.task_id == task_id,
             models.TaskInstance.execution_date == dttm).first()
-        try:
+
+        def _get_logs_with_metadata(try_number, metadata):
             if ti is None:
                 logs = ["*** Task instance did not exist in the DB\n"]
                 metadata['end_of_log'] = True
             else:
-                dag = dagbag.get_dag(dag_id)
-                ti.task = dag.get_task(ti.task_id)
                 logs, metadatas = handler.read(ti, try_number, metadata=metadata)
                 metadata = metadatas[0]
+            return logs, metadata
 
+        try:
+            if ti is not None:
+                dag = dagbag.get_dag(dag_id)
+                ti.task = dag.get_task(ti.task_id)
             if response_format == 'json':
+                logs, metadata = _get_logs_with_metadata(try_number, metadata)
                 message = logs[0] if try_number is not None else logs
                 return jsonify(message=message, metadata=metadata)
 
-            file_obj = BytesIO(b'\n'.join(
-                log.encode('utf-8') for log in logs
-            ))
             filename_template = conf.get('core', 'LOG_FILENAME_TEMPLATE')
             attachment_filename = render_log_filename(
                 ti=ti,
                 try_number="all" if try_number is None else try_number,
                 filename_template=filename_template)
-            return send_file(file_obj, as_attachment=True,
-                             attachment_filename=attachment_filename)
+            metadata['download_logs'] = True
+
+            def _generate_log_stream(try_number, metadata):
+                if try_number is None and ti is not None:
+                    next_try = ti.next_try_number
+                    try_numbers = list(range(1, next_try))
+                else:
+                    try_numbers = [try_number]
+                for try_number in try_numbers:
+                    metadata.pop('end_of_log', None)
+                    metadata.pop('max_offset', None)
+                    metadata.pop('offset', None)
+                    while 'end_of_log' not in metadata or not metadata['end_of_log']:
+                        logs, metadata = _get_logs_with_metadata(try_number, metadata)
+                        yield "\n".join(logs) + "\n"
+            return Response(_generate_log_stream(try_number, metadata),
+                            mimetype="text/plain",
+                            headers={"Content-Disposition": "attachment; filename={}".format(
+                                attachment_filename)})
         except AttributeError as e:
             error_message = ["Task log handler {} does not support read logs.\n{}\n"
                              .format(task_log_reader, str(e))]
diff --git a/tests/utils/log/test_es_task_handler.py b/tests/utils/log/test_es_task_handler.py
@@ -190,6 +190,22 @@ def test_read_timeout(self):
         self.assertEqual(0, metadatas[0]['offset'])
         self.assertTrue(timezone.parse(metadatas[0]['last_log_timestamp']) == ts)
 
+    def test_read_as_download_logs(self):
+        ts = pendulum.now()
+        logs, metadatas = self.es_task_handler.read(self.ti,
+                                                    1,
+                                                    {'offset': 0,
+                                                     'last_log_timestamp': str(ts),
+                                                     'download_logs': True,
+                                                     'end_of_log': False})
+        self.assertEqual(1, len(logs))
+        self.assertEqual(len(logs), len(metadatas))
+        self.assertEqual(self.test_message, logs[0])
+        self.assertFalse(metadatas[0]['end_of_log'])
+        self.assertTrue(metadatas[0]['download_logs'])
+        self.assertEqual(1, metadatas[0]['offset'])
+        self.assertTrue(timezone.parse(metadatas[0]['last_log_timestamp']) > ts)
+
     def test_read_raises(self):
         with mock.patch.object(self.es_task_handler.log, 'exception') as mock_exception:
             with mock.patch("elasticsearch_dsl.Search.execute") as mock_execute:
diff --git a/tests/www/test_views.py b/tests/www/test_views.py
@@ -741,6 +741,29 @@ def test_get_logs_with_metadata_as_download_file(self):
         self.assertEqual(200, response.status_code)
         self.assertIn('Log for testing.', response.data.decode('utf-8'))
 
+    def test_get_logs_with_metadata_as_download_large_file(self):
+        with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.read") as read_mock:
+            first_return = (['1st line'], [{}])
+            second_return = (['2nd line'], [{'end_of_log': False}])
+            third_return = (['3rd line'], [{'end_of_log': True}])
+            fourth_return = (['should never be read'], [{'end_of_log': True}])
+            read_mock.side_effect = [first_return, second_return, third_return, fourth_return]
+            url_template = "get_logs_with_metadata?dag_id={}&" \
+                           "task_id={}&execution_date={}&" \
+                           "try_number={}&metadata={}&format=file"
+            try_number = 1
+            url = url_template.format(self.DAG_ID,
+                                      self.TASK_ID,
+                                      quote_plus(self.DEFAULT_DATE.isoformat()),
+                                      try_number,
+                                      json.dumps({}))
+            response = self.client.get(url)
+
+            self.assertIn('1st line', response.data.decode('utf-8'))
+            self.assertIn('2nd line', response.data.decode('utf-8'))
+            self.assertIn('3rd line', response.data.decode('utf-8'))
+            self.assertNotIn('should never be read', response.data.decode('utf-8'))
+
     def test_get_logs_with_metadata(self):
         url_template = "get_logs_with_metadata?dag_id={}&" \
                        "task_id={}&execution_date={}&" \
