Dev-iL commented on PR #69546:
URL: https://github.com/apache/airflow/pull/69546#issuecomment-4914736277
## 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)
```python
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)
```python
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
---
## Which is Better for PR #69546?
**PR #69526's approach is objectively better** because:
1. **User experience:** A clear `AirflowOptionalProviderFeatureException`
with install instructions beats a cryptic import error
2. **Consistency:** Matches the pattern used elsewhere in PR #69526 (Google,
Amazon, Postgres providers)
3. **Maintenance:** Follows Airflow's conventions for optional features
**However, PR #69546 is acceptable if:**
- The team assumes pgvector is always installed as a hard dependency
- They want the minimal diff for easier review/backporting
- Error handling is delegated to the operator's caller
**Recommendation:** Update PR #69546 to include the error handling from PR
#69526—it's just a few extra lines but significantly improves user experience.
## Test Comparison: PR #69546 vs PR #69526
### **PR #69546 Tests** (Current in linked commit)
```python
@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:**
```python
@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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
