improve dbt plugin code (#20)

Author: ismailsimsek

opendbt/__init__.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import sys
+import tempfile
 from pathlib import Path
 
 from dbt.cli.main import dbtRunner as DbtCliRunner
@@ -94,15 +95,18 @@ def run(self, command: str = "build", target: str = None, args: list = None, use
         if write_json:
             run_args.remove("--no-write-json")
 
-        if use_subprocess:
+        if False:
             shell = False
             self.log.info("Working dir is %s" % os.getcwd())
             self.log.info("Running command (shell=%s) `%s`" % (shell, " ".join(command)))
             Utils.runcommand(command=['opendbt'] + run_args)
             return None
         else:
-            self.log.info(f"Running `dbt {' '.join(run_args)}`")
-            return OpenDbtCli.run(args=run_args)
+            with tempfile.TemporaryDirectory() as tmp_working_dir:
+                os.chdir(tmp_working_dir)
+                self.log.info(f"Running `dbt {' '.join(run_args)}`")
+                self.log.info("CWD is %s" % os.getcwd())
+                return OpenDbtCli.run(args=run_args)
 
     def manifest(self, partial_parse=True, no_write_manifest=True) -> Manifest:
         args = []

opendbt/airflow/__init__.py

@@ -46,9 +46,9 @@ def execute(self, context):
         """
         Execute the dbt command.
         """
-        runner = opendbt.OpenDbtProject(project_dir=self.project_dir,
-                                        profiles_dir=self.profiles_dir,
-                                        target=self.target)
+        runner = OpenDbtAirflowProject(project_dir=self.project_dir,
+                                       profiles_dir=self.profiles_dir,
+                                       target=self.target)
         runner.run(command=self.command, args=self.args, use_subprocess=self.use_subprocess)
 
 

opendbt/airflow/dbtdocs.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+
+def init_plugins_dbtdocs_page(dbt_docs_dir: Path):
+    from airflow.plugins_manager import AirflowPlugin
+    from flask import Blueprint
+    from flask_appbuilder import BaseView, expose
+    from flask import abort
+    from airflow.www.auth import has_access
+    from airflow.security import permissions
+
+    class DBTDocsView(BaseView):
+        route_base = "/dbt"
+        default_view = "dbt_docs_index"
+
+        @expose("/dbt_docs_index.html")  # type: ignore[misc]
+        @has_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE)])
+        def dbt_docs_index(self):
+            if not dbt_docs_dir.joinpath("index.html").is_file():
+                abort(404)
+            else:
+                return dbt_docs_dir.joinpath("index.html").read_text()
+            # return self.render_template("index.html", content="")
+
+        @expose("/catalog.json")  # type: ignore[misc]
+        @has_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE)])
+        def catalog(self):
+            if not dbt_docs_dir.joinpath("catalog.json").is_file():
+                abort(404)
+            else:
+                data = dbt_docs_dir.joinpath("catalog.json").read_text()
+                return data, 200, {"Content-Type": "application/json"}
+            # return self.render_template("index.html", content="")
+
+        @expose("/manifest.json")  # type: ignore[misc]
+        @has_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE)])
+        def manifest(self):
+            if not dbt_docs_dir.joinpath("manifest.json").is_file():
+                abort(404)
+            else:
+                data = dbt_docs_dir.joinpath("manifest.json").read_text()
+                return data, 200, {"Content-Type": "application/json"}
+            # return self.render_template("index.html", content="")
+
+    # Creating a flask blueprint to integrate the templates and static folder
+    bp = Blueprint(
+        "DBT Plugin",
+        __name__,
+        template_folder=dbt_docs_dir.as_posix(),
+        static_folder=dbt_docs_dir.as_posix(),
+        # static_url_path='/dbtdocsview'
+    )
+
+
+    class AirflowDbtDocsPlugin(AirflowPlugin):
+        name = "DBT Docs Plugin"
+        flask_blueprints = [bp]
+        appbuilder_views = [{"name": "DBT Docs", "category": "", "view": DBTDocsView()}]
+
+    return AirflowDbtDocsPlugin

opendbt/airflow/plugin.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from opendbt.airflow import dbtdocs
+from opendbt.airflow import plugin
 
 # create public page on airflow server to serve DBT docs
-airflow_dbtdocs_page = dbtdocs.init_plugins_dbtdocs_page(Path("/opt/dbttest/target"))
+airflow_dbtdocs_page = plugin.init_plugins_dbtdocs_page(Path("/opt/dbttest/target"))

opendbt/dbt/shared/__init__.py

@@ -19,9 +19,9 @@ class TestAirflowBase(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        cls._compose = DockerCompose(filepath=cls.resources_dir.joinpath('airflow').as_posix(),
+        cls._compose = DockerCompose(cls.resources_dir.joinpath('airflow').as_posix(),
                                      compose_file_name="docker-compose.yaml",
-                                     # build=True
+                                     build=True
                                      )
         cls._compose.stop()
         cls._compose.start()