Source code for operators.dbt.dbt_docs

# dbt_test.py
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context

from operators.dbt.dbt_base import BaseDBTOperator

from airflow.utils.decorators import apply_defaults
from airflow.hooks.base import BaseHook

from config import Config

import os
from typing import List, Optional, Dict, Any




class DBTDocsOperator(BaseDBTOperator):
    """
    Airflow operator to generate dbt docs, package them, and upload to S3.

        :param select: List of models to include in the docs generation.
        :param exclude: List of models to exclude from the docs generation.
        :param bucket: S3 bucket to upload the generated docs.
        :param s3_prefix: S3 prefix for the generated docs.
        :param args: Additional arguments for the BaseDBTOperator.
        :param kwargs: Additional keyword arguments for the BaseDBTOperator.

    """

    ui_color = "#FFB433"
    ui_fgcolor = "#000000"
    custom_operator_name = "DBT Docs Operator"

    @apply_defaults
    def __init__(
        self,
        select: Optional[List[str]] = None,
        exclude: Optional[List[str]] = None,
        bucket: str = Config.DATAHUB_S3_BUCKET,
        s3_prefix: str = "{env}/data_docs/dag={dag_id}/docs/{version}",
        *args,
        **kwargs,
    ) -> None:
        """
        :param select: List of models to include in the docs generation.
        :param exclude: List of models to exclude from the docs generation.
        :param bucket: S3 bucket to upload the generated docs.
        :param s3_prefix: S3 prefix for the generated docs.
        :param args: Additional arguments for the BaseDBTOperator.
        :param kwargs: Additional keyword arguments for the BaseDBTOperator.
        """
        super().__init__(*args, **kwargs)
        self.bucket = bucket
        self.s3_prefix = s3_prefix
        self.select = select
        self.exclude = exclude
        self.proxy_lambda = "https://5w3irbfxp5.execute-api.us-east-1.amazonaws.com/deploy/{proxy}/static_index.html"



    def execute(self, context: Dict[str, Any]) -> None:
        """Execute dbt docs generate and upload to S3."""

        super().execute(context)

        dbt_compile_cmd = " ".join(self.build_command("dbt compile"))
        dbt_build_cmd = " ".join(self.build_command("dbt build"))

        dbt_docs_cmd = " ".join(self.build_command("dbt docs generate --static"))

        try:
            self.log.info(f"Running: {dbt_compile_cmd}")
            self.run_ssh_command(dbt_compile_cmd)
            self.log.info(f"Running: {dbt_build_cmd}")
            self.run_ssh_command(dbt_build_cmd)
            self.run_ssh_command(f"cp {self.target_path}/run_results.json {self.target_path}/run_results_build.json")
            self.log.info(f"Running: {dbt_docs_cmd}")
            self.run_ssh_command(dbt_docs_cmd)
            self.log.info("dbt docs generation completed successfully.")
            self.red_datahub_recipe(context)
        except Exception as e:
            self.log.error(f"dbt docs generation failed! {e}")
            raise

        s3_key = self._get_s3_key()
        self._upload_to_s3(self.target_path, s3_key)

        # Generate and store public URL
        public_url = self._generate_public_url(s3_key)
        self.xcom_push(context, key="dbt_docs_url", value=public_url)


    def _get_s3_key(self) -> str:
        """Generate S3 key for upload."""
        return self.s3_prefix.format(
            env=self.target, dag_id=self.dag_id, version=self.job_id
        )

    def _upload_to_s3(self, docs_path: str, s3_key: str) -> None:
        """Upload packaged docs to S3."""
        self.run_ssh_command(
            f"ls {docs_path} ; aws s3 sync {docs_path} s3://{self.bucket}/{s3_key}"
        )

    def _generate_public_url(self, s3_key: str) -> str:
        """Generate public URL using GeneratePresignedUrl Lambda."""
        base_url = self.proxy_lambda.format(proxy=s3_key)
        print(f"Generated public URL: {base_url}")
        return base_url



    def replace_dbt_paths(self, obj, dbt_project_root):
        if isinstance(obj, dict):
            return {
                k: self.replace_dbt_paths(v, dbt_project_root) for k, v in obj.items()
            }
        elif isinstance(obj, list):
            return [self.replace_dbt_paths(item, dbt_project_root) for item in obj]
        elif isinstance(obj, str):
            return obj.replace("{DBT_PROJECT_ROOT}", dbt_project_root)
        return obj




    def red_datahub_recipe(self, context: Context) -> Dict[str, Any]:
        """Return datahub recipe for the operator."""

        import yaml
        import tempfile

        filename = "dbt_recipe.dhub.yaml"
        recipe_path = f"{self.project_directory}/{filename}"

        dag_name = context["dag"].dag_id
        if os.path.exists(recipe_path):
            with open(recipe_path, "r") as file:
                recipe = yaml.safe_load(file)
        else:
            recipe = {}

        # modify it
        recipe = self.replace_dbt_paths(recipe, self.target_path)
        print(f"Recipe: {recipe}")
        if 'source' in recipe:
            recipe["source"]['config']['env'] = "DEV"

        # recipe['meta_mapping'] = {
        #     "dbt_project_root": self.project_directory,
        #     "dag_name": dag_name,
        #     "target_platform": self.target_platform,
        # }
        # write it back
        recipe_temp_path= f"{self.target_path}/{filename}"

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as temp_file:
            yaml.dump(recipe, temp_file)
            temp_path = temp_file.name
        
            ssh_target = f"{self.dbt_ssh_user}@{self.dbt_ssh_host}"
            scp_command =[
                'scp', temp_path, f'{ssh_target}:{recipe_temp_path}'
            ]
            self._run_bash_command(' '.join(scp_command))
            

        self.log.info(
            f"Recipe written to {recipe_temp_path}"
        )
        self.run_ssh_command(
            f"datahub ingest run -c {recipe_temp_path}"
        )

        self.run_ssh_command(f"datahub ingest list-runs")