Use the correct pgvector register_vector for psycopg3 connections (#69443)

[shashbha14]

closes: #69443
'PgVectorIngestOperator' always called 'pgvector.psycopg2.register_vector', so it broke
when 'PostgresHook' returned a psycopg3 connection. Now it checks the hook's 'USE_PSYCOPG3'
flag and picks the matching helper ('pgvector.psycopg' for psycopg3, 'pgvector.psycopg2'
otherwise), importing only the driver in use — same pattern as 'postgres_to_gcs.py'.

Added tests for both driver paths.

---

Important

🛠️ Maintainer triage note for @shashbha14 · by @potiuk · 2026-08-13 12:55 UTC

Helpful heads-up from the maintainers — please address before this PR can be reviewed:

• ❌ Merge conflicts. See docs.

Full list of what we check: Pull Request quality criteria.

The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.

_{Automated triage — may be imperfect; a maintainer takes the next look.}

[Dev-iL]

This bug is also handled in #69526. According to my earlier attempts, it shouldn't be possible to solve this in isolation from the the postgres provider. Can you please compare your solution to my latest and see which is better?

[Dev-iL]

Business Logic Comparison: PR #69546 vs PR #69526 (pgvector-only)

Both PRs fix the same issue in PgVectorIngestOperator._register_vector(), but with different error handling:

PR #69546 (Simple)

ifUSE_PSYCOPG3:
frompgvector.psycopgimportregister_vectorelse:
frompgvector.psycopg2importregister_vectorregister_vector(conn)

• Pro: Clean, minimal
• Con: Will fail silently with bare ImportError if pgvector or its driver variant isn't installed

---

PR #69526 (Defensive)

ifUSE_PSYCOPG3:
try:
frompgvector.psycopgimportregister_vectorexcept (ImportError, ModuleNotFoundError) aserr:
raiseAirflowOptionalProviderFeatureException(
"pgvector's psycopg (v3) integration is not installed. Please install it with ""`pip install pgvector`."
) fromerrelse:
try:
frompgvector.psycopg2importregister_vectorexcept (ImportError, ModuleNotFoundError) aserr:
raiseAirflowOptionalProviderFeatureException(
"psycopg2 is not installed. Please install it with ""`pip install apache-airflow-providers-postgres[psycopg2]`."
) fromerrregister_vector(conn)

• Pro: Clear, actionable error messages with installation instructions
• Con: More verbose; assumes pgvector/psycopg2 are optional

---

Test Comparison: PR #69546 vs PR #69526

PR #69546 Tests

@patch("airflow.providers.pgvector.operators.pgvector.USE_PSYCOPG3", False)@patch("airflow.providers.pgvector.operators.pgvector.PgVectorIngestOperator.get_db_hook")deftest_register_vector_psycopg2(mock_get_db_hook, pg_vector_ingest_operator):
mock_db_hook=Mock()
mock_get_db_hook.return_value=mock_db_hookmock_register_vector=MagicMock()
withpatch.dict("sys.modules", {"pgvector.psycopg2": MagicMock(register_vector=mock_register_vector)}):
pg_vector_ingest_operator._register_vector()
mock_register_vector.assert_called_once_with(mock_db_hook.get_conn())
@patch("airflow.providers.pgvector.operators.pgvector.USE_PSYCOPG3", True)@patch("airflow.providers.pgvector.operators.pgvector.PgVectorIngestOperator.get_db_hook")deftest_register_vector_psycopg3(mock_get_db_hook, pg_vector_ingest_operator):
mock_db_hook=Mock()
mock_get_db_hook.return_value=mock_db_hookmock_register_vector=MagicMock()
withpatch.dict("sys.modules", {"pgvector.psycopg": MagicMock(register_vector=mock_register_vector)}):
pg_vector_ingest_operator._register_vector()
mock_register_vector.assert_called_once_with(mock_db_hook.get_conn())
@patch("airflow.providers.pgvector.operators.pgvector.PgVectorIngestOperator._register_vector")@patch("airflow.providers.pgvector.operators.pgvector.SQLExecuteQueryOperator.execute")deftest_execute(mock_execute_query_operator_execute, mock_register_vector, pg_vector_ingest_operator):
pg_vector_ingest_operator.execute(None)
mock_register_vector.assert_called_once()
mock_execute_query_operator_execute.assert_called_once()

Coverage: ✅ Tests both driver paths (psycopg2 + psycopg3)
Approach: Patches sys.modules to avoid requiring either driver
Gap: ❌ No error handling tests—doesn't verify behavior when imports fail

---

PR #69526 Tests (Same file, but more comprehensive)

Everything from PR #69546plus:

@patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", False)deftest_register_vector_raises_clear_error_without_psycopg2(monkeypatch, pg_vector_ingest_operator):
monkeypatch.setitem(sys.modules, "pgvector.psycopg2", None)
withpytest.raises(AirflowOptionalProviderFeatureException, match="psycopg2 is not installed"):
pg_vector_ingest_operator._register_vector()
@patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", True)deftest_register_vector_raises_clear_error_without_psycopg3(monkeypatch, pg_vector_ingest_operator):
monkeypatch.setitem(sys.modules, "pgvector.psycopg", None)
withpytest.raises(AirflowOptionalProviderFeatureException, match=r"psycopg \(v3\) integration"):
pg_vector_ingest_operator._register_vector()
deftest_pgvector_module_imports_without_psycopg2(monkeypatch):
"""The module must import cleanly even when psycopg2/pgvector.psycopg2 isn't installed."""monkeypatch.setitem(sys.modules, "psycopg2", None)
monkeypatch.setitem(sys.modules, "pgvector.psycopg2", None)
module_name="airflow.providers.pgvector.operators.pgvector"monkeypatch.delitem(sys.modules, module_name, raising=False)
try:
module=importlib.import_module(module_name)
assertmodule.PgVectorIngestOperatorisnotNonefinally:
monkeypatch.delitem(sys.modules, module_name, raising=False)

Coverage: ✅ Tests both driver paths + error scenarios
Approach: Same sys.modules patching + explicit error cases
Plus: ✅ Tests module import robustness when dependencies missing

---

Verdict: Test Quality

Aspect	PR #69546	PR #69526
Happy path (psycopg2)	✅	✅
Happy path (psycopg3)	✅	✅
Error: missing psycopg2	❌	✅
Error: missing psycopg3	❌	✅
Module imports cleanly	❌	✅
Total test cases	3	6

PR #69526 is significantly better tested. It validates not just the happy path, but also:

• Clear error messages when dependencies are missing
• The module itself can be imported safely without pgvector installed

For PR #69546 to be production-ready, it should add the three missing test cases from PR #69526.