Repository: apache/airflow

Issue:
Allow to set statement behavior for PostgresOperator
### Body

Add the ability to pass parameters like `statement_timeout` from PostgresOperator.
https://www.postgresql.org/docs/14/runtime-config-client.html#GUC-STATEMENT-TIMEOUT

The goal is to allow to control over specific query rather than setting the parameters on the connection level.

### Committer

- [X] I acknowledge that I am a maintainer/committer of the Apache Airflow project.


Discussion:
Cool. I assume this is about the possibility of overriding it for individual statements? Because we already can add "extra options" in the `connection" and they will be passed to the "client.connect" ?  
Ah yeah. You wrote about :) . Note to self: read before you comment @potiuk 
Hi, May I work on this issue?
@victorphoenix3 sure! assigned to you

Resolving patch:
diff --git a/airflow/providers/postgres/example_dags/example_postgres.py b/airflow/providers/postgres/example_dags/example_postgres.py
--- a/airflow/providers/postgres/example_dags/example_postgres.py
+++ b/airflow/providers/postgres/example_dags/example_postgres.py
@@ -64,14 +64,11 @@
     # [START postgres_operator_howto_guide_get_birth_date]
     get_birth_date = PostgresOperator(
         task_id="get_birth_date",
-        sql="""
-            SELECT * FROM pet
-            WHERE birth_date
-            BETWEEN SYMMETRIC DATE '{{ params.begin_date }}' AND DATE '{{ params.end_date }}';
-            """,
-        params={'begin_date': '2020-01-01', 'end_date': '2020-12-31'},
+        sql="SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC %(begin_date)s AND %(end_date)s",
+        parameters={"begin_date": "2020-01-01", "end_date": "2020-12-31"},
+        runtime_parameters={'statement_timeout': '3000ms'},
     )
-    # [START postgres_operator_howto_guide_get_birth_date]
+    # [END postgres_operator_howto_guide_get_birth_date]
 
     create_pet_table >> populate_pet_table >> get_all_pets >> get_birth_date
     # [END postgres_operator_howto_guide]
diff --git a/airflow/providers/postgres/operators/postgres.py b/airflow/providers/postgres/operators/postgres.py
--- a/airflow/providers/postgres/operators/postgres.py
+++ b/airflow/providers/postgres/operators/postgres.py
@@ -17,6 +17,8 @@
 # under the License.
 from typing import TYPE_CHECKING, Iterable, List, Mapping, Optional, Sequence, Union
 
+from psycopg2.sql import SQL, Identifier
+
 from airflow.models import BaseOperator
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 
@@ -52,6 +54,7 @@ def __init__(
         autocommit: bool = False,
         parameters: Optional[Union[Mapping, Iterable]] = None,
         database: Optional[str] = None,
+        runtime_parameters: Optional[Mapping] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -60,10 +63,28 @@ def __init__(
         self.autocommit = autocommit
         self.parameters = parameters
         self.database = database
+        self.runtime_parameters = runtime_parameters
         self.hook: Optional[PostgresHook] = None
 
     def execute(self, context: 'Context'):
         self.hook = PostgresHook(postgres_conn_id=self.postgres_conn_id, schema=self.database)
-        self.hook.run(self.sql, self.autocommit, parameters=self.parameters)
+        if self.runtime_parameters:
+            final_sql = []
+            sql_param = {}
+            for param in self.runtime_parameters:
+                set_param_sql = f"SET {{}} TO %({param})s;"
+                dynamic_sql = SQL(set_param_sql).format(Identifier(f"{param}"))
+                final_sql.append(dynamic_sql)
+            for param, val in self.runtime_parameters.items():
+                sql_param.update({f"{param}": f"{val}"})
+            if self.parameters:
+                sql_param.update(self.parameters)
+            if isinstance(self.sql, str):
+                final_sql.append(SQL(self.sql))
+            else:
+                final_sql.extend(list(map(SQL, self.sql)))
+            self.hook.run(final_sql, self.autocommit, parameters=sql_param)
+        else:
+            self.hook.run(self.sql, self.autocommit, parameters=self.parameters)
         for output in self.hook.conn.notices:
             self.log.info(output)


Test patch:
diff --git a/tests/providers/postgres/operators/test_postgres.py b/tests/providers/postgres/operators/test_postgres.py
--- a/tests/providers/postgres/operators/test_postgres.py
+++ b/tests/providers/postgres/operators/test_postgres.py
@@ -98,3 +98,18 @@ def test_overwrite_schema(self):
             op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
         except OperationalError as e:
             assert 'database "foobar" does not exist' in str(e)
+
+    def test_runtime_parameter_setting(self):
+        """
+        Verifies ability to pass server configuration parameters to
+        PostgresOperator
+        """
+
+        sql = "SELECT 1;"
+        op = PostgresOperator(
+            task_id='postgres_operator_test_runtime_parameter_setting',
+            sql=sql,
+            dag=self.dag,
+            runtime_parameters={'statement_timeout': '3000ms'},
+        )
+        op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)