Replace deprecated PostgresOperator

Author: greenape

flowetl/flowetl/flowetl/mixins/fixed_sql_mixin.py

@@ -34,8 +34,8 @@ def fixed_sql_operator(*, class_name: str, sql: str, is_sensor: bool = False) ->
     if is_sensor:
         from airflow.sensors.sql import SqlSensor as op_base
     else:
-        from airflow.providers.postgres.operators.postgres import (
-            PostgresOperator as op_base,
+        from airflow.providers.common.sql.operators.sql import (
+            SQLExecuteQueryOperator as op_base,
         )
 
     return type(

flowetl/flowetl/flowetl/mixins/fixed_sql_with_params_mixin.py

@@ -43,8 +43,8 @@ def fixed_sql_operator_with_params(
     if is_sensor:
         from airflow.sensors.sql import SqlSensor as op_base
     else:
-        from airflow.providers.postgres.operators.postgres import (
-            PostgresOperator as op_base,
+        from airflow.providers.common.sql.operators.sql import (
+            SQLExecuteQueryOperator as op_base,
         )
 
     return type(

flowetl/flowetl/flowetl/mixins/wrapping_sql_mixin.py

@@ -36,11 +36,11 @@ def wrapped_sql_operator(*, class_name: str, sql: str) -> Type:
 
     """
     from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
-    from airflow.providers.postgres.operators.postgres import PostgresOperator
+    from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 
     return type(
         class_name,
-        (TableNameMacrosMixin, WrappingSQLMixin, PostgresOperator),
+        (TableNameMacrosMixin, WrappingSQLMixin, SQLExecuteQueryOperator),
         dict(
             wrapper_sql=sql,
         ),

flowetl/flowetl/flowetl/operators/analyze_operator.py

@@ -1,11 +1,11 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
-from airflow.providers.postgres.operators.postgres import PostgresOperator
+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
 
 
-class AnalyzeOperator(TableNameMacrosMixin, PostgresOperator):
+class AnalyzeOperator(TableNameMacrosMixin, SQLExecuteQueryOperator):
     """
     The analyze operator triggers the postgres analyze command on a table.
 
@@ -14,7 +14,7 @@ class AnalyzeOperator(TableNameMacrosMixin, PostgresOperator):
     target : str
         jinja templated schema qualified table name.
     kwargs : dict
-        Passed to airflow.operators.postgres_operator.PostgresOperator
+        Passed to airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator
     """
 
     def __init__(self, *, target: str, **kwargs) -> None:

flowetl/flowetl/flowetl/operators/create_foreign_staging_table_operator.py

@@ -4,11 +4,11 @@
 
 from typing import Dict, Optional
 
-from airflow.providers.postgres.operators.postgres import PostgresOperator
+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
 
 
-class CreateForeignStagingTableOperator(TableNameMacrosMixin, PostgresOperator):
+class CreateForeignStagingTableOperator(TableNameMacrosMixin, SQLExecuteQueryOperator):
     def __init__(
         self,
         *,
@@ -49,7 +49,7 @@ def __init__(
         encoding : str or None, default None
             String giving encoding type. Uses system locale by default.
         kwargs : dict
-            Passed to airflow.operators.postgres_operator.PostgresOperator
+            Passed to airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator
         """
         # Using an f-string here because filename needs to be templated, which is won't be if it is a param
         sql = f"""

flowetl/mounts/dags/test_dag.py

@@ -10,12 +10,12 @@
     catchup=True,
 ) as dag:
     from airflow.operators.bash import BashOperator
-    from airflow.providers.postgres.operators.postgres import PostgresOperator
+    from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 
     flowetl_test_op = BashOperator(
         task_id="flowetl_install_test_op", bash_command="date"
     )
-    flowdb_test_op = PostgresOperator(
+    flowdb_test_op = SQLExecuteQueryOperator(
         task_id="flowdb_connect_test_op",
         sql="SELECT * FROM geography.geo_kinds",
         postgres_conn_id="flowdb",

flowetl/tests/unit/test_fixed_sql_mixin.py

@@ -4,15 +4,15 @@
 
 
 def test_fixed_sql():
-    from airflow.providers.postgres.operators.postgres import PostgresOperator
+    from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
     from airflow.sensors.sql import SqlSensor
     from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
     from flowetl.mixins.fixed_sql_mixin import fixed_sql_operator
 
     new_type = fixed_sql_operator(class_name="DUMMY_TYPE", sql="FIXED_SQL")
     new_instance = new_type(task_id="DUMMY")
     assert new_type.fixed_sql == "FIXED_SQL"
-    assert isinstance(new_instance, PostgresOperator)
+    assert isinstance(new_instance, SQLExecuteQueryOperator)
     assert isinstance(new_instance, TableNameMacrosMixin)
     assert type(new_instance).__name__ == "DUMMY_TYPE"
     new_instance.prepare_template()

flowetl/tests/unit/test_params_mixin.py

@@ -4,7 +4,7 @@
 
 
 def test_fixed_sql_with_params():
-    from airflow.providers.postgres.operators.postgres import PostgresOperator
+    from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
     from airflow.sensors.sql import SqlSensor
 
     from flowetl.mixins.fixed_sql_with_params_mixin import (
@@ -18,7 +18,7 @@ def test_fixed_sql_with_params():
     new_instance = new_type(task_id="DUMMY", DUMMY_PARAM="DUMMY_PARAM_VALUE")
     assert new_type.fixed_sql == "FIXED_SQL"
     assert new_type.named_params == ["DUMMY_PARAM"]
-    assert isinstance(new_instance, PostgresOperator)
+    assert isinstance(new_instance, SQLExecuteQueryOperator)
     assert isinstance(new_instance, TableNameMacrosMixin)
     assert type(new_instance).__name__ == "DUMMY_TYPE"
     assert {k: v.value for k, v in new_instance.params.items()} == dict(

flowetl/tests/unit/test_wrapping_sql_mixin.py

@@ -4,14 +4,14 @@
 
 
 def test_wrapped_sql(mock_basic_dag):
-    from airflow.providers.postgres.operators.postgres import PostgresOperator
+    from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
     from flowetl.mixins.table_name_macros_mixin import TableNameMacrosMixin
     from flowetl.mixins.wrapping_sql_mixin import wrapped_sql_operator
 
     new_type = wrapped_sql_operator(class_name="DUMMY_TYPE", sql="WRAPS({sql})")
     new_instance = new_type(sql="WRAPPED", task_id="DUMMY", dag=mock_basic_dag)
     assert new_type.wrapper_sql == "WRAPS({sql})"
-    assert isinstance(new_instance, PostgresOperator)
+    assert isinstance(new_instance, SQLExecuteQueryOperator)
     assert isinstance(new_instance, TableNameMacrosMixin)
     assert type(new_instance).__name__ == "DUMMY_TYPE"
     new_instance.prepare_template()