Feat: add PostgreSQL extensions for vector and AGE

- Ensure VECTOR extension is available when PostgreSQL init
- Ensure AGE extension is available when PGGraphStorage init
This commit is contained in:
yangdx 2025-07-21 01:46:41 +08:00
parent e6de125eb4
commit 19a38d9310

View file

@ -73,6 +73,10 @@ class PostgreSQLDB:
max_size=self.max,
)
# Ensure VECTOR extension is available
async with self.pool.acquire() as connection:
await self.configure_vector_extension(connection)
logger.info(
f"PostgreSQL, Connected to database at {self.host}:{self.port}/{self.database}"
)
@ -82,6 +86,26 @@ class PostgreSQLDB:
)
raise
@staticmethod
async def configure_vector_extension(connection: asyncpg.Connection) -> None:
"""Create VECTOR extension if it doesn't exist for vector similarity operations."""
try:
await connection.execute("CREATE EXTENSION IF NOT EXISTS vector") # type: ignore
logger.info("VECTOR extension ensured for PostgreSQL")
except Exception as e:
logger.warning(f"Could not create VECTOR extension: {e}")
# Don't raise - let the system continue without vector extension
@staticmethod
async def configure_age_extension(connection: asyncpg.Connection) -> None:
"""Create AGE extension if it doesn't exist for graph operations."""
try:
await connection.execute("CREATE EXTENSION IF NOT EXISTS age") # type: ignore
logger.info("AGE extension ensured for PostgreSQL")
except Exception as e:
logger.warning(f"Could not create AGE extension: {e}")
# Don't raise - let the system continue without AGE extension
@staticmethod
async def configure_age(connection: asyncpg.Connection, graph_name: str) -> None:
"""Set the Apache AGE environment and creates a graph if it does not exist.
@ -1851,6 +1875,11 @@ class PGGraphStorage(BaseGraphStorage):
f"PostgreSQL Graph initialized: workspace='{self.workspace}', graph_name='{self.graph_name}'"
)
# Create AGE extension and configure graph environment once at initialization
async with self.db.pool.acquire() as connection:
# First ensure AGE extension is created
await PostgreSQLDB.configure_age_extension(connection)
# Execute each statement separately and ignore errors
queries = [
f"SELECT create_graph('{self.graph_name}')",