Skip to content

Commit

Permalink
Extend DbtDocsLocalOperator with static flag (astronomer#759)
Browse files Browse the repository at this point in the history
## Description

<!-- Add a brief but complete description of the change. -->
Extends the docsOperator to make use of dbt `--static` flag. The static
flag is available from dbt 1.7>

I decided to add the flag through `dbt_cmd_flags`. Other options could
be a parameter or the `operator_args`.
There is no official documentation page from DBT on the flag
[yet](dbt-labs/docs.getdbt.com#4599). When
they do add it, I can trim down our documentation and link directly to
dbt.

I contribute to learn and am appreciative of any feedback.

## Related Issue(s)
closes astronomer#746 
<!-- If this PR closes an issue, you can use a keyword to auto-close.
-->
<!-- i.e. "closes #0000" -->

## Breaking Change?

<!-- If this introduces a breaking change, specify that here. -->

## Checklist

- [x] I have made corresponding changes to the documentation (if
required)
- [x] I have added tests that prove my fix is effective or that my
feature works

---------

Co-authored-by: Tatiana Al-Chueyr <[email protected]>
Co-authored-by: Justin Bandoro <[email protected]>
  • Loading branch information
3 people authored and arojasb3 committed Jul 14, 2024
1 parent 9890f23 commit 48dce2b
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 1 deletion.
11 changes: 10 additions & 1 deletion cosmos/operators/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -548,6 +548,15 @@ def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self.base_cmd = ["docs", "generate"]

self.check_static_flag()

def check_static_flag(self) -> None:
flag = "--static"
if self.dbt_cmd_flags:
if flag in self.dbt_cmd_flags:
# For the --static flag we only upload the generated static_index.html file
self.required_files = ["static_index.html"]


class DbtDocsCloudLocalOperator(DbtDocsLocalOperator, ABC):
"""
Expand Down Expand Up @@ -578,7 +587,7 @@ def upload_to_cloud_storage(self, project_dir: str) -> None:

class DbtDocsS3LocalOperator(DbtDocsCloudLocalOperator):
"""
Executes `dbt docs generate` command and upload to S3 storage. Returns the S3 path to the generated documentation.
Executes `dbt docs generate` command and upload to S3 storage.
:param connection_id: S3's Airflow connection ID
:param bucket_name: S3's bucket name
Expand Down
28 changes: 28 additions & 0 deletions docs/configuration/generating-docs.rst
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,34 @@ You can use the :class:`~cosmos.operators.DbtDocsGCSOperator` to generate and up
bucket_name="test_bucket",
)
Static Flag
~~~~~~~~~~~~~~~~~~~~~~~

All of the DbtDocsOperator accept the ``--static`` flag. To learn more about the static flag, check out the `original PR on dbt-core <https://github.com/dbt-labs/dbt-docs/pull/465>`_.
The static flag is used to generate a single doc file that can be hosted directly from cloud storage.
By having a single documentation file, you can make use of Access control can be configured through Identity-Aware Proxy (IAP), and making it easy to host.

.. note::
The static flag is only available from dbt-core >=1.7

The following code snippet shows how to provide this flag with the default jaffle_shop project:


.. code-block:: python
from cosmos.operators import DbtDocsGCSOperator
# then, in your DAG code:
generate_dbt_docs_aws = DbtDocsGCSOperator(
task_id="generate_dbt_docs_gcs",
project_dir="path/to/jaffle_shop",
profile_config=profile_config,
# docs-specific arguments
connection_id="test_gcs",
bucket_name="test_bucket",
dbt_cmd_flags=["--static"],
)
Custom Callback
~~~~~~~~~~~~~~~~~~~~~~~

Expand Down
11 changes: 11 additions & 0 deletions tests/operators/test_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -488,3 +488,14 @@ def test_operator_execute_deps_parameters(
mock_ensure_profile.return_value.__enter__.return_value = (Path("/path/to/profile"), {"ENV_VAR": "value"})
task.execute(context={"task_instance": MagicMock()})
assert mock_build_and_run_cmd.call_args_list[0].kwargs["command"] == expected_call_kwargs


def test_dbt_docs_local_operator_with_static_flag():
# Check when static flag is passed, the required files are correctly adjusted to a single file
operator = DbtDocsLocalOperator(
task_id="fake-task",
project_dir="fake-dir",
profile_config=profile_config,
dbt_cmd_flags=["--static"],
)
assert operator.required_files == ["static_index.html"]

0 comments on commit 48dce2b

Please sign in to comment.