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

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 freeto 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

@Dev-iL