Skip to content

Commit

Permalink
Make conn_id parameter match the DbtDocsS3Operator
Browse files Browse the repository at this point in the history
  • Loading branch information
chrishronek committed Nov 7, 2023
1 parent c118fc1 commit 35b630e
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 4 deletions.
8 changes: 4 additions & 4 deletions cosmos/operators/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -572,13 +572,13 @@ class DbtFreshnessS3LocalOperator(DbtFreshnessLocalOperator):

def __init__(
self,
aws_conn_id: str,
connection_id: str,
bucket_name: str,
folder_dir: str | None = None,
**kwargs: str,
) -> None:
"Initializes the operator."
self.aws_conn_id = aws_conn_id
self.connection_id = connection_id
self.bucket_name = bucket_name
self.folder_dir = folder_dir

Expand All @@ -591,15 +591,15 @@ def upload_to_s3(self, project_dir: str) -> None:
"Uploads the generated documentation to S3."
logger.info(
'Attempting to upload generated docs to S3 using S3Hook("%s")',
self.aws_conn_id,
self.connection_id,
)

from airflow.providers.amazon.aws.hooks.s3 import S3Hook

target_dir = f"{project_dir}/target"

hook = S3Hook(
self.aws_conn_id,
self.connection_id,
extra_args={
"ContentType": "text/html",
},
Expand Down
9 changes: 9 additions & 0 deletions dev/dags/dbt_docs.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from cosmos.operators import (
DbtDocsAzureStorageOperator,
DbtDocsS3Operator,
DbtFreshnessS3Operator,
DbtDocsGCSOperator,
)
from cosmos.profiles import PostgresUserPasswordProfileMapping
Expand Down Expand Up @@ -83,6 +84,14 @@ def which_upload():
bucket_name="cosmos-docs",
)

generate_dbt_freshness_aws = DbtFreshnessS3Operator(
task_id="generate_dbt_freshness_aws",
project_dir=DBT_ROOT_PATH / "jaffle_shop",
profile_config=profile_config,
connection_id=S3_CONN_ID,
bucket_name="cosmos-docs",
)

generate_dbt_docs_azure = DbtDocsAzureStorageOperator(
task_id="generate_dbt_docs_azure",
project_dir=DBT_ROOT_PATH / "jaffle_shop",
Expand Down

0 comments on commit 35b630e

Please sign in to comment.