From 4b8a175ba40ef5d445a7eff4b7399584709438e0 Mon Sep 17 00:00:00 2001 From: Dev-iL <6509619+Dev-iL@users.noreply.github.com> Date: Mon, 6 Jul 2026 11:01:26 +0300 Subject: [PATCH] Guard pgvector operator's psycopg2-specific register_vector import from pgvector.psycopg2 import register_vector ran unconditionally at module level, forcing psycopg2 to be importable regardless of which driver the underlying PostgresHook connection actually uses. Make the import lazy and raise a clear AirflowOptionalProviderFeatureException if it's genuinely needed but missing. This does not fix the separate, pre-existing bug where register_vector is always called with the psycopg2-specific helper regardless of whether the connection is actually psycopg2 or psycopg3; that is tracked at https://github.com/apache/airflow/issues/69443. The finally block was reimporting the module while the fake psycopg2 absence was still monkeypatched, leaving it cached in a degraded state for the rest of the test session. Also documents that this file's two register_vector-mocking tests depend on psycopg2 staying installed in this provider's dev environment, which is not guaranteed once apache-airflow-providers-postgres stops hard-requiring it (a separate, independently-branched fix for the same migration) -- requesting that provider's [psycopg2] extra here isn't possible yet since the extra doesn't exist until that fix lands. Part of the migration tracked in apache/airflow#68453. --- .../providers/pgvector/operators/pgvector.py | 13 ++++++-- .../unit/pgvector/operators/test_pgvector.py | 31 +++++++++++++++++-- 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py b/providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py index ae48a377ea7a7..6911d1ae69b9b 100644 --- a/providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py +++ b/providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py @@ -17,8 +17,7 @@ # under the License. from __future__ import annotations -from pgvector.psycopg2 import register_vector - +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator @@ -41,6 +40,16 @@ def __init__(self, *args, **kwargs) -> None: def _register_vector(self) -> None: """Register the vector type with your connection.""" + # This always uses the psycopg2-specific registration helper, regardless of whether the + # underlying connection is actually psycopg2 or psycopg3; tracked separately at + # https://github.com/apache/airflow/issues/69443 + 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 conn = self.get_db_hook().get_conn() register_vector(conn) diff --git a/providers/pgvector/tests/unit/pgvector/operators/test_pgvector.py b/providers/pgvector/tests/unit/pgvector/operators/test_pgvector.py index 6fe144d5bee99..33a6f95bc77c2 100644 --- a/providers/pgvector/tests/unit/pgvector/operators/test_pgvector.py +++ b/providers/pgvector/tests/unit/pgvector/operators/test_pgvector.py @@ -16,10 +16,13 @@ # under the License. from __future__ import annotations +import importlib +import sys from unittest.mock import Mock, patch import pytest +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.pgvector.operators.pgvector import PgVectorIngestOperator @@ -32,7 +35,12 @@ def pg_vector_ingest_operator(): ) -@patch("airflow.providers.pgvector.operators.pgvector.register_vector") +# These two tests patch pgvector.psycopg2.register_vector, which requires psycopg2 to be +# installed. Once apache-airflow-providers-postgres stops hard-requiring psycopg2 (a separate, +# independently-branched fix for the same migration, apache/airflow#68453), this provider's own +# dev environment could lose psycopg2 unless its dependency on that provider requests the +# [psycopg2] extra explicitly -- currently not declared because that extra doesn't exist yet. +@patch("pgvector.psycopg2.register_vector") @patch("airflow.providers.pgvector.operators.pgvector.PgVectorIngestOperator.get_db_hook") def test_register_vector(mock_get_db_hook, mock_register_vector, pg_vector_ingest_operator): # Create a mock database connection @@ -43,7 +51,7 @@ def test_register_vector(mock_get_db_hook, mock_register_vector, pg_vector_inges mock_register_vector.assert_called_with(mock_db_hook.get_conn()) -@patch("airflow.providers.pgvector.operators.pgvector.register_vector") +@patch("pgvector.psycopg2.register_vector") @patch("airflow.providers.pgvector.operators.pgvector.SQLExecuteQueryOperator.execute") @patch("airflow.providers.pgvector.operators.pgvector.PgVectorIngestOperator.get_db_hook") def test_execute( @@ -54,3 +62,22 @@ def test_execute( pg_vector_ingest_operator.execute(None) mock_execute_query_operator_execute.assert_called_once() + + +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() + + +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)