SnowflakeSqlApiOperator - Change the base class (#31751)

* Change base class for SnowflakeSqlApiOperator * Remove unwanted check
apache · Jun 11, 2023 · 35845f9 · 35845f9
1 parent 3206719
commit 35845f9
Showing 1 changed file with 13 additions and 20 deletions.
diff --git a/airflow/providers/snowflake/operators/snowflake.py b/airflow/providers/snowflake/operators/snowflake.py
@@ -370,7 +370,7 @@ def __init__(
         self.query_ids: list[str] = []
 
 
-class SnowflakeSqlApiOperator(SnowflakeOperator):
+class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
     """
     Implemented Snowflake SQL API Operator to support multiple SQL statements sequentially,
     which is the behavior of the SnowflakeOperator, the Snowflake SQL API allows submitting
@@ -459,25 +459,18 @@ def __init__(
         self.token_renewal_delta = token_renewal_delta
         self.bindings = bindings
         self.execute_async = False
-        if self.__class__.__base__.__name__ != "SnowflakeOperator":
-            # It's better to do str check of the parent class name because currently SnowflakeOperator
-            # is deprecated and in future OSS SnowflakeOperator may be removed
-            if any(
-                [warehouse, database, role, schema, authenticator, session_parameters]
-            ):  # pragma: no cover
-                hook_params = kwargs.pop("hook_params", {})  # pragma: no cover
-                kwargs["hook_params"] = {
-                    "warehouse": warehouse,
-                    "database": database,
-                    "role": role,
-                    "schema": schema,
-                    "authenticator": authenticator,
-                    "session_parameters": session_parameters,
-                    **hook_params,
-                }
-            super().__init__(conn_id=snowflake_conn_id, **kwargs)  # pragma: no cover
-        else:
-            super().__init__(**kwargs)
+        if any([warehouse, database, role, schema, authenticator, session_parameters]):  # pragma: no cover
+            hook_params = kwargs.pop("hook_params", {})  # pragma: no cover
+            kwargs["hook_params"] = {
+                "warehouse": warehouse,
+                "database": database,
+                "role": role,
+                "schema": schema,
+                "authenticator": authenticator,
+                "session_parameters": session_parameters,
+                **hook_params,
+            }
+        super().__init__(conn_id=snowflake_conn_id, **kwargs)  # pragma: no cover
 
     def execute(self, context: Context) -> None:
         """