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.

[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)

if USE_PSYCOPG3:
    from pgvector.psycopg import register_vector
else:
    from pgvector.psycopg2 import register_vector
register_vector(conn)

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

---

PR #69526 (Defensive)

if USE_PSYCOPG3:
    try:
        from pgvector.psycopg import register_vector
    except (ImportError, ModuleNotFoundError) as err:
        raise AirflowOptionalProviderFeatureException(
            "pgvector's psycopg (v3) integration is not installed. Please install it with "
            "`pip install pgvector`."
        ) from err
else:
    try:
        from pgvector.psycopg2 import register_vector
    except (ImportError, ModuleNotFoundError) as err:
        raise AirflowOptionalProviderFeatureException(
            "psycopg2 is not installed. Please install it with "
            "`pip install apache-airflow-providers-postgres[psycopg2]`."
        ) from err
register_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")
def test_register_vector_psycopg2(mock_get_db_hook, pg_vector_ingest_operator):
    mock_db_hook = Mock()
    mock_get_db_hook.return_value = mock_db_hook
    mock_register_vector = MagicMock()

    with patch.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")
def test_register_vector_psycopg3(mock_get_db_hook, pg_vector_ingest_operator):
    mock_db_hook = Mock()
    mock_get_db_hook.return_value = mock_db_hook
    mock_register_vector = MagicMock()

    with patch.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")
def test_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 #69546 plus:

@patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", False)
def test_register_vector_raises_clear_error_without_psycopg2(monkeypatch, pg_vector_ingest_operator):
    monkeypatch.setitem(sys.modules, "pgvector.psycopg2", None)
    with pytest.raises(AirflowOptionalProviderFeatureException, match="psycopg2 is not installed"):
        pg_vector_ingest_operator._register_vector()


@patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", True)
def test_register_vector_raises_clear_error_without_psycopg3(monkeypatch, pg_vector_ingest_operator):
    monkeypatch.setitem(sys.modules, "pgvector.psycopg", None)
    with pytest.raises(AirflowOptionalProviderFeatureException, match=r"psycopg \(v3\) integration"):
        pg_vector_ingest_operator._register_vector()


def test_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)
        assert module.PgVectorIngestOperator is not None
    finally:
        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.