Skip to content

pgvector operator's register_vector always assumes a psycopg2 connection #69443

Description

@Dev-iL

PgVectorIngestOperator._register_vector (in providers/pgvector/src/airflow/providers/pgvector/operators/pgvector.py) always calls the psycopg2-specific pgvector.psycopg2.register_vector(conn), regardless of whether the underlying PostgresHook connection is actually using psycopg2 or psycopg3.

Why this is wrong: PostgresHook (in apache-airflow-providers-postgres) picks psycopg2 or psycopg3 based on USE_PSYCOPG3 (SQLAlchemy version + psycopg package availability). pgvector.psycopg2.register_vector expects a psycopg2 connection object; calling it with a psycopg3 connection is likely incorrect (psycopg3 has its own pgvector.psycopg registration helper with a different API).

Context: Surfaced while auditing all psycopg2-specific code paths across the codebase as part of #68453 (migrating the sync Postgres driver default from psycopg2 to psycopg3). That work only added an import guard around the top-level from pgvector.psycopg2 import register_vector so the module doesn't hard-require psycopg2 at import time — it does not fix this connection-type mismatch, which is a separate, pre-existing bug.

Follow-up work needed:

  • Detect whether the hook's connection is psycopg2 or psycopg3 (e.g. via PostgresHook's USE_PSYCOPG3).
  • Call the matching registration helper: pgvector.psycopg2.register_vector for psycopg2, pgvector.psycopg.register_vector for psycopg3.
  • Add test coverage for both paths.

Acceptance criteria: PgVectorIngestOperator correctly registers the pgvector type for both a psycopg2-backed and a psycopg3-backed PostgresHook connection, with tests covering both.

Metadata

Metadata

Assignees

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions