rhwang10 commented on a change in pull request #4303: [AIRFLOW-3370]
Elasticsearch log task handler additional features
URL: https://github.com/apache/airflow/pull/4303#discussion_r253541579
##########
File path: airflow/utils/log/es_task_handler.py
##########
@@ -68,12 +118,115 @@ def __init__(self, base_log_folder, filename_template,
self.mark_end_on_close = True
self.end_of_log_mark = end_of_log_mark
+ def set_context(self, ti):
+ if self.write_stdout:
+ self.writer = ParentStdout()
+ sys.stdout = self.writer
+
+ self.taskInstance = self._process_task_instance(ti)
+
+ self.handler = logging.StreamHandler(stream=sys.stdout)
+
+ if self.json_format:
+ self.handler.setFormatter(JsonFormatter(self.record_labels,
+ self.taskInstance))
+ elif not self.json_format:
+ self.handler.setFormatter(self.formatter)
+ self.handler.setLevel(self.level)
+
+ elif not self.write_stdout:
+ super(ElasticsearchTaskHandler, self).set_context(ti)
+ self.mark_end_on_close = not ti.raw
+
+ def emit(self, record):
+ if self.write_stdout:
+ self.formatter.format(record)
+ if self.handler is not None:
+ self.handler.emit(record)
+ elif not self.write_stdout:
+ super(ElasticsearchTaskHandler).emit(record)
+
+ def flush(self):
+ if self.handler is not None:
+ self.handler.flush()
+
+ def _process_task_instance(self, ti):
+ """
+ Process task instance information to create a log_id
+ key for Elasticsearch
+ """
+ ti_info = {'dag_id': str(ti.dag_id),
+ 'task_id': str(ti.task_id),
+ 'execution_date': str(ti.execution_date),
+ 'try_number': str(ti.try_number)}
+ return ti_info
+
+ def read(self, task_instance, try_number=None, metadata=None):
+ """
+ Read logs of a given task instance from elasticsearch.
+ :param task_instance: task instance object
+ :param try_number: task instance try_number to read logs from.
+ If None,it will start at 1.
+ """
Review comment:
added. added param types as well- think that's definitely a good idea
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
