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

Author: vshshjn7

airflow/models/dag.py

@@ -1475,6 +1475,24 @@ class DagModel(Base, LoggingMixin):
     def __repr__(self):
         return "<DAG: {self.dag_id}>".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,

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))

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))