From 4ac32cffc4124a5702dade5db0f832870496e0d1 Mon Sep 17 00:00:00 2001
From: Pankaj Singh <98807258+pankajastro@users.noreply.github.com>
Date: Thu, 26 Dec 2024 12:29:42 +0530
Subject: [PATCH] Add `kwargs` param in DocsOperator method
 upload_to_cloud_storage (#1422)

This PR addresses an issue where the `callback` callable parameter's
handling was modified in PR
https://github.com/astronomer/astronomer-cosmos/pull/1389 leading to a
change in
the way keyword arguments like context were passed to the callable.

However, the Docs Operator has its own implementation of the callback
callable,
which only expects a single parameter, project_dir. As a result, passing
extra
keyword arguments like context is causing a mismatch and resulting in
the error described in #1420.

This PR adds the `kwargs` param in `upload_to_cloud_storage` method of the Docs operators

We have integration tests for these operator but look like CI does have
not required setup and it get ignored.

https://github.com/astronomer/astronomer-cosmos/blob/main/dev/dags/dbt_docs.py

related: #1420
---
 CHANGELOG.rst             | 7 +++++++
 cosmos/__init__.py        | 2 +-
 cosmos/operators/local.py | 8 ++++----
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 168d8fd0b..3cff1bbb3 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,13 @@
 Changelog
 =========
 
+1.8.1a1 (2024-12-23)
+--------------------
+
+Bug Fixes
+
+* Add ``kwargs`` param in DocsOperator method ``upload_to_cloud_storage`` by @pankajastro in #1422
+
 1.8.0 (2024-12-20)
 --------------------
 
diff --git a/cosmos/__init__.py b/cosmos/__init__.py
index 81da409ab..9577b9dbd 100644
--- a/cosmos/__init__.py
+++ b/cosmos/__init__.py
@@ -6,7 +6,7 @@
 Contains dags, task groups, and operators.
 """
 
-__version__ = "1.8.0"
+__version__ = "1.8.1a1"
 
 
 from cosmos.airflow.dag import DbtDag
diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
index dee00114c..c5d47dd48 100644
--- a/cosmos/operators/local.py
+++ b/cosmos/operators/local.py
@@ -862,7 +862,7 @@ def __init__(
         self.callback = self.upload_to_cloud_storage
 
     @abstractmethod
-    def upload_to_cloud_storage(self, project_dir: str) -> None:
+    def upload_to_cloud_storage(self, project_dir: str, **kwargs: Any) -> None:
         """Abstract method to upload the generated documentation to cloud storage."""
 
 
@@ -893,7 +893,7 @@ def __init__(
             kwargs["connection_id"] = aws_conn_id
         super().__init__(*args, **kwargs)
 
-    def upload_to_cloud_storage(self, project_dir: str) -> None:
+    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")',
@@ -959,7 +959,7 @@ def __init__(
             kwargs["bucket_name"] = container_name
         super().__init__(*args, **kwargs)
 
-    def upload_to_cloud_storage(self, project_dir: str) -> None:
+    def upload_to_cloud_storage(self, project_dir: str, **kwargs: Any) -> None:
         """Uploads the generated documentation to Azure Blob Storage."""
         self.log.info(
             'Attempting to upload generated docs to Azure Blob Storage using WasbHook(conn_id="%s")',
@@ -1003,7 +1003,7 @@ class DbtDocsGCSLocalOperator(DbtDocsCloudLocalOperator):
 
     ui_color = "#4772d5"
 
-    def upload_to_cloud_storage(self, project_dir: str) -> None:
+    def upload_to_cloud_storage(self, project_dir: str, **kwargs: Any) -> None:
         """Uploads the generated documentation to Google Cloud Storage"""
         self.log.info(
             'Attempting to upload generated docs to Storage using GCSHook(conn_id="%s")',