Skip to content

pgvector: Guard the psycopg2-specific register_vector import#69491

Closed
Dev-iL wants to merge 1 commit into
apache:mainfrom
Dev-iL:psycopg3-sync-pgvector
Closed

pgvector: Guard the psycopg2-specific register_vector import#69491
Dev-iL wants to merge 1 commit into
apache:mainfrom
Dev-iL:psycopg3-sync-pgvector

Conversation

@Dev-iL

@Dev-iL Dev-iL commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator

related:

depends on:

from pgvector.psycopg2 import register_vector ran unconditionally at module level in PgVectorIngestOperator, forcing psycopg2 to be importable regardless of which driver the underlying PostgresHook connection actually uses. Make the import lazy and raise a clear, actionable error if it's genuinely needed but missing.

Note: this module only imports cleanly without psycopg2 installed once apache-airflow-providers-postgres itself stops hard-requiring psycopg2 — that's the companion change in psycopg3-sync-postgres. This PR should be merged after (or at least released alongside) that one; see the code comment at the import site and in the test file.

This does not fix a 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's tracked separately at #69443, referenced in a code comment at the guard site so it isn't lost.

What changed

  • providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py: from pgvector.psycopg2 import register_vector moved inside _register_vector(), guarded with AirflowOptionalProviderFeatureException if pgvector.psycopg2/psycopg2 is unavailable.

Test plan

  • New tests: test_register_vector_raises_clear_error_without_psycopg2, test_pgvector_module_imports_without_psycopg2 (simulates psycopg2 absent via sys.modules patching, confirms clean import and a clear error on actual use).
  • Existing register_vector-mocking tests still pass; a code comment documents that they currently rely on psycopg2 staying installed in this provider's dev environment until the postgres provider's [psycopg2] extra exists to depend on explicitly.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: Claude Code (Sonnet 5) following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

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 apache#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#68453.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant