From a1c532eacfeddcbefaa3e565a0522e25315286c4 Mon Sep 17 00:00:00 2001 From: subkanthi Date: Mon, 22 Nov 2021 13:36:39 -0500 Subject: [PATCH 1/8] [16185] Added LocalKubernetesExecutor to breeze supported executors --- BREEZE.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/BREEZE.rst b/BREEZE.rst index a62a9734cd3c1..b6a14240ffd33 100644 --- a/BREEZE.rst +++ b/BREEZE.rst @@ -2422,6 +2422,7 @@ This is the current syntax for `./breeze <./breeze>`_: One of: KubernetesExecutor CeleryExecutor LocalExecutor CeleryKubernetesExecutor + LocalKubernetesExecutor Default: KubernetesExecutor From d16a343c1478eb068e01d7ef363d91bb704918f8 Mon Sep 17 00:00:00 2001 From: subkanthi Date: Mon, 22 Nov 2021 13:38:51 -0500 Subject: [PATCH 2/8] Revert "[16185] Added LocalKubernetesExecutor to breeze supported executors" This reverts commit a1c532eacfeddcbefaa3e565a0522e25315286c4. --- BREEZE.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/BREEZE.rst b/BREEZE.rst index b6a14240ffd33..a62a9734cd3c1 100644 --- a/BREEZE.rst +++ b/BREEZE.rst @@ -2422,7 +2422,6 @@ This is the current syntax for `./breeze <./breeze>`_: One of: KubernetesExecutor CeleryExecutor LocalExecutor CeleryKubernetesExecutor - LocalKubernetesExecutor Default: KubernetesExecutor From e7faadd5d65c68978cba1e80d19be5a16680eeca Mon Sep 17 00:00:00 2001 From: subkanthi Date: Mon, 22 Nov 2021 20:18:34 -0500 Subject: [PATCH 3/8] Override get_uri in PostgresHook to support extra configs --- airflow/providers/postgres/hooks/postgres.py | 9 +++++++++ tests/providers/postgres/hooks/test_postgres.py | 11 ++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index 93688d74cd67c..54c20c2148577 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -138,6 +138,15 @@ 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('charset', False): + charset = conn.extra_dejson["charset"] + return f"{uri}?charset={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/tests/providers/postgres/hooks/test_postgres.py b/tests/providers/postgres/hooks/test_postgres.py index c22a7223cefe1..a9f32a81fcb96 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,15 @@ 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({'charset': '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?charset=utf-8") + @mock.patch('airflow.providers.postgres.hooks.postgres.psycopg2.connect') def test_get_conn_cursor(self, mock_connect): self.connection.extra = '{"cursor": "dictcursor"}' From 441c9918a513c7155fcf1f51fa7917815876e6ad Mon Sep 17 00:00:00 2001 From: subkanthi Date: Mon, 22 Nov 2021 20:23:46 -0500 Subject: [PATCH 4/8] Override get_uri in PostgresHook to support extra configs --- airflow/providers/postgres/hooks/postgres.py | 6 +++--- tests/providers/postgres/hooks/test_postgres.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index 54c20c2148577..178276999c1c5 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -142,9 +142,9 @@ def copy_expert(self, sql: str, filename: str) -> None: def get_uri(self) -> str: conn = self.get_connection(getattr(self, self.conn_name_attr)) uri = super().get_uri() - if conn.extra_dejson.get('charset', False): - charset = conn.extra_dejson["charset"] - return f"{uri}?charset={charset}" + 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: diff --git a/tests/providers/postgres/hooks/test_postgres.py b/tests/providers/postgres/hooks/test_postgres.py index a9f32a81fcb96..990d0fd3609e9 100644 --- a/tests/providers/postgres/hooks/test_postgres.py +++ b/tests/providers/postgres/hooks/test_postgres.py @@ -59,12 +59,12 @@ def test_get_conn(self, mock_connect): @mock.patch('airflow.providers.postgres.hooks.postgres.psycopg2.connect') def test_get_uri(self, mock_connect): - self.connection.extra = json.dumps({'charset': 'utf-8'}) + 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?charset=utf-8") + 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): From 4d4deb0163d3fbc549c7a1aeec320f4e760fa2e0 Mon Sep 17 00:00:00 2001 From: subkanthi Date: Fri, 26 Nov 2021 13:37:25 -0500 Subject: [PATCH 5/8] Fixed flake8 error --- tests/providers/postgres/hooks/test_postgres.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/providers/postgres/hooks/test_postgres.py b/tests/providers/postgres/hooks/test_postgres.py index 990d0fd3609e9..56f8288518828 100644 --- a/tests/providers/postgres/hooks/test_postgres.py +++ b/tests/providers/postgres/hooks/test_postgres.py @@ -64,7 +64,9 @@ def test_get_uri(self, mock_connect): 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") + 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): From d0af5e178c0c58f394bcdf924a467d2ca75db09e Mon Sep 17 00:00:00 2001 From: subkanthi Date: Fri, 26 Nov 2021 13:50:46 -0500 Subject: [PATCH 6/8] Add client_encoding to postgres document --- .../apache-airflow-providers-postgres/connections/postgres.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/apache-airflow-providers-postgres/connections/postgres.rst b/docs/apache-airflow-providers-postgres/connections/postgres.rst index 11e469d8a4cbc..00fb1649f0ff8 100644 --- a/docs/apache-airflow-providers-postgres/connections/postgres.rst +++ b/docs/apache-airflow-providers-postgres/connections/postgres.rst @@ -57,6 +57,9 @@ 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 https://www.postgresql.org/docs/9.3/multibyte.html for supported character set + values. More details on all Postgres parameters supported can be found in `Postgres documentation `_. From 5f973732aa37dac7f0092cef9ef711aa5a09ba43 Mon Sep 17 00:00:00 2001 From: subkanthi Date: Fri, 26 Nov 2021 13:53:28 -0500 Subject: [PATCH 7/8] Add client_encoding to postgres document --- .../apache-airflow-providers-postgres/connections/postgres.rst | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-postgres/connections/postgres.rst b/docs/apache-airflow-providers-postgres/connections/postgres.rst index 00fb1649f0ff8..f777c851dd716 100644 --- a/docs/apache-airflow-providers-postgres/connections/postgres.rst +++ b/docs/apache-airflow-providers-postgres/connections/postgres.rst @@ -58,8 +58,7 @@ Extra (optional) * ``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 https://www.postgresql.org/docs/9.3/multibyte.html for supported character set - values. + Refer to `Postgres supported character sets `_ More details on all Postgres parameters supported can be found in `Postgres documentation `_. From 0f6d4a58c629d821fd3c2a4895ee49898d104374 Mon Sep 17 00:00:00 2001 From: subkanthi Date: Fri, 26 Nov 2021 14:26:40 -0500 Subject: [PATCH 8/8] Fixed static checks --- airflow/providers/postgres/hooks/postgres.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index 178276999c1c5..22bd76babb3fd 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -138,7 +138,6 @@ 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()