From 92314e8c8f7f6df8a4f81b8ca4641e97b3b55c17 Mon Sep 17 00:00:00 2001
From: Pankaj Koti <pankajkoti699@gmail.com>
Date: Fri, 24 Jan 2025 14:20:28 +0530
Subject: [PATCH] Resolve type-check errrors with respect to update method
 signatures

---
 cosmos/operators/azure_container_instance.py | 8 +++++++-
 cosmos/operators/docker.py                   | 8 +++++++-
 cosmos/operators/gcp_cloud_run_job.py        | 8 +++++++-
 cosmos/operators/kubernetes.py               | 8 +++++++-
 cosmos/operators/virtualenv.py               | 2 ++
 5 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/cosmos/operators/azure_container_instance.py b/cosmos/operators/azure_container_instance.py
index 7f335bd99..39c39590b 100644
--- a/cosmos/operators/azure_container_instance.py
+++ b/cosmos/operators/azure_container_instance.py
@@ -63,7 +63,13 @@ def __init__(
             **kwargs,
         )
 
-    def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> None:
+    def build_and_run_cmd(
+        self,
+        context: Context,
+        cmd_flags: list[str] | None = None,
+        run_as_async: bool = False,
+        async_context: dict[str, Any] | None = None,
+    ) -> Any:
         self.build_command(context, cmd_flags)
         self.log.info(f"Running command: {self.command}")
         result = AzureContainerInstancesOperator.execute(self, context)
diff --git a/cosmos/operators/docker.py b/cosmos/operators/docker.py
index 8dc614cfc..795e00410 100644
--- a/cosmos/operators/docker.py
+++ b/cosmos/operators/docker.py
@@ -57,7 +57,13 @@ def __init__(
 
         super().__init__(image=image, **kwargs)
 
-    def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> Any:
+    def build_and_run_cmd(
+        self,
+        context: Context,
+        cmd_flags: list[str] | None = None,
+        run_as_async: bool = False,
+        async_context: dict[str, Any] | None = None,
+    ) -> Any:
         self.build_command(context, cmd_flags)
         self.log.info(f"Running command: {self.command}")
         result = DockerOperator.execute(self, context)
diff --git a/cosmos/operators/gcp_cloud_run_job.py b/cosmos/operators/gcp_cloud_run_job.py
index ef47db2cc..edb9d4954 100644
--- a/cosmos/operators/gcp_cloud_run_job.py
+++ b/cosmos/operators/gcp_cloud_run_job.py
@@ -70,7 +70,13 @@ def __init__(
         self.environment_variables = environment_variables or DEFAULT_ENVIRONMENT_VARIABLES
         super().__init__(project_id=project_id, region=region, job_name=job_name, **kwargs)
 
-    def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> Any:
+    def build_and_run_cmd(
+        self,
+        context: Context,
+        cmd_flags: list[str] | None = None,
+        run_as_async: bool = False,
+        async_context: dict[str, Any] | None = None,
+    ) -> Any:
         self.build_command(context, cmd_flags)
         self.log.info(f"Running command: {self.command}")
         result = CloudRunExecuteJobOperator.execute(self, context)
diff --git a/cosmos/operators/kubernetes.py b/cosmos/operators/kubernetes.py
index f86925fde..8230b3dd4 100644
--- a/cosmos/operators/kubernetes.py
+++ b/cosmos/operators/kubernetes.py
@@ -68,7 +68,13 @@ def build_env_args(self, env: dict[str, str | bytes | PathLike[Any]]) -> None:
 
         self.env_vars: list[Any] = convert_env_vars(env_vars_dict)
 
-    def build_and_run_cmd(self, context: Context, cmd_flags: list[str] | None = None) -> Any:
+    def build_and_run_cmd(
+        self,
+        context: Context,
+        cmd_flags: list[str] | None = None,
+        run_as_async: bool = False,
+        async_context: dict[str, Any] | None = None,
+    ) -> Any:
         self.build_kube_args(context, cmd_flags)
         self.log.info(f"Running command: {self.arguments}")
         result = KubernetesPodOperator.execute(self, context)
diff --git a/cosmos/operators/virtualenv.py b/cosmos/operators/virtualenv.py
index 3bd54da99..2e8b70f3c 100644
--- a/cosmos/operators/virtualenv.py
+++ b/cosmos/operators/virtualenv.py
@@ -96,6 +96,8 @@ def run_command(
         cmd: list[str],
         env: dict[str, str | bytes | os.PathLike[Any]],
         context: Context,
+        run_as_async: bool = False,
+        async_context: dict[str, Any] | None = None,
     ) -> FullOutputSubprocessResult | dbtRunnerResult:
         # No virtualenv_dir set, so create a temporary virtualenv
         if self.virtualenv_dir is None or self.is_virtualenv_dir_temporary: