Merge pull request gocardless#28 from tucows/master

Add DbtDocsGenerateOperator and extend DbtTestOperator
tactivos · Sep 28, 2020 · bdfc3f3 · bdfc3f3
2 parents 054891e + 38e5a0e
commit bdfc3f3
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/airflow_dbt/__init__.py b/airflow_dbt/__init__.py
@@ -3,5 +3,6 @@
     DbtSeedOperator,
     DbtSnapshotOperator,
     DbtRunOperator,
-    DbtTestOperator
+    DbtTestOperator,
+    DbtDocsGenerateOperator
 )
diff --git a/airflow_dbt/hooks/dbt_hook.py b/airflow_dbt/hooks/dbt_hook.py
@@ -41,6 +41,8 @@ def __init__(self,
                  dir='.',
                  vars=None,
                  full_refresh=False,
+                 data=False,
+                 schema=False,
                  models=None,
                  exclude=None,
                  select=None,
@@ -52,6 +54,8 @@ def __init__(self,
         self.target = target
         self.vars = vars
         self.full_refresh = full_refresh
+        self.data = data
+        self.schema = schema
         self.models = models
         self.exclude = exclude
         self.select = select
@@ -84,6 +88,12 @@ def run_cli(self, *command):
         if self.vars is not None:
             dbt_cmd.extend(['--vars', self._dump_vars()])
 
+        if self.data:
+            dbt_cmd.extend(['--data'])
+
+        if self.schema:
+            dbt_cmd.extend(['--schema'])
+
         if self.models is not None:
             dbt_cmd.extend(['--models', self.models])
 

diff --git a/airflow_dbt/operators/__init__.py b/airflow_dbt/operators/__init__.py
@@ -2,5 +2,6 @@
     DbtSeedOperator,
     DbtSnapshotOperator,
     DbtRunOperator,
-    DbtTestOperator
+    DbtTestOperator,
+    DbtDocsGenerateOperator
 )
diff --git a/airflow_dbt/operators/dbt_operator.py b/airflow_dbt/operators/dbt_operator.py
@@ -46,6 +46,8 @@ def __init__(self,
                  dbt_bin='dbt',
                  verbose=True,
                  full_refresh=False,
+                 data=False,
+                 schema=False,
                  *args,
                  **kwargs):
         super(DbtBaseOperator, self).__init__(*args, **kwargs)
@@ -56,6 +58,8 @@ def __init__(self,
         self.vars = vars
         self.models = models
         self.full_refresh = full_refresh
+        self.data = data
+        self.schema = schema
         self.exclude = exclude
         self.select = select
         self.dbt_bin = dbt_bin
@@ -69,6 +73,8 @@ def create_hook(self):
             dir=self.dir,
             vars=self.vars,
             full_refresh=self.full_refresh,
+            data=self.data,
+            schema=self.schema,
             models=self.models,
             exclude=self.exclude,
             select=self.select,
@@ -96,6 +102,16 @@ def execute(self, context):
         self.create_hook().run_cli('test')
 
 
+class DbtDocsGenerateOperator(DbtBaseOperator):
+    @apply_defaults
+    def __init__(self, profiles_dir=None, target=None, *args, **kwargs):
+        super(DbtDocsGenerateOperator, self).__init__(profiles_dir=profiles_dir, target=target, *args,
+                                                      **kwargs)
+
+    def execute(self, context):
+        self.create_hook().run_cli('docs', 'generate')
+
+
 class DbtSnapshotOperator(DbtBaseOperator):
     @apply_defaults
     def __init__(self, profiles_dir=None, target=None, *args, **kwargs):