[TWTTR] Fix for rendering code on UI (twitter-forks#34)

Sampreeth24 · Apr 24, 2022 · d99ccc0 · d99ccc0
1 parent f5de1c1
commit d99ccc0
Show file tree

Hide file tree

Showing 3 changed files with 2,653 additions and 16 deletions.
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
@@ -2773,6 +2773,24 @@ def next_dagrun_data_interval(self, value: Optional[Tuple[datetime, datetime]])
         else:
             self.next_dagrun_data_interval_start, self.next_dagrun_data_interval_end = value
 
+    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
@@ -2794,21 +2812,7 @@ def get_last_dagrun(self, session=None, include_externally_triggered=False):
         )
 
     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)
 
     @staticmethod
     @provide_session

diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -998,7 +998,7 @@ def code(self, session=None):
         try:
             dag_id = request.args.get('dag_id')
             dag_orm = DagModel.get_dagmodel(dag_id, session=session)
-            code = DagCode.get_code_by_fileloc(dag_orm.fileloc)
+            code = DagCode.get_code_by_fileloc(dag_orm.get_local_fileloc())
             html_code = Markup(highlight(code, lexers.PythonLexer(), HtmlFormatter(linenos=True)))
 
         except Exception as e: