From 4ce8d4cfb24e3c9e50d3ff582b97ff79ab509168 Mon Sep 17 00:00:00 2001 From: vshshjn7 Date: Mon, 9 Mar 2020 18:04:20 +0530 Subject: [PATCH] [TWTTR] Fix for rendering code on UI (#34) --- airflow/models/dag.py | 34 +++++++++++++++++++--------------- airflow/www/views.py | 2 +- airflow/www_rbac/views.py | 2 +- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/airflow/models/dag.py b/airflow/models/dag.py index a82fb8002141a..a621ab88461cd 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -1475,6 +1475,24 @@ class DagModel(Base, LoggingMixin): def __repr__(self): return "".format(self=self) + def get_local_fileloc(self): + # TODO: [CX-16591] Resolve this in upstream by storing relative path in db (config driven) + try: + # Fix for DAGs that are manually triggered in the UI, as the DAG path in the DB is + # stored by the scheduler which has a different path than the webserver due to absolute + # paths in aurora including randomly generated job-specific directories. Due to this + # the path the webserver uses when it tries to trigger a DAG does not match the + # existing scheduler path and the DAG can not be found. + # Also, fix for render code on UI by changing "/code" in views.py + path_regex = "airflow_scheduler-.-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[" \ + "0-9a-f]{12}/runs/.*/sandbox/airflow_home" + path_split = re.split(path_regex, self.fileloc)[1] + return os.environ.get("AIRFLOW_HOME") + path_split + except IndexError: + self.log.info("No airflow_home in path: " + self.fileloc) + + return self.fileloc + @property def timezone(self): return settings.TIMEZONE @@ -1505,21 +1523,7 @@ def safe_dag_id(self): return self.dag_id.replace('.', '__dot__') def get_dag(self): - # TODO: [CX-16591] Resolve this in upstream by storing relative path in db (config driven) - try: - # Fix for DAGs that are manually triggered in the UI, as the DAG path in the DB is - # stored by the scheduler which has a different path than the webserver due to absolute - # paths in aurora including randomly generated job-specific directories. Due to this - # the path the webserver uses when it tries to trigger a DAG does not match the - # existing scheduler path and the DAG can not be found. - path_regex = "airflow_scheduler-.-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[" \ - "0-9a-f]{12}/runs/.*/sandbox/airflow_home" - path_split = re.split(path_regex, self.fileloc)[1] - self.fileloc = os.environ.get("AIRFLOW_HOME") + path_split - except IndexError: - self.log.info("No airflow_home in path: " + self.fileloc) - - return DagBag(dag_folder=self.fileloc).get_dag(self.dag_id) + return DagBag(dag_folder=self.get_local_fileloc()).get_dag(self.dag_id) @provide_session def create_dagrun(self, diff --git a/airflow/www/views.py b/airflow/www/views.py index e0b4097159ff4..0c3759d165cd7 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -696,7 +696,7 @@ def code(self, session=None): dm = models.DagModel dag = session.query(dm).filter(dm.dag_id == dag_id).first() try: - with wwwutils.open_maybe_zipped(dag.fileloc, 'r') as f: + with wwwutils.open_maybe_zipped(dag.get_local_fileloc(), 'r') as f: code = f.read() html_code = highlight( code, lexers.PythonLexer(), HtmlFormatter(linenos=True)) diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py index 3a8020fb2324c..7139d7d237d1f 100644 --- a/airflow/www_rbac/views.py +++ b/airflow/www_rbac/views.py @@ -446,7 +446,7 @@ def code(self, session=None): dag_id = request.args.get('dag_id') dag = session.query(dm).filter(dm.dag_id == dag_id).first() try: - with wwwutils.open_maybe_zipped(dag.fileloc, 'r') as f: + with wwwutils.open_maybe_zipped(dag.get_local_fileloc(), 'r') as f: code = f.read() html_code = highlight( code, lexers.PythonLexer(), HtmlFormatter(linenos=True))