From 35b630eb8a606a954faef99f13ad34ef933bae83 Mon Sep 17 00:00:00 2001
From: Chris Hronek <chronek@linqto.com>
Date: Tue, 7 Nov 2023 12:40:56 -0700
Subject: [PATCH] Make conn_id parameter match the DbtDocsS3Operator

---
 cosmos/operators/local.py | 8 ++++----
 dev/dags/dbt_docs.py      | 9 +++++++++
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
index c0bdbbc00..0b132134d 100644
--- a/cosmos/operators/local.py
+++ b/cosmos/operators/local.py
@@ -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
 
@@ -591,7 +591,7 @@ 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
@@ -599,7 +599,7 @@ def upload_to_s3(self, project_dir: str) -> None:
         target_dir = f"{project_dir}/target"
 
         hook = S3Hook(
-            self.aws_conn_id,
+            self.connection_id,
             extra_args={
                 "ContentType": "text/html",
             },
diff --git a/dev/dags/dbt_docs.py b/dev/dags/dbt_docs.py
index edf89bdab..42a4eb39d 100644
--- a/dev/dags/dbt_docs.py
+++ b/dev/dags/dbt_docs.py
@@ -20,6 +20,7 @@
 from cosmos.operators import (
     DbtDocsAzureStorageOperator,
     DbtDocsS3Operator,
+    DbtFreshnessS3Operator,
     DbtDocsGCSOperator,
 )
 from cosmos.profiles import PostgresUserPasswordProfileMapping
@@ -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",