[Bug] DbtDocsGCSLocalOperator.upload_to_cloud_storage() got an unexpected keyword argument 'context'

[tuantran0910]

Astronomer Cosmos Version

1.8.0

dbt-core version

1.9.0

Versions of dbt adapters

No response

LoadMode

AUTOMATIC

ExecutionMode

LOCAL

InvocationMode

None

airflow version

2.10.3

Operating System

linux

If a you think it's an UI issue, what browsers are you seeing the problem on?

No response

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

What happened?

Since astronomer-cosmos version v1.8.0 has been released recently, I have found a new bug related to dbt docs

In the file cosmos/operators/local.py, the context is passed as parameter to the callback function:

class DbtLocalBaseOperator(AbstractDbtBaseOperator):

      # Other logic 

      def run_command(
          self,
          cmd: list[str],
          env: dict[str, str | bytes | os.PathLike[Any]],
          context: Context,
      ) -> FullOutputSubprocessResult | dbtRunnerResult:

              # Other logic

              if self.callback:
                  self.callback_args.update({"context": context})
                  self.callback(tmp_project_dir, **self.callback_args)

The class DbtDocsCloudLocalOperator has the method upload_to_cloud_storage and assign it as the callback function:

class DbtDocsCloudLocalOperator(DbtDocsLocalOperator, ABC):
    """
    Abstract class for operators that upload the generated documentation to cloud storage.
    """

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

        super().__init__(**kwargs)

        # override the callback with our own
        self.callback = self.upload_to_cloud_storage

    @abstractmethod
    def upload_to_cloud_storage(self, project_dir: str) -> None:
        """Abstract method to upload the generated documentation to cloud storage."""

Here the upload_to_cloud_storage method is missing the parameter context, then it causes the error: TypeError: DbtDocsGCSLocalOperator.upload_to_cloud_storage() got an unexpected keyword argument 'context'

Relevant log output

  File "/opt/airflow/dag.py", line 68, in docs_dag
    DbtDocsGCSOperator(
  File "/home/airflow/.local/lib/python3.10/site-packages/airflow/models/baseoperator.py", line 406, in wrapper
    return func(self, *args, **kwargs)
  File "/home/airflow/.local/lib/python3.10/site-packages/cosmos/operators/base.py", line 268, in execute
    self.build_and_run_cmd(context=context, cmd_flags=self.add_cmd_flags())
  File "/home/airflow/.local/lib/python3.10/site-packages/cosmos/operators/local.py", line 662, in build_and_run_cmd
    result = self.run_command(cmd=dbt_cmd, env=env, context=context)
  File "/home/airflow/.local/lib/python3.10/site-packages/cosmos/operators/local.py", line 516, in run_command
    self.callback(tmp_project_dir, **self.callback_args)
TypeError: DbtDocsGCSLocalOperator.upload_to_cloud_storage() got an unexpected keyword argument 'context'

How to reproduce

In the code, I call this as a task in my dag:

DbtDocsGCSOperator(
    task_id="generate_dbt_docs_gcs",
    project_dir="path/to/jaffle_shop",
    profile_config=profile_config,
    # docs-specific arguments
    connection_id="test_gcs",
    bucket_name="test_bucket",
)

Anything else :)?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Contact Details

[email protected]

[dosubot]

Hey there, @tuantran0910! I'm here to help you with any bugs, questions, and contributions you have. Let's work together to solve this issue!

The error you're encountering is due to the upload_to_cloud_storage method in the DbtDocsGCSLocalOperator class not accepting the context parameter, which is being passed to it as part of the callback function. To resolve this issue, you need to modify the upload_to_cloud_storage method to accept the context parameter.

Here's how you can adjust the method signature:

def upload_to_cloud_storage(self, project_dir: str, context: Context) -> None:
    """Method to upload the generated documentation to cloud storage."""
    # Your implementation here

By adding the context parameter, the method will match the expected signature when used as a callback, preventing the TypeError from occurring.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[pankajastro]

Thank you very much for reporting it. Indeed this look like a bug may be caused by #1389. Could you please try https://pypi.org/project/astronomer-cosmos/1.8.1a1/ and see if it fixes.

[tuantran0910]

Hi, thank you for replying me. I have tested it and the issue has been fixed.

[pankajastro]

Thanks @tuantran0910 for quick response. We will release a stable version soon.