apache-airflow-providers-postgres 6.7.0__tar.gz → 6.7.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (42) hide show
  1. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/PKG-INFO +6 -6
  2. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/README.rst +3 -3
  3. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/changelog.rst +12 -0
  4. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/index.rst +3 -3
  5. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/provider.yaml +2 -1
  6. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/pyproject.toml +3 -3
  7. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/__init__.py +1 -1
  8. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/hooks/postgres.py +28 -22
  9. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/hooks/test_postgres.py +135 -300
  10. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/.gitignore +0 -0
  11. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/LICENSE +0 -0
  12. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/NOTICE +0 -0
  13. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/.latest-doc-only-change.txt +0 -0
  14. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/commits.rst +0 -0
  15. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/conf.py +0 -0
  16. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/configurations-ref.rst +0 -0
  17. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/connections/postgres.rst +0 -0
  18. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/dialects.rst +0 -0
  19. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/installing-providers-from-sources.rst +0 -0
  20. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/integration-logos/Postgres.png +0 -0
  21. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/operators.rst +0 -0
  22. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/redirects.txt +0 -0
  23. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/docs/security.rst +0 -0
  24. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/__init__.py +0 -0
  25. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/__init__.py +0 -0
  26. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/assets/__init__.py +0 -0
  27. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/assets/postgres.py +0 -0
  28. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/dialects/__init__.py +0 -0
  29. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/dialects/postgres.py +0 -0
  30. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/get_provider_info.py +0 -0
  31. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/src/airflow/providers/postgres/hooks/__init__.py +0 -0
  32. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/conftest.py +0 -0
  33. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/system/__init__.py +0 -0
  34. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/system/postgres/__init__.py +0 -0
  35. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/system/postgres/example_postgres.py +0 -0
  36. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/__init__.py +0 -0
  37. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/__init__.py +0 -0
  38. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/assets/__init__.py +0 -0
  39. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/assets/test_postgres.py +0 -0
  40. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/dialects/__init__.py +0 -0
  41. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/dialects/test_postgres.py +0 -0
  42. {apache_airflow_providers_postgres-6.7.0 → apache_airflow_providers_postgres-6.7.1}/tests/unit/postgres/hooks/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: apache-airflow-providers-postgres
3
- Version: 6.7.0
3
+ Version: 6.7.1
4
4
  Summary: Provider package apache-airflow-providers-postgres for Apache Airflow
5
5
  Keywords: airflow-provider,postgres,airflow,integration
6
6
  Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -40,8 +40,8 @@ Requires-Dist: psycopg[binary]>=3.2.9 ; extra == "psycopg" and ( python_version
40
40
  Requires-Dist: psycopg[binary]>=3.3.3 ; extra == "psycopg" and ( python_version >= '3.14')
41
41
  Requires-Dist: sqlalchemy>=1.4.54 ; extra == "sqlalchemy"
42
42
  Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
43
- Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/changelog.html
44
- Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0
43
+ Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/changelog.html
44
+ Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1
45
45
  Project-URL: Mastodon, https://fosstodon.org/@airflow
46
46
  Project-URL: Slack Chat, https://s.apache.org/airflow-slack
47
47
  Project-URL: Source Code, https://github.com/apache/airflow
@@ -79,7 +79,7 @@ Provides-Extra: sqlalchemy
79
79
 
80
80
  Package ``apache-airflow-providers-postgres``
81
81
 
82
- Release: ``6.7.0``
82
+ Release: ``6.7.1``
83
83
 
84
84
 
85
85
  `PostgreSQL <https://www.postgresql.org/>`__
@@ -92,7 +92,7 @@ This is a provider package for ``postgres`` provider. All classes for this provi
92
92
  are in ``airflow.providers.postgres`` python package.
93
93
 
94
94
  You can find package information and changelog for the provider
95
- in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/>`_.
95
+ in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/>`_.
96
96
 
97
97
  Installation
98
98
  ------------
@@ -156,5 +156,5 @@ Extra Dependencies
156
156
  =================== ============================================================================================================================================================
157
157
 
158
158
  The changelog for the provider package can be found in the
159
- `changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/changelog.html>`_.
159
+ `changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/changelog.html>`_.
160
160
 
@@ -23,7 +23,7 @@
23
23
 
24
24
  Package ``apache-airflow-providers-postgres``
25
25
 
26
- Release: ``6.7.0``
26
+ Release: ``6.7.1``
27
27
 
28
28
 
29
29
  `PostgreSQL <https://www.postgresql.org/>`__
@@ -36,7 +36,7 @@ This is a provider package for ``postgres`` provider. All classes for this provi
36
36
  are in ``airflow.providers.postgres`` python package.
37
37
 
38
38
  You can find package information and changelog for the provider
39
- in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/>`_.
39
+ in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/>`_.
40
40
 
41
41
  Installation
42
42
  ------------
@@ -100,4 +100,4 @@ Extra Dependencies
100
100
  =================== ============================================================================================================================================================
101
101
 
102
102
  The changelog for the provider package can be found in the
103
- `changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/changelog.html>`_.
103
+ `changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/changelog.html>`_.
@@ -27,6 +27,18 @@
27
27
  Changelog
28
28
  ---------
29
29
 
30
+ 6.7.1
31
+ .....
32
+
33
+ Misc
34
+ ~~~~
35
+
36
+ * ``Refactor PostgresHook and associated runtime tests (#66893)``
37
+
38
+ .. Below changes are excluded from the changelog. Move them to
39
+ appropriate section above if needed. Do not delete the lines(!):
40
+
41
+
30
42
  6.7.0
31
43
  .....
32
44
 
@@ -78,7 +78,7 @@ apache-airflow-providers-postgres package
78
78
  `PostgreSQL <https://www.postgresql.org/>`__
79
79
 
80
80
 
81
- Release: 6.7.0
81
+ Release: 6.7.1
82
82
 
83
83
  Provider package
84
84
  ----------------
@@ -138,5 +138,5 @@ Downloading official packages
138
138
  You can download officially released packages and verify their checksums and signatures from the
139
139
  `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
140
140
 
141
- * `The apache-airflow-providers-postgres 6.7.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0.tar.gz.sha512>`__)
142
- * `The apache-airflow-providers-postgres 6.7.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.0-py3-none-any.whl.sha512>`__)
141
+ * `The apache-airflow-providers-postgres 6.7.1 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1.tar.gz.sha512>`__)
142
+ * `The apache-airflow-providers-postgres 6.7.1 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_postgres-6.7.1-py3-none-any.whl.sha512>`__)
@@ -23,12 +23,13 @@ description: |
23
23
 
24
24
  state: ready
25
25
  lifecycle: production
26
- source-date-epoch: 1779138287
26
+ source-date-epoch: 1780426378
27
27
  # Note that those versions are maintained by release manager - do not update them manually
28
28
  # with the exception of case where other provider in sources has >= new provider version.
29
29
  # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
30
30
  # to be done in the same PR
31
31
  versions:
32
+ - 6.7.1
32
33
  - 6.7.0
33
34
  - 6.6.3
34
35
  - 6.6.2
@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
25
25
 
26
26
  [project]
27
27
  name = "apache-airflow-providers-postgres"
28
- version = "6.7.0"
28
+ version = "6.7.1"
29
29
  description = "Provider package apache-airflow-providers-postgres for Apache Airflow"
30
30
  readme = "README.rst"
31
31
  license = "Apache-2.0"
@@ -137,8 +137,8 @@ apache-airflow-providers-common-sql = {workspace = true}
137
137
  apache-airflow-providers-standard = {workspace = true}
138
138
 
139
139
  [project.urls]
140
- "Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0"
141
- "Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.0/changelog.html"
140
+ "Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1"
141
+ "Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.7.1/changelog.html"
142
142
  "Bug Tracker" = "https://github.com/apache/airflow/issues"
143
143
  "Source Code" = "https://github.com/apache/airflow"
144
144
  "Slack Chat" = "https://s.apache.org/airflow-slack"
@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
29
29
 
30
30
  __all__ = ["__version__"]
31
31
 
32
- __version__ = "6.7.0"
32
+ __version__ = "6.7.1"
33
33
 
34
34
  if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
35
35
  "2.11.0"
@@ -230,6 +230,29 @@ class PostgresHook(DbApiHook):
230
230
  valid_cursors = ", ".join(cursor_types.keys())
231
231
  raise ValueError(f"Invalid cursor passed {_cursor}. Valid options are: {valid_cursors}")
232
232
 
233
+ def _get_cursor_config(self, raw_cursor: str) -> tuple[str, Any]:
234
+ cursor = self._get_cursor(raw_cursor)
235
+
236
+ if USE_PSYCOPG3:
237
+ return "row_factory", cursor
238
+
239
+ return "cursor_factory", cursor
240
+
241
+ def _create_connection(self, conn_args: dict[str, Any]) -> CompatConnection:
242
+ if USE_PSYCOPG3:
243
+ from psycopg.connection import Connection as pgConnection
244
+
245
+ connection = pgConnection.connect(**cast("Any", conn_args))
246
+
247
+ register_default_adapters(connection)
248
+
249
+ if self.enable_log_db_messages and hasattr(connection, "add_notice_handler"):
250
+ connection.add_notice_handler(self._notice_handler)
251
+
252
+ return connection
253
+
254
+ return ppg2_connect(**conn_args)
255
+
233
256
  def _generate_cursor_name(self):
234
257
  """Generate a unique name for server-side cursor."""
235
258
  import uuid
@@ -262,30 +285,13 @@ class PostgresHook(DbApiHook):
262
285
  if arg_name not in self.ignored_extra_options:
263
286
  conn_args[arg_name] = arg_val
264
287
 
265
- if USE_PSYCOPG3:
266
- from psycopg.connection import Connection as pgConnection
267
-
268
- raw_cursor = conn.extra_dejson.get("cursor")
269
- if raw_cursor:
270
- conn_args["row_factory"] = self._get_cursor(raw_cursor)
271
-
272
- # Use Any type for the connection args to avoid type conflicts
273
- connection = pgConnection.connect(**cast("Any", conn_args))
274
- self.conn = cast("CompatConnection", connection)
275
-
276
- # Register JSON handlers for both json and jsonb types
277
- # This ensures JSON data is properly decoded from bytes to Python objects
278
- register_default_adapters(connection)
288
+ raw_cursor = conn.extra_dejson.get("cursor")
279
289
 
280
- # Add the notice handler AFTER the connection is established
281
- if self.enable_log_db_messages and hasattr(self.conn, "add_notice_handler"):
282
- self.conn.add_notice_handler(self._notice_handler)
283
- else: # psycopg2
284
- raw_cursor = conn.extra_dejson.get("cursor", False)
285
- if raw_cursor:
286
- conn_args["cursor_factory"] = self._get_cursor(raw_cursor)
290
+ if raw_cursor:
291
+ key, value = self._get_cursor_config(raw_cursor)
292
+ conn_args[key] = value
287
293
 
288
- self.conn = cast("CompatConnection", ppg2_connect(**conn_args))
294
+ self.conn = self._create_connection(conn_args)
289
295
 
290
296
  return self.conn
291
297
 
@@ -19,7 +19,6 @@ from __future__ import annotations
19
19
 
20
20
  import json
21
21
  import os
22
- from types import SimpleNamespace
23
22
  from unittest import mock
24
23
 
25
24
  import pandas as pd
@@ -30,7 +29,7 @@ import sqlalchemy
30
29
  from airflow.models import Connection
31
30
  from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException
32
31
  from airflow.providers.postgres.dialects.postgres import PostgresDialect
33
- from airflow.providers.postgres.hooks.postgres import CompatConnection, PostgresHook
32
+ from airflow.providers.postgres.hooks.postgres import PostgresHook
34
33
 
35
34
  from tests_common.test_utils.common_sql import mock_db_hook
36
35
  from tests_common.test_utils.version_compat import NOTSET
@@ -59,36 +58,6 @@ else:
59
58
  import psycopg2.extras
60
59
 
61
60
 
62
- @pytest.fixture
63
- def postgres_hook_setup():
64
- """Set up mock PostgresHook for testing."""
65
- table = "test_postgres_hook_table"
66
- cur = mock.MagicMock(rowcount=0)
67
- conn = mock.MagicMock(spec=CompatConnection)
68
- conn.cursor.return_value = cur
69
-
70
- class UnitTestPostgresHook(PostgresHook):
71
- conn_name_attr = "test_conn_id"
72
-
73
- def get_conn(self):
74
- return conn
75
-
76
- db_hook = UnitTestPostgresHook()
77
-
78
- # Return a namespace with all the objects
79
- setup = SimpleNamespace(table=table, cur=cur, conn=conn, db_hook=db_hook)
80
-
81
- yield setup
82
-
83
- # Teardown - only for real database tests
84
- try:
85
- with PostgresHook().get_conn() as real_conn:
86
- with real_conn.cursor() as real_cur:
87
- real_cur.execute(f"DROP TABLE IF EXISTS {table}")
88
- except Exception:
89
- pass # Ignore cleanup errors for unit tests
90
-
91
-
92
61
  @pytest.fixture
93
62
  def mock_connect(mocker):
94
63
  """Mock the connection object according to the correct psycopg version."""
@@ -816,10 +785,8 @@ class TestPostgresHook:
816
785
  ) == INSERT_SQL_STATEMENT.format('"schema"')
817
786
 
818
787
 
819
- @pytest.mark.backend("postgres")
820
- @pytest.mark.skipif(USE_PSYCOPG3, reason="psycopg v3 is available")
821
- class TestPostgresHookPPG2:
822
- """PostgresHook tests that are specific to psycopg2."""
788
+ class _BasePostgresHookRuntimeTests:
789
+ """Shared runtime tests for psycopg2 and psycopg3."""
823
790
 
824
791
  table = "test_postgres_hook_table"
825
792
 
@@ -841,6 +808,121 @@ class TestPostgresHookPPG2:
841
808
  with conn.cursor() as cur:
842
809
  cur.execute(f"DROP TABLE IF EXISTS {self.table}")
843
810
 
811
+ def test_insert_rows(self):
812
+ table = "table"
813
+ rows = [("hello",), ("world",)]
814
+
815
+ self.db_hook.insert_rows(table, rows)
816
+
817
+ assert self.conn.close.call_count == 1
818
+ assert self.cur.close.call_count == 1
819
+ assert self.conn.commit.call_count == 2
820
+
821
+ sql = f"INSERT INTO {table} VALUES (%s)"
822
+ self.cur.executemany.assert_any_call(sql, rows)
823
+
824
+ def test_insert_rows_replace(self):
825
+ table = "table"
826
+ rows = [
827
+ (1, "hello"),
828
+ (2, "world"),
829
+ ]
830
+ fields = ("id", "value")
831
+
832
+ self.db_hook.insert_rows(
833
+ table,
834
+ rows,
835
+ fields,
836
+ replace=True,
837
+ replace_index=fields[0],
838
+ )
839
+
840
+ assert self.conn.close.call_count == 1
841
+ assert self.cur.close.call_count == 1
842
+ assert self.conn.commit.call_count == 2
843
+
844
+ sql = (
845
+ f"INSERT INTO {table} ({fields[0]}, {fields[1]}) VALUES (%s,%s) "
846
+ f"ON CONFLICT ({fields[0]}) DO UPDATE SET {fields[1]} = excluded.{fields[1]}"
847
+ )
848
+ self.cur.executemany.assert_any_call(sql, rows)
849
+
850
+ def test_insert_rows_replace_missing_target_field_arg(self):
851
+ table = "table"
852
+ rows = [
853
+ (1, "hello"),
854
+ (2, "world"),
855
+ ]
856
+ fields = ("id", "value")
857
+
858
+ with pytest.raises(
859
+ ValueError,
860
+ match="PostgreSQL ON CONFLICT upsert syntax requires column names",
861
+ ):
862
+ self.db_hook.insert_rows(
863
+ table,
864
+ rows,
865
+ replace=True,
866
+ replace_index=fields[0],
867
+ )
868
+
869
+ def test_insert_rows_replace_missing_replace_index_arg(self):
870
+ table = "table"
871
+ rows = [
872
+ (1, "hello"),
873
+ (2, "world"),
874
+ ]
875
+ fields = ("id", "value")
876
+
877
+ with pytest.raises(
878
+ ValueError,
879
+ match="PostgreSQL ON CONFLICT upsert syntax requires an unique index",
880
+ ):
881
+ self.db_hook.insert_rows(
882
+ table,
883
+ rows,
884
+ fields,
885
+ replace=True,
886
+ )
887
+
888
+ def test_insert_rows_replace_all_index(self):
889
+ table = "table"
890
+ rows = [
891
+ (1, "hello"),
892
+ (2, "world"),
893
+ ]
894
+ fields = ("id", "value")
895
+
896
+ self.db_hook.insert_rows(
897
+ table,
898
+ rows,
899
+ fields,
900
+ replace=True,
901
+ replace_index=fields,
902
+ )
903
+
904
+ assert self.conn.close.call_count == 1
905
+ assert self.cur.close.call_count == 1
906
+ assert self.conn.commit.call_count == 2
907
+
908
+ sql = (
909
+ f"INSERT INTO {table} ({', '.join(fields)}) VALUES (%s,%s) "
910
+ f"ON CONFLICT ({', '.join(fields)}) DO NOTHING"
911
+ )
912
+ self.cur.executemany.assert_any_call(sql, rows)
913
+
914
+ def test_dialect_name(self):
915
+ assert self.db_hook.dialect_name == "postgresql"
916
+
917
+ def test_dialect(self):
918
+ assert isinstance(self.db_hook.dialect, PostgresDialect)
919
+
920
+
921
+ @pytest.mark.backend("postgres")
922
+ @pytest.mark.skipif(USE_PSYCOPG3, reason="psycopg v3 is available")
923
+ class TestPostgresHookPPG2(_BasePostgresHookRuntimeTests):
924
+ """PostgresHook tests that are specific to psycopg2."""
925
+
844
926
  def test_copy_expert(self, mocker):
845
927
  open_mock = mocker.mock_open(read_data='{"some": "json"}')
846
928
  mocker.patch("airflow.providers.postgres.hooks.postgres.open", open_mock)
@@ -915,169 +997,59 @@ class TestPostgresHookPPG2:
915
997
  assert call_kw["sql"] == sql
916
998
  assert call_kw["sql_parameters"] == parameters
917
999
 
918
- def test_insert_rows(self, postgres_hook_setup):
919
- setup = postgres_hook_setup
920
- table = "table"
921
- rows = [("hello",), ("world",)]
922
-
923
- setup.db_hook.insert_rows(table, rows)
924
-
925
- assert setup.conn.close.call_count == 1
926
- assert setup.cur.close.call_count == 1
927
-
928
- commit_count = 2 # The first and last commit
929
- assert commit_count == setup.conn.commit.call_count
930
-
931
- sql = f"INSERT INTO {table} VALUES (%s)"
932
- setup.cur.executemany.assert_any_call(sql, rows)
933
-
934
1000
  @mock.patch("airflow.providers.common.sql.hooks.sql.send_sql_hook_lineage")
935
- def test_insert_rows_hook_lineage(self, mock_send_lineage, postgres_hook_setup):
936
- setup = postgres_hook_setup
1001
+ def test_insert_rows_hook_lineage(self, mock_send_lineage):
937
1002
  table = "table"
938
1003
  rows = [("hello",), ("world",)]
939
1004
 
940
- setup.db_hook.insert_rows(table, rows)
1005
+ self.db_hook.insert_rows(table, rows)
941
1006
 
942
1007
  mock_send_lineage.assert_called_once()
1008
+
943
1009
  call_kw = mock_send_lineage.call_args.kwargs
944
- assert call_kw["context"] is setup.db_hook
1010
+
1011
+ assert call_kw["context"] is self.db_hook
945
1012
  assert call_kw["sql"] == f"INSERT INTO {table} VALUES (%s)"
946
1013
  assert call_kw["row_count"] == 2
947
1014
 
948
1015
  @mock.patch("airflow.providers.postgres.hooks.postgres.execute_batch")
949
- def test_insert_rows_fast_executemany(self, mock_execute_batch, postgres_hook_setup):
950
- setup = postgres_hook_setup
1016
+ def test_insert_rows_fast_executemany(self, mock_execute_batch):
951
1017
  table = "table"
952
1018
  rows = [("hello",), ("world",)]
953
1019
 
954
- setup.db_hook.insert_rows(table, rows, fast_executemany=True)
1020
+ self.db_hook.insert_rows(table, rows, fast_executemany=True)
955
1021
 
956
- assert setup.conn.close.call_count == 1
957
- assert setup.cur.close.call_count == 1
1022
+ assert self.conn.close.call_count == 1
1023
+ assert self.cur.close.call_count == 1
958
1024
 
959
1025
  commit_count = 2 # The first and last commit
960
- assert setup.conn.commit.call_count == commit_count
1026
+ assert self.conn.commit.call_count == commit_count
961
1027
 
962
1028
  mock_execute_batch.assert_called_once_with(
963
- setup.cur,
1029
+ self.cur,
964
1030
  f"INSERT INTO {table} VALUES (%s)", # expected SQL
965
1031
  [("hello",), ("world",)], # expected values
966
1032
  page_size=1000,
967
1033
  )
968
1034
 
969
1035
  # executemany should NOT be called in this mode
970
- setup.cur.executemany.assert_not_called()
1036
+ self.cur.executemany.assert_not_called()
971
1037
 
972
1038
  @mock.patch("airflow.providers.postgres.hooks.postgres.send_sql_hook_lineage")
973
1039
  @mock.patch("airflow.providers.postgres.hooks.postgres.execute_batch")
974
- def test_insert_rows_fast_executemany_hook_lineage(
975
- self, mock_execute_batch, mock_send_lineage, postgres_hook_setup
976
- ):
977
- setup = postgres_hook_setup
1040
+ def test_insert_rows_fast_executemany_hook_lineage(self, mock_execute_batch, mock_send_lineage):
1041
+
978
1042
  table = "table"
979
1043
  rows = [("hello",), ("world",)]
980
1044
 
981
- setup.db_hook.insert_rows(table, rows, fast_executemany=True)
1045
+ self.db_hook.insert_rows(table, rows, fast_executemany=True)
982
1046
 
983
1047
  mock_send_lineage.assert_called_once()
984
1048
  call_kw = mock_send_lineage.call_args.kwargs
985
- assert call_kw["context"] is setup.db_hook
1049
+ assert call_kw["context"] is self.db_hook
986
1050
  assert call_kw["sql"] == f"INSERT INTO {table} VALUES (%s)"
987
1051
  assert call_kw["row_count"] == 2
988
1052
 
989
- def test_insert_rows_replace(self, postgres_hook_setup):
990
- setup = postgres_hook_setup
991
- table = "table"
992
- rows = [
993
- (
994
- 1,
995
- "hello",
996
- ),
997
- (
998
- 2,
999
- "world",
1000
- ),
1001
- ]
1002
- fields = ("id", "value")
1003
-
1004
- setup.db_hook.insert_rows(table, rows, fields, replace=True, replace_index=fields[0])
1005
-
1006
- assert setup.conn.close.call_count == 1
1007
- assert setup.cur.close.call_count == 1
1008
-
1009
- commit_count = 2 # The first and last commit
1010
- assert commit_count == setup.conn.commit.call_count
1011
-
1012
- sql = (
1013
- f"INSERT INTO {table} ({fields[0]}, {fields[1]}) VALUES (%s,%s) "
1014
- f"ON CONFLICT ({fields[0]}) DO UPDATE SET {fields[1]} = excluded.{fields[1]}"
1015
- )
1016
- setup.cur.executemany.assert_any_call(sql, rows)
1017
-
1018
- def test_insert_rows_replace_missing_target_field_arg(self, postgres_hook_setup):
1019
- setup = postgres_hook_setup
1020
- table = "table"
1021
- rows = [
1022
- (
1023
- 1,
1024
- "hello",
1025
- ),
1026
- (
1027
- 2,
1028
- "world",
1029
- ),
1030
- ]
1031
- fields = ("id", "value")
1032
- with pytest.raises(ValueError, match="PostgreSQL ON CONFLICT upsert syntax requires column names"):
1033
- setup.db_hook.insert_rows(table, rows, replace=True, replace_index=fields[0])
1034
-
1035
- def test_insert_rows_replace_missing_replace_index_arg(self, postgres_hook_setup):
1036
- setup = postgres_hook_setup
1037
- table = "table"
1038
- rows = [
1039
- (
1040
- 1,
1041
- "hello",
1042
- ),
1043
- (
1044
- 2,
1045
- "world",
1046
- ),
1047
- ]
1048
- fields = ("id", "value")
1049
- with pytest.raises(ValueError, match="PostgreSQL ON CONFLICT upsert syntax requires an unique index"):
1050
- setup.db_hook.insert_rows(table, rows, fields, replace=True)
1051
-
1052
- def test_insert_rows_replace_all_index(self, postgres_hook_setup):
1053
- setup = postgres_hook_setup
1054
- table = "table"
1055
- rows = [
1056
- (
1057
- 1,
1058
- "hello",
1059
- ),
1060
- (
1061
- 2,
1062
- "world",
1063
- ),
1064
- ]
1065
- fields = ("id", "value")
1066
-
1067
- setup.db_hook.insert_rows(table, rows, fields, replace=True, replace_index=fields)
1068
-
1069
- assert setup.conn.close.call_count == 1
1070
- assert setup.cur.close.call_count == 1
1071
-
1072
- commit_count = 2 # The first and last commit
1073
- assert commit_count == setup.conn.commit.call_count
1074
-
1075
- sql = (
1076
- f"INSERT INTO {table} ({', '.join(fields)}) VALUES (%s,%s) "
1077
- f"ON CONFLICT ({', '.join(fields)}) DO NOTHING"
1078
- )
1079
- setup.cur.executemany.assert_any_call(sql, rows)
1080
-
1081
1053
  @pytest.mark.usefixtures("reset_logging_config")
1082
1054
  def test_get_all_db_log_messages(self, mocker):
1083
1055
  messages = ["a", "b", "c"]
@@ -1120,40 +1092,12 @@ class TestPostgresHookPPG2:
1120
1092
  finally:
1121
1093
  hook.run(sql=f"DROP PROCEDURE {proc_name} (s text)")
1122
1094
 
1123
- def test_dialect_name(self, postgres_hook_setup):
1124
- setup = postgres_hook_setup
1125
- assert setup.db_hook.dialect_name == "postgresql"
1126
-
1127
- def test_dialect(self, postgres_hook_setup):
1128
- setup = postgres_hook_setup
1129
- assert isinstance(setup.db_hook.dialect, PostgresDialect)
1130
-
1131
1095
 
1132
1096
  @pytest.mark.backend("postgres")
1133
1097
  @pytest.mark.skipif(not USE_PSYCOPG3, reason="psycopg v3 or sqlalchemy v2 are not available")
1134
- class TestPostgresHookPPG3:
1098
+ class TestPostgresHookPPG3(_BasePostgresHookRuntimeTests):
1135
1099
  """PostgresHook tests that are specific to psycopg3."""
1136
1100
 
1137
- table = "test_postgres_hook_table"
1138
-
1139
- def setup_method(self):
1140
- self.cur = mock.MagicMock(rowcount=0)
1141
- self.conn = conn = mock.MagicMock()
1142
- self.conn.cursor.return_value = self.cur
1143
-
1144
- class UnitTestPostgresHook(PostgresHook):
1145
- conn_name_attr = "test_conn_id"
1146
-
1147
- def get_conn(self):
1148
- return conn
1149
-
1150
- self.db_hook = UnitTestPostgresHook()
1151
-
1152
- def teardown_method(self):
1153
- with PostgresHook().get_conn() as conn:
1154
- with conn.cursor() as cur:
1155
- cur.execute(f"DROP TABLE IF EXISTS {self.table}")
1156
-
1157
1101
  def test_copy_expert_from(self, mocker):
1158
1102
  """Tests copy_expert with a 'COPY FROM STDIN' operation."""
1159
1103
  statement = "COPY test_table FROM STDIN"
@@ -1235,109 +1179,6 @@ class TestPostgresHookPPG3:
1235
1179
  )
1236
1180
  self.conn.commit.assert_called_once()
1237
1181
 
1238
- def test_insert_rows(self):
1239
- table = "table"
1240
- rows = [("hello",), ("world",)]
1241
-
1242
- self.db_hook.insert_rows(table, rows)
1243
-
1244
- assert self.conn.close.call_count == 1
1245
- assert self.cur.close.call_count == 1
1246
-
1247
- commit_count = 2 # The first and last commit
1248
- assert commit_count == self.conn.commit.call_count
1249
-
1250
- sql = f"INSERT INTO {table} VALUES (%s)"
1251
- self.cur.executemany.assert_any_call(sql, rows)
1252
-
1253
- def test_insert_rows_replace(self):
1254
- table = "table"
1255
- rows = [
1256
- (
1257
- 1,
1258
- "hello",
1259
- ),
1260
- (
1261
- 2,
1262
- "world",
1263
- ),
1264
- ]
1265
- fields = ("id", "value")
1266
-
1267
- self.db_hook.insert_rows(table, rows, fields, replace=True, replace_index=fields[0])
1268
-
1269
- assert self.conn.close.call_count == 1
1270
- assert self.cur.close.call_count == 1
1271
-
1272
- commit_count = 2 # The first and last commit
1273
- assert commit_count == self.conn.commit.call_count
1274
-
1275
- sql = (
1276
- f"INSERT INTO {table} ({fields[0]}, {fields[1]}) VALUES (%s,%s) "
1277
- f"ON CONFLICT ({fields[0]}) DO UPDATE SET {fields[1]} = excluded.{fields[1]}"
1278
- )
1279
- self.cur.executemany.assert_any_call(sql, rows)
1280
-
1281
- def test_insert_rows_replace_missing_target_field_arg(self):
1282
- table = "table"
1283
- rows = [
1284
- (
1285
- 1,
1286
- "hello",
1287
- ),
1288
- (
1289
- 2,
1290
- "world",
1291
- ),
1292
- ]
1293
- fields = ("id", "value")
1294
- with pytest.raises(ValueError, match="PostgreSQL ON CONFLICT upsert syntax requires column names"):
1295
- self.db_hook.insert_rows(table, rows, replace=True, replace_index=fields[0])
1296
-
1297
- def test_insert_rows_replace_missing_replace_index_arg(self):
1298
- table = "table"
1299
- rows = [
1300
- (
1301
- 1,
1302
- "hello",
1303
- ),
1304
- (
1305
- 2,
1306
- "world",
1307
- ),
1308
- ]
1309
- fields = ("id", "value")
1310
- with pytest.raises(ValueError, match="PostgreSQL ON CONFLICT upsert syntax requires an unique index"):
1311
- self.db_hook.insert_rows(table, rows, fields, replace=True)
1312
-
1313
- def test_insert_rows_replace_all_index(self):
1314
- table = "table"
1315
- rows = [
1316
- (
1317
- 1,
1318
- "hello",
1319
- ),
1320
- (
1321
- 2,
1322
- "world",
1323
- ),
1324
- ]
1325
- fields = ("id", "value")
1326
-
1327
- self.db_hook.insert_rows(table, rows, fields, replace=True, replace_index=fields)
1328
-
1329
- assert self.conn.close.call_count == 1
1330
- assert self.cur.close.call_count == 1
1331
-
1332
- commit_count = 2 # The first and last commit
1333
- assert commit_count == self.conn.commit.call_count
1334
-
1335
- sql = (
1336
- f"INSERT INTO {table} ({', '.join(fields)}) VALUES (%s,%s) "
1337
- f"ON CONFLICT ({', '.join(fields)}) DO NOTHING"
1338
- )
1339
- self.cur.executemany.assert_any_call(sql, rows)
1340
-
1341
1182
  @pytest.mark.skip(reason="Notice handling is callback-based in psycopg3 and cannot be tested this way.")
1342
1183
  def test_get_all_db_log_messages(self, mocker):
1343
1184
  pass
@@ -1366,9 +1207,3 @@ class TestPostgresHookPPG3:
1366
1207
  mock_logger.info.assert_any_call("Message from db: 42")
1367
1208
  finally:
1368
1209
  hook.run(sql=f"DROP PROCEDURE {proc_name} (s text)")
1369
-
1370
- def test_dialect_name(self):
1371
- assert self.db_hook.dialect_name == "postgresql"
1372
-
1373
- def test_dialect(self):
1374
- assert isinstance(self.db_hook.dialect, PostgresDialect)