fix validate-operators-init static checks

apache · Oct 17, 2024 · 2d01221 · 2d01221
1 parent af20925
commit 2d01221
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py
@@ -123,6 +123,8 @@ def __init__(
         forecast_horizon: int,
         data_granularity_unit: str,
         data_granularity_count: int,
+        display_name: str,
+        model_display_name: str | None = None,
         optimization_objective: str | None = None,
         column_specs: dict[str, str] | None = None,
         column_transformations: list[dict[str, dict[str, str]]] | None = None,
@@ -145,7 +147,12 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
-            region=region, impersonation_chain=impersonation_chain, parent_model=parent_model, **kwargs
+            display_name=display_name,
+            model_display_name=model_display_name,
+            region=region,
+            impersonation_chain=impersonation_chain,
+            parent_model=parent_model,
+            **kwargs,
         )
         self.dataset_id = dataset_id
         self.target_column = target_column

diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py
@@ -509,6 +509,8 @@ def __init__(
         *,
         command: Sequence[str] = [],
         region: str,
+        display_name: str,
+        model_display_name: str | None = None,
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
@@ -517,6 +519,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
+            display_name=display_name,
+            model_display_name=model_display_name,
             region=region,
             parent_model=parent_model,
             impersonation_chain=impersonation_chain,
@@ -962,6 +966,8 @@ def __init__(
         python_package_gcs_uri: str,
         python_module_name: str,
         region: str,
+        display_name: str,
+        model_display_name: str | None = None,
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
@@ -970,6 +976,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
+            display_name=display_name,
+            model_display_name=model_display_name,
             region=region,
             parent_model=parent_model,
             impersonation_chain=impersonation_chain,
@@ -1423,6 +1431,8 @@ def __init__(
         script_path: str,
         requirements: Sequence[str] | None = None,
         region: str,
+        display_name: str,
+        model_display_name: str | None = None,
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
@@ -1431,6 +1441,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(
+            display_name=display_name,
+            model_display_name=model_display_name,
             region=region,
             parent_model=parent_model,
             impersonation_chain=impersonation_chain,