feat(operator): Add new operator for dbt docs generate

Support for dbt docs generate itself is pretty trivial as we only needed a
new configuration and dbt docs generate does not take many
configuration arguments. The biggest challenge of supporting dbt docs generate
is solving the issue of running dbt docs generate in a temporary
directory: there isn't much sense in executing dbt docs generate if we
cannot use the results of the execution. The alternative of pushing
the dbt artifacts to XCom is, of course, still available, but I wanted
to use this an excuse to develop something else.

So, we also added a push_dbt_project method to our S3
hook. This method is called right before exiting the temporary
directory, and can be controlled with a push_dbt_project argument to
the operator. The argument is by default True for the new
DbtDocsGenerateOperator so that we can push the documentation files to
S3. It is also True by default when running DbtDepsOperator since
there was no way for us install dependencies.

In the future, I'd like us to support more "dbt backends" besides S3:
maybe support for pulling projects from Github? SSH and SFTP
connections? Lots of possibilities!

Author: tomasfarias

airflow_dbt_python/hooks/dbt.py

@@ -25,6 +25,7 @@
 from dbt.task.debug import DebugTask
 from dbt.task.deps import DepsTask
 from dbt.task.freshness import FreshnessTask
+from dbt.task.generate import GenerateTask
 from dbt.task.list import ListTask
 from dbt.task.parse import ParseTask
 from dbt.task.run import RunTask
@@ -321,6 +322,15 @@ class DepsTaskConfig(BaseConfig):
     which: str = dataclasses.field(default="deps", init=False)
 
 
+@dataclass
+class GenerateTaskConfig(SelectionConfig):
+    """Generate task arguments."""
+
+    cls: BaseTask = dataclasses.field(default=GenerateTask, init=False)
+    compile: bool = True
+    which: str = dataclasses.field(default="generate", init=False)
+
+
 @dataclass
 class ListTaskConfig(SelectionConfig):
     """Dbt list task arguments."""
@@ -428,6 +438,7 @@ class ConfigFactory(FromStrMixin, Enum):
     CLEAN = CleanTaskConfig
     DEBUG = DebugTaskConfig
     DEPS = DepsTaskConfig
+    GENERATE = GenerateTaskConfig
     LIST = ListTaskConfig
     PARSE = ParseTaskConfig
     RUN = RunTaskConfig

airflow_dbt_python/hooks/s3.py

@@ -15,13 +15,13 @@ class DbtS3Hook(S3Hook):
     all the files corresponding to a project.
     """
 
-    def get_dbt_profiles(
+    def pull_dbt_profiles(
         self, s3_profiles_url: str, profiles_dir: Optional[str] = None
     ) -> Path:
-        """Fetch a dbt profiles file from S3.
+        """Pull a dbt profiles file from S3.
 
-        Fetches dbt profiles.yml file from the directory given by s3_profiles_url
-        and pulls it to profiles_dir/profiles.yml.
+        Pulls dbt profiles.yml file from the directory given by s3_profiles_url
+        and saves it to profiles_dir/profiles.yml.
 
         Args:
             s3_profiles_url: An S3 URL to a directory containing the dbt profiles file.
@@ -63,13 +63,13 @@ def download_one_s3_object(self, target: Path, s3_object):
             # exists.
             self.log.warning("A file with no name was found in S3 at %s", s3_object)
 
-    def get_dbt_project(
+    def pull_dbt_project(
         self, s3_project_url: str, project_dir: Optional[str] = None
     ) -> Path:
-        """Fetch all dbt project files from S3.
+        """Pull all dbt project files from S3.
 
-        Fetches the dbt project files from the directory given by s3_project_url
-        and pulls them to project_dir. However, if the URL points to a zip file,
+        Pulls the dbt project files from the directory given by s3_project_url
+        and saves them to project_dir. However, if the URL points to a zip file,
         we assume it contains all the project files, and only download and unzip that
         instead.
 
@@ -116,7 +116,7 @@ def get_dbt_project(
     def download_many_s3_keys(
         self, bucket_name: str, s3_keys: list[str], target_dir: Path, prefix: str
     ):
-        """Download multiple s3 keys."""
+        """Download multiple S3 keys."""
         for s3_object_key in s3_keys:
             s3_object = self.get_key(key=s3_object_key, bucket_name=bucket_name)
             path_file = Path(s3_object_key).relative_to(prefix)
@@ -131,3 +131,35 @@ def download_many_s3_keys(
             local_project_file.parent.mkdir(parents=True, exist_ok=True)
 
             self.download_one_s3_object(local_project_file, s3_object)
+
+    def push_dbt_project(self, s3_project_url: str, project_dir: str):
+        """Push a dbt project to S3."""
+        bucket_name, key = self.parse_s3_url(s3_project_url)
+        dbt_project_files = Path(project_dir).glob("**/*")
+
+        if key.endswith(".zip"):
+            zip_file_path = Path(project_dir) / "dbt_project.zip"
+            with ZipFile(zip_file_path, "w") as zf:
+                for _file in dbt_project_files:
+                    zf.write(_file, arcname=_file.relative_to(project_dir))
+
+            self.load_file(
+                zip_file_path, key=s3_project_url, bucket_name=bucket_name, replace=True
+            )
+            zip_file_path.unlink()
+
+        else:
+            for _file in dbt_project_files:
+                if _file.is_dir():
+                    continue
+
+                s3_key = f"s3://{bucket_name}/{key}{ _file.relative_to(project_dir)}"
+
+                self.load_file(
+                    filename=_file,
+                    key=s3_key,
+                    bucket_name=bucket_name,
+                    replace=True,
+                )
+
+        self.log.info("Pushed dbt project to: %s", s3_project_url)

airflow_dbt_python/operators/dbt.py

@@ -88,6 +88,7 @@ def __init__(
         # Extra features configuration
         s3_conn_id: str = "aws_default",
         do_xcom_push_artifacts: Optional[list[str]] = None,
+        push_dbt_project: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -132,6 +133,7 @@ def __init__(
 
         self.s3_conn_id = s3_conn_id
         self.do_xcom_push_artifacts = do_xcom_push_artifacts
+        self.push_dbt_project = push_dbt_project
         self._s3_hook = None
         self._dbt_hook = None
 
@@ -222,7 +224,10 @@ def dbt_directory(self) -> Iterator[str]:
         """Provides a temporary directory to execute dbt.
 
         Creates a temporary directory for dbt to run in and prepares the dbt files
-        if they need to be pulled from S3.
+        if they need to be pulled from S3. If a S3 backend is being used, and
+        self.push_dbt_project is True, before leaving the temporary directory, we push
+        back the project to S3. Pushing back a project enables commands like deps or
+        docs generate.
 
         Yields:
             The temporary directory's name.
@@ -242,22 +247,29 @@ def dbt_directory(self) -> Iterator[str]:
 
             yield tmp_dir
 
+            if (
+                self.push_dbt_project is True
+                and urlparse(str(store_project_dir)).scheme == "s3"
+            ):
+                self.log.info("Pushing dbt project back to S3: %s", store_project_dir)
+                self.s3_hook.push_dbt_project(store_project_dir, tmp_dir)
+
         self.profiles_dir = store_profiles_dir
         self.project_dir = store_project_dir
 
     def prepare_directory(self, tmp_dir: str):
         """Prepares a dbt directory by pulling files from S3."""
         if urlparse(str(self.profiles_dir)).scheme == "s3":
             self.log.info("Fetching profiles.yml from S3: %s", self.profiles_dir)
-            profiles_file_path = self.s3_hook.get_dbt_profiles(
+            profiles_file_path = self.s3_hook.pull_dbt_profiles(
                 self.profiles_dir,
                 tmp_dir,
             )
             self.profiles_dir = str(profiles_file_path.parent) + "/"
 
         if urlparse(str(self.project_dir)).scheme == "s3":
             self.log.info("Fetching dbt project from S3: %s", self.project_dir)
-            project_dir_path = self.s3_hook.get_dbt_project(
+            project_dir_path = self.s3_hook.pull_dbt_project(
                 self.project_dir,
                 tmp_dir,
             )
@@ -463,15 +475,34 @@ class DbtDepsOperator(DbtBaseOperator):
     https://docs.getdbt.com/reference/commands/deps.
     """
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, push_dbt_project=True, **kwargs) -> None:
         super().__init__(**kwargs)
+        self.push_dbt_project = push_dbt_project
 
     @property
     def command(self) -> str:
         """Return the deps command."""
         return "deps"
 
 
+class DbtDocsGenerateOperator(DbtBaseOperator):
+    """Executes a dbt docs generate command.
+
+    The documentation for the dbt command can be found here:
+    https://docs.getdbt.com/reference/commands/cmd-docs.
+    """
+
+    def __init__(self, compile=True, push_dbt_project=True, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.compile = compile
+        self.push_dbt_project = push_dbt_project
+
+    @property
+    def command(self) -> str:
+        """Return the generate command."""
+        return "generate"
+
+
 class DbtCleanOperator(DbtBaseOperator):
     """Executes a dbt clean command.
 

tests/hooks/dbt/test_dbt_generate.py

@@ -0,0 +1,64 @@
+"""Unit test module for running dbt docs generate with the DbtHook."""
+
+
+def test_dbt_docs_generate_task(hook, profiles_file, dbt_project_file, model_files):
+    """Test a dbt docs generate task."""
+    import shutil
+
+    target_dir = dbt_project_file.parent / "target"
+    if target_dir.exists() is True:
+        shutil.rmtree(target_dir)
+    assert target_dir.exists() is False
+
+    factory = hook.get_config_factory("generate")
+    config = factory.create_config(
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file.parent,
+    )
+    success, results = hook.run_dbt_task(config)
+
+    assert success is True
+    assert results is not None
+    assert target_dir.exists() is True
+
+    index = target_dir / "index.html"
+    assert index.exists() is True
+
+    manifest = target_dir / "manifest.json"
+    assert manifest.exists() is True
+
+    catalog = target_dir / "catalog.json"
+    assert catalog.exists() is True
+
+
+def test_dbt_docs_generate_task_no_compile(
+    hook, profiles_file, dbt_project_file, model_files
+):
+    """Test a dbt docs generate task without compiling."""
+    import shutil
+
+    target_dir = dbt_project_file.parent / "target"
+    if target_dir.exists() is True:
+        shutil.rmtree(target_dir)
+    assert target_dir.exists() is False
+
+    factory = hook.get_config_factory("generate")
+    config = factory.create_config(
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file.parent,
+        compile=False,
+    )
+    success, results = hook.run_dbt_task(config)
+
+    assert success is True
+    assert results is not None
+    assert target_dir.exists() is True
+
+    index = target_dir / "index.html"
+    assert index.exists() is True
+
+    manifest = target_dir / "manifest.json"
+    assert manifest.exists() is False
+
+    catalog = target_dir / "catalog.json"
+    assert catalog.exists() is True