Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable Docs DAG in CI leveraging existing CI connections #1428

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
+9 −40
Open

Enable Docs DAG in CI leveraging existing CI connections #1428

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
72e7885
Enable Docs DAG in CI leveraging existing CI connections
pankajkoti Dec 26, 2024
6379a8b
Update dev/dags/dbt_docs.py
pankajkoti Dec 27, 2024
5797a52
List files in temporary directory
pankajkoti Dec 27, 2024
caa1ed6
More prints
pankajkoti Dec 27, 2024
84703fc
Install deps
pankajkoti Dec 27, 2024
b4c4f65
Update Azure connection ID
pankajkoti Dec 27, 2024
bdde29f
Remove prints
pankajkoti Dec 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 9 additions & 40 deletions dev/dags/dbt_docs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,6 @@
from pathlib import Path

from airflow import DAG
from airflow.decorators import task
from airflow.exceptions import AirflowNotFoundException
from airflow.hooks.base import BaseHook
from pendulum import datetime

from cosmos import ProfileConfig
Expand All @@ -27,9 +24,6 @@
DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt"
DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH))

S3_CONN_ID = "aws_docs"
AZURE_CONN_ID = "azure_docs"
GCS_CONN_ID = "gcs_docs"

profile_config = ProfileConfig(
profile_name="default",
Expand All @@ -41,32 +35,6 @@
)


@task.branch(task_id="which_upload")
def which_upload():
"""Only run the docs tasks if we have the proper connections set up"""
downstream_tasks_to_run = []

try:
BaseHook.get_connection(S3_CONN_ID)
downstream_tasks_to_run += ["generate_dbt_docs_aws"]
except AirflowNotFoundException:
pass

# if we have an AZURE_CONN_ID, check if it's valid
try:
BaseHook.get_connection(AZURE_CONN_ID)
downstream_tasks_to_run += ["generate_dbt_docs_azure"]
except AirflowNotFoundException:
pass
try:
BaseHook.get_connection(GCS_CONN_ID)
downstream_tasks_to_run += ["generate_dbt_docs_gcs"]
except AirflowNotFoundException:
pass

return downstream_tasks_to_run


with DAG(
dag_id="docs_dag",
start_date=datetime(2023, 1, 1),
Expand All @@ -79,24 +47,25 @@ def which_upload():
task_id="generate_dbt_docs_aws",
project_dir=DBT_ROOT_PATH / "jaffle_shop",
profile_config=profile_config,
connection_id=S3_CONN_ID,
bucket_name="cosmos-docs",
connection_id="aws_s3_conn",
bucket_name="cosmos-ci-docs",
install_deps=True,
)

generate_dbt_docs_azure = DbtDocsAzureStorageOperator(
task_id="generate_dbt_docs_azure",
project_dir=DBT_ROOT_PATH / "jaffle_shop",
profile_config=profile_config,
connection_id=AZURE_CONN_ID,
bucket_name="$web",
connection_id="azure_abfs_conn",
bucket_name="cosmos-ci-docs",
install_deps=True,
)

generate_dbt_docs_gcs = DbtDocsGCSOperator(
task_id="generate_dbt_docs_gcs",
project_dir=DBT_ROOT_PATH / "jaffle_shop",
profile_config=profile_config,
connection_id=GCS_CONN_ID,
bucket_name="cosmos-docs",
connection_id="gcp_gs_conn",
bucket_name="cosmos-ci-docs",
install_deps=True,
)

which_upload() >> [generate_dbt_docs_aws, generate_dbt_docs_azure, generate_dbt_docs_gcs]
Loading