Generating Docs using DbtDocsS3Operator

[rahul-ve]

Hi team,

I am trying to use DbtDocsS3Operator to generate the docs but I dont see an option to specify image or pod_template_file.yaml like for DbtTestKubernetesOperator. Currently the DbtDocsS3Operator is using the Airflow image but it does not contain the project code.

is there a way to set the image and volumes used by this operator.

    generate_dbt_docs_aws = DbtDocsS3Operator(
        task_id="generate_dbt_docs_aws",
        name="generate-dbt-docs-aws",
        project_dir=POD_PROJECT_PATH,
        profile_config=ProfileConfig(
            profile_name="test_profile",
            target_name="test"
            profiles_yml_filepath = "/project/profiles.yml",
        ),
        # docs-specific arguments
        connection_id="aws_s3_conn",
        bucket_name="datalake",
        folder_dir="/dbt_docs/test",
        dbt_cmd_flags=["--static"],
    )

Thanks
Rahul

[tatiana]

Hi @rahul-ve thanks for creating this issue. This looks like a brilliant opportunity for a first contribution to Cosmos!

The current implementation for DbtDocsS3Operator is:

astronomer-cosmos/cosmos/operators/local.py

Lines 1166 to 1221 in f698166

	class DbtDocsS3LocalOperator(DbtDocsCloudLocalOperator):
	"""
	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
	:param folder_dir: This can be used to specify under which directory the generated DBT documentation should be
	uploaded.
	"""
	ui_color = "#FF9900"
	def __init__(
	self,
	*args: Any,
	aws_conn_id: str | None = None,
	**kwargs: Any,
	) -> None:
	if aws_conn_id:
	warnings.warn(
	"Please, use `connection_id` instead of `aws_conn_id`. The argument `aws_conn_id` will be"
	" deprecated in Cosmos 2.0",
	DeprecationWarning,
	)
	kwargs["connection_id"] = aws_conn_id
	super().__init__(*args, **kwargs)
	def upload_to_cloud_storage(self, project_dir: str, **kwargs: Any) -> None:
	"""Uploads the generated documentation to S3."""
	self.log.info(
	'Attempting to upload generated docs to S3 using S3Hook("%s")',
	self.connection_id,
	)
	from airflow.providers.amazon.aws.hooks.s3 import S3Hook
	target_dir = f"{project_dir}/target"
	hook = S3Hook(
	self.connection_id,
	extra_args={
	"ContentType": "text/html",
	},
	)
	for filename in self.required_files:
	key = f"{self.folder_dir}/{filename}" if self.folder_dir else filename
	s3_path = f"s3://{self.bucket_name}/{key}"
	self.log.info("Uploading %s to %s", filename, s3_path)
	hook.load_file(
	filename=f"{target_dir}/{filename}",
	bucket_name=self.bucket_name,
	key=key,
	replace=True,
	)

It would be possible to add a Kubernetes version in cosmos/operators/kubernetes.py, similar to:

astronomer-cosmos/cosmos/operators/kubernetes.py

Lines 151 to 159 in f698166

	class DbtBuildKubernetesOperator(DbtBuildMixin, DbtKubernetesBaseOperator):
	"""
	Executes a dbt core build command.
	"""
	template_fields: Sequence[str] = DbtKubernetesBaseOperator.template_fields + DbtBuildMixin.template_fields # type: ignore[operator]
	def __init__(self, *args: Any, **kwargs: Any) -> None:
	super().__init__(*args, **kwargs)

Would you like to contribute to this feature?

[rahul-ve]

@tatiana
Happy to try but need few clarifications.
running dbt docs generate is straight forward, copying the data form the pod to cloud storage is going to be tricky. Local operator uses a callback function, is that how k8soperator should work as well?