diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index 93688d74cd67c..22bd76babb3fd 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -138,6 +138,14 @@ def copy_expert(self, sql: str, filename: str) -> None: file.truncate(file.tell()) conn.commit() + def get_uri(self) -> str: + conn = self.get_connection(getattr(self, self.conn_name_attr)) + uri = super().get_uri() + if conn.extra_dejson.get('client_encoding', False): + charset = conn.extra_dejson["client_encoding"] + return f"{uri}?client_encoding={charset}" + return uri + def bulk_load(self, table: str, tmp_file: str) -> None: """Loads a tab-delimited file into a database table""" self.copy_expert(f"COPY {table} FROM STDIN", tmp_file) diff --git a/docs/apache-airflow-providers-postgres/connections/postgres.rst b/docs/apache-airflow-providers-postgres/connections/postgres.rst index 11e469d8a4cbc..f777c851dd716 100644 --- a/docs/apache-airflow-providers-postgres/connections/postgres.rst +++ b/docs/apache-airflow-providers-postgres/connections/postgres.rst @@ -57,6 +57,8 @@ Extra (optional) configuration parameter. * ``keepalives_idle`` - Controls the number of seconds of inactivity after which TCP should send a keepalive message to the server. + * ``client_encoding``: specifies client encoding(character set) of the client connection. + Refer to `Postgres supported character sets `_ More details on all Postgres parameters supported can be found in `Postgres documentation `_. diff --git a/tests/providers/postgres/hooks/test_postgres.py b/tests/providers/postgres/hooks/test_postgres.py index c22a7223cefe1..56f8288518828 100644 --- a/tests/providers/postgres/hooks/test_postgres.py +++ b/tests/providers/postgres/hooks/test_postgres.py @@ -16,7 +16,7 @@ # specific language governing permissions and limitations # under the License. # - +import json import unittest from tempfile import NamedTemporaryFile from unittest import mock @@ -57,6 +57,17 @@ def test_get_conn(self, mock_connect): user='login', password='password', host='host', dbname='schema', port=None ) + @mock.patch('airflow.providers.postgres.hooks.postgres.psycopg2.connect') + def test_get_uri(self, mock_connect): + self.connection.extra = json.dumps({'client_encoding': 'utf-8'}) + self.connection.conn_type = 'postgres' + self.db_hook.get_conn() + assert mock_connect.call_count == 1 + + self.assertEqual( + self.db_hook.get_uri(), "postgres://login:password@host/schema?client_encoding=utf-8" + ) + @mock.patch('airflow.providers.postgres.hooks.postgres.psycopg2.connect') def test_get_conn_cursor(self, mock_connect): self.connection.extra = '{"cursor": "dictcursor"}'