Operator argument fix

cosmos/operators/_asynchronous/databricks.py

@@ -3,12 +3,13 @@
 
 from typing import Any
 
-try:  # Airflow 3
-    from airflow.sdk.bases.operator import BaseOperator
-except ImportError:  # Airflow 2
-    from airflow.models import BaseOperator
 from airflow.utils.context import Context
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
 
 class DbtRunAirflowAsyncDatabricksOperator(BaseOperator):  # type: ignore[misc]
     def __init__(self, *args: Any, **kwargs: Any):

cosmos/operators/aws_ecs.py

@@ -29,6 +29,11 @@
 DEFAULT_CONTAINER_NAME = "dbt"
 DEFAULT_ENVIRONMENT_VARIABLES: dict[str, str] = {}
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
 try:
     from airflow.providers.amazon.aws.operators.ecs import EcsRunTaskOperator
 except ImportError:  # pragma: no cover
@@ -76,21 +81,38 @@ def __init__(
                 "overrides": None,
             }
         )
-        super().__init__(**kwargs)
+
         # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
         # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
         # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
         # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
         # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
         # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
-        base_operator_args = set(inspect.signature(EcsRunTaskOperator.__init__).parameters.keys())
+        default_args = kwargs.get("default_args", {})
+        operator_kwargs = {}
+
+        operator_args: set[str] = set()
+        for clazz in EcsRunTaskOperator.__mro__:
+            operator_args.update(inspect.signature(clazz.__init__).parameters.keys())
+            if clazz == BaseOperator:
+                break
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
         base_kwargs = {}
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in base_operator_args:
-                base_kwargs[arg_key] = arg_value
-        base_kwargs["task_id"] = kwargs["task_id"]
-        base_kwargs["aws_conn_id"] = aws_conn_id
-        EcsRunTaskOperator.__init__(self, **base_kwargs)
+        for arg in {*inspect.signature(AbstractDbtBase.__init__).parameters.keys()}:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+        AbstractDbtBase.__init__(self, **base_kwargs)
+        EcsRunTaskOperator.__init__(self, **operator_kwargs)
 
     def build_and_run_cmd(
         self,

cosmos/operators/azure_container_instance.py

@@ -3,6 +3,10 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
 if TYPE_CHECKING:  # pragma: no cover
     try:
         from airflow.sdk.definitions.context import Context
@@ -68,20 +72,37 @@ def __init__(
                 "registry_conn_id": registry_conn_id,
             }
         )
-        super().__init__(**kwargs)
         # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
         # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
         # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
         # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
         # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
         # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
-        base_operator_args = set(inspect.signature(AzureContainerInstancesOperator.__init__).parameters.keys())
+
+        default_args = kwargs.get("default_args", {})
+        operator_kwargs = {}
+        operator_args: set[str] = set()
+        for clazz in AzureContainerInstancesOperator.__mro__:
+            operator_args.update(inspect.signature(clazz.__init__).parameters.keys())
+            if clazz == BaseOperator:
+                break
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
         base_kwargs = {}
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in base_operator_args:
-                base_kwargs[arg_key] = arg_value
-        base_kwargs["task_id"] = kwargs["task_id"]
-        AzureContainerInstancesOperator.__init__(self, **base_kwargs)
+        for arg in {*inspect.signature(AbstractDbtBase.__init__).parameters.keys()}:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+        AbstractDbtBase.__init__(self, **base_kwargs)
+        AzureContainerInstancesOperator.__init__(self, **operator_kwargs)
 
     def build_and_run_cmd(
         self,

cosmos/operators/docker.py

@@ -3,6 +3,10 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Sequence
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
 if TYPE_CHECKING:  # pragma: no cover
     try:
         from airflow.sdk.definitions.context import Context
@@ -58,21 +62,38 @@ def __init__(
                 "Airflow connections are not available in the Docker container for the mapping to work."
             )
 
-        super().__init__(image=image, **kwargs)
         # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
         # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
         # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
         # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
         # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
         # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
         kwargs["image"] = image
-        base_operator_args = set(inspect.signature(DockerOperator.__init__).parameters.keys())
+
+        default_args = kwargs.get("default_args", {})
+        operator_kwargs = {}
+        operator_args: set[str] = set()
+        for clazz in DockerOperator.__mro__:
+            operator_args.update(inspect.signature(clazz.__init__).parameters.keys())
+            if clazz == BaseOperator:
+                break
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
         base_kwargs = {}
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in base_operator_args:
-                base_kwargs[arg_key] = arg_value
-        base_kwargs["task_id"] = kwargs["task_id"]
-        DockerOperator.__init__(self, **base_kwargs)
+        for arg in {*inspect.signature(AbstractDbtBase.__init__).parameters.keys()}:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+        AbstractDbtBase.__init__(self, **base_kwargs)
+        DockerOperator.__init__(self, **operator_kwargs)
 
     def build_and_run_cmd(
         self,

cosmos/operators/gcp_cloud_run_job.py

@@ -28,6 +28,11 @@
 
 DEFAULT_ENVIRONMENT_VARIABLES: dict[str, str] = {}
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
 try:
     from airflow.providers.google.cloud.operators.cloud_run import CloudRunExecuteJobOperator
 
@@ -72,7 +77,6 @@ def __init__(
         self.profile_config = profile_config
         self.command = command
         self.environment_variables = environment_variables or DEFAULT_ENVIRONMENT_VARIABLES
-        super().__init__(project_id=project_id, region=region, job_name=job_name, **kwargs)
         # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
         # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
         # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
@@ -84,17 +88,33 @@ def __init__(
                 "project_id": project_id,
                 "region": region,
                 "job_name": job_name,
-                "command": command,
-                "environment_variables": environment_variables,
             }
         )
-        base_operator_args = set(inspect.signature(CloudRunExecuteJobOperator.__init__).parameters.keys())
+
+        default_args = kwargs.get("default_args", {})
+        operator_kwargs = {}
+        operator_args: set[str] = set()
+        for clazz in CloudRunExecuteJobOperator.__mro__:
+            operator_args.update(inspect.signature(clazz.__init__).parameters.keys())
+            if clazz == BaseOperator:
+                break
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
         base_kwargs = {}
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in base_operator_args:
-                base_kwargs[arg_key] = arg_value
-        base_kwargs["task_id"] = kwargs["task_id"]
-        CloudRunExecuteJobOperator.__init__(self, **base_kwargs)
+        for arg in {*inspect.signature(AbstractDbtBase.__init__).parameters.keys()}:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+        AbstractDbtBase.__init__(self, **base_kwargs)
+        CloudRunExecuteJobOperator.__init__(self, **operator_kwargs)
 
     def build_and_run_cmd(
         self,

cosmos/operators/kubernetes.py

@@ -29,6 +29,11 @@
     DbtTestMixin,
 )
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
 DBT_NO_TESTS_MSG = "Nothing to do"
 DBT_WARN_MSG = "WARN"
 
@@ -64,20 +69,38 @@ class DbtKubernetesBaseOperator(AbstractDbtBase, KubernetesPodOperator):  # type
 
     def __init__(self, profile_config: ProfileConfig | None = None, **kwargs: Any) -> None:
         self.profile_config = profile_config
-        super().__init__(**kwargs)
+
         # In PR #1474, we refactored cosmos.operators.base.AbstractDbtBase to remove its inheritance from BaseOperator
         # and eliminated the super().__init__() call. This change was made to resolve conflicts in parent class
         # initializations while adding support for ExecutionMode.AIRFLOW_ASYNC. Operators under this mode inherit
         # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
         # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
         # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
-        base_operator_args = set(inspect.signature(KubernetesPodOperator.__init__).parameters.keys())
+        default_args = kwargs.get("default_args", {})
+        operator_kwargs = {}
+        operator_args: set[str] = set()
+        for clazz in KubernetesPodOperator.__mro__:
+            operator_args.update(inspect.signature(clazz.__init__).parameters.keys())
+            if clazz == BaseOperator:
+                break
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
         base_kwargs = {}
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in base_operator_args:
-                base_kwargs[arg_key] = arg_value
-        base_kwargs["task_id"] = kwargs["task_id"]
-        KubernetesPodOperator.__init__(self, **base_kwargs)
+        for arg in {*inspect.signature(AbstractDbtBase.__init__).parameters.keys()}:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+
+        AbstractDbtBase.__init__(self, **base_kwargs)
+        KubernetesPodOperator.__init__(self, **operator_kwargs)
 
     def build_env_args(self, env: dict[str, str | bytes | PathLike[Any]]) -> None:
         env_vars_dict: dict[str, str] = dict()

cosmos/operators/local.py

@@ -17,11 +17,6 @@
 import jinja2
 from airflow import DAG
 from airflow.exceptions import AirflowException, AirflowSkipException
-
-try:  # Airflow 3
-    from airflow.sdk.bases.operator import BaseOperator
-except ImportError:  # Airflow 2
-    from airflow.models import BaseOperator
 from airflow.models.taskinstance import TaskInstance
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -49,6 +44,11 @@
     remote_target_path_conn_id,
 )
 
+try:
+    from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
+except ImportError:
+    from airflow.models import BaseOperator  # Airflow 2
+
 try:
     from airflow.datasets import Dataset
     from openlineage.common.provider.dbt.local import DbtLocalArtifactProcessor
@@ -766,19 +766,31 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Airflow provider operators that enable deferrable SQL query execution. Since super().__init__() was removed
         # from AbstractDbtBase and different parent classes require distinct initialization arguments, we explicitly
         # initialize them (including the BaseOperator) here by segregating the required arguments for each parent class.
-        abstract_dbt_local_base_kwargs = {}
-        base_operator_kwargs = {}
-        abstract_dbt_local_base_args_keys = (
-            inspect.getfullargspec(AbstractDbtBase.__init__).args
-            + inspect.getfullargspec(AbstractDbtLocalBase.__init__).args
-        )
-        base_operator_args = set(inspect.signature(BaseOperator.__init__).parameters.keys())
-        for arg_key, arg_value in kwargs.items():
-            if arg_key in abstract_dbt_local_base_args_keys:
-                abstract_dbt_local_base_kwargs[arg_key] = arg_value
-            if arg_key in base_operator_args:
-                base_operator_kwargs[arg_key] = arg_value
-        AbstractDbtLocalBase.__init__(self, **abstract_dbt_local_base_kwargs)
+        base_kwargs = {}
+        operator_kwargs = {}
+        operator_args = {*inspect.signature(BaseOperator.__init__).parameters.keys()}
+
+        default_args = kwargs.get("default_args", {})
+
+        for arg in operator_args:
+            try:
+                operator_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                pass
+
+        for arg in {
+            *inspect.getfullargspec(AbstractDbtBase.__init__).args,
+            *inspect.getfullargspec(AbstractDbtLocalBase.__init__).args,
+        }:
+            try:
+                base_kwargs[arg] = kwargs[arg]
+            except KeyError:
+                try:
+                    base_kwargs[arg] = default_args[arg]
+                except KeyError:
+                    pass
+
+        AbstractDbtLocalBase.__init__(self, **base_kwargs)
         if AIRFLOW_VERSION.major < _AIRFLOW3_MAJOR_VERSION:
             if (
                 kwargs.get("emit_datasets", True)
@@ -791,12 +803,12 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 # error: Incompatible types in assignment (expression has type "list[DatasetAlias]", target has type "str")
                 dag_id = kwargs.get("dag")
                 task_group_id = kwargs.get("task_group")
-                base_operator_kwargs["outlets"] = [
+                operator_kwargs["outlets"] = [
                     DatasetAlias(name=get_dataset_alias_name(dag_id, task_group_id, self.task_id))
                 ]  # type: ignore
-        if "task_id" in base_operator_kwargs:
-            base_operator_kwargs.pop("task_id")
-        BaseOperator.__init__(self, task_id=self.task_id, **base_operator_kwargs)
+        if "task_id" in operator_kwargs:
+            operator_kwargs.pop("task_id")
+        BaseOperator.__init__(self, task_id=self.task_id, **operator_kwargs)
 
 
 class DbtBuildLocalOperator(DbtBuildMixin, DbtLocalBaseOperator):