Skip to content

Commit

Permalink
improve dbt plugin code (#20)
Browse files Browse the repository at this point in the history
  • Loading branch information
ismailsimsek authored Oct 16, 2024
1 parent 16b181c commit 8b5ed23
Show file tree
Hide file tree
Showing 7 changed files with 73 additions and 45 deletions.
10 changes: 7 additions & 3 deletions opendbt/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import os
import sys
import tempfile
from pathlib import Path

from dbt.cli.main import dbtRunner as DbtCliRunner
Expand Down Expand Up @@ -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 = []
Expand Down
6 changes: 3 additions & 3 deletions opendbt/airflow/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
35 changes: 0 additions & 35 deletions opendbt/airflow/dbtdocs.py

This file was deleted.

59 changes: 59 additions & 0 deletions opendbt/airflow/plugin.py
Original file line number Diff line number Diff line change
@@ -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
Empty file added opendbt/dbt/shared/__init__.py
Empty file.
4 changes: 2 additions & 2 deletions tests/resources/airflow/plugins/airflow_dbtdocs_page.py
Original file line number Diff line number Diff line change
@@ -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"))
4 changes: 2 additions & 2 deletions tests/test_airflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 8b5ed23

Please sign in to comment.