From 5a2fcfed7d4a55b14608f26cbe41091e9ba37b91 Mon Sep 17 00:00:00 2001 From: doria <93405247+dr-frmr@users.noreply.github.com> Date: Thu, 15 May 2025 10:45:14 -0400 Subject: [PATCH] fix: ensure pgvector installed before running alembic migrations (#109) * fix: ensure pgvector installed before running alembic migrations * fix: include extensions in search path --- migrations/env.py | 6 +++++- src/db.py | 11 ++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/migrations/env.py b/migrations/env.py index e8ad8ae0..234ce529 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -107,8 +107,12 @@ def run_migrations_online() -> None: connection.execute( text(f"GRANT ALL ON SCHEMA {target_metadata.schema} TO current_user") ) + # Install pgvector extension if it doesn't exist + connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) # Set and verify search_path - connection.execute(text(f"SET search_path TO {target_metadata.schema}, public")) + connection.execute( + text(f"SET search_path TO {target_metadata.schema}, public, extensions") + ) connection.commit() context.configure( diff --git a/src/db.py b/src/db.py index 4960cc0a..ee6d544e 100644 --- a/src/db.py +++ b/src/db.py @@ -50,11 +50,12 @@ def init_db(): echo=os.getenv("SQL_DEBUG", "false").lower() == "true", ) - # Create schema if it doesn't exist - if table_schema: - with sync_engine.connect() as connection: - connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"')) - connection.commit() + with sync_engine.connect() as connection: + # Create schema if it doesn't exist + connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"')) + # Install pgvector extension if it doesn't exist + connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + connection.commit() # Run Alembic migrations alembic_cfg = Config("alembic.ini")