Repository: run-llama/llama_index
Pull request: Create PostgresKVStore from existing SQLAlchemy Engine
# Description

Add support for the creation of `PostgresKVStore` instances using existing SQLAlchemy `Engine` and/or `AsyncEngine` objects. This enables developers to reuse preconfigured connection pools, apply custom engine settings, and integrate more seamlessly with applications that already manage their own database connections.

`PostgresKVStore` objects can now be instantiated as follows:

```python
engine = create_engine(...)
async_engine = create_async_engine(...)

kv_store = PostgresKVStore(engine=engine, async_engine=async_engine, table_name="kv_store")
```

Fixes #18797 

## New Package?

Not a new package.

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [x] Yes
- [ ] No

## Type of Change

New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

Your pull-request will likely not be merged unless it is covered by some form of impactful unit testing.

- [ ] I added new unit tests to cover this change
- [x] I believe this change is already covered by existing unit tests

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `uv run make format; uv run make lint` to appease the lint gods


Problem statement:
[Feature Request]: Create PostgresKVStore and PostgresDocumentStore from existing SQLAlchemy Engine
Hello,

### Feature Description

It would be highly beneficial to support the creation of `PostgresKVStore` and `PostgresDocumentStore` instances using existing SQLAlchemy `Engine` and/or `AsyncEngine` objects. This would enable developers to reuse preconfigured connection pools, apply custom engine settings, and integrate more seamlessly with applications that already manage their own database connections.

### Proposed Usage:

```python
engine = create_engine(...)
async_engine = create_async_engine(...)

kv_store = PostgresKVStore(engine=engine, async_engine=async_engine, table_name="kv_store")
doc_store = PostgresDocumentStore(postgres_kvstore=kv_store)
```

### Reason

Currently, both `PostgresKVStore` and its wrapper `PostgresDocumentStore` only support two factory methods — `from_uri(...)` and `from_params(...)` — for instantiation. 

This design forces every consumer to recreate their own connection pools and gives developers no way to inject or share an existing SQLAlchemy `Engine` or `AsyncEngine`. 

As a result, you lose out on centralized pool management, custom connection settings, and any middleware or event hooks you’ve already configured on your engines.

### Value of Feature

Allowing `PostgresKVStore` and `PostgresDocumentStore` to be created from existing SQLAlchemy `Engine` and/or `AsyncEngine` instances provides greater flexibility and control for developers. 

It would enable efficient reuse of connection pools, seamless integration with existing database infrastructure, and better alignment with application-level configurations.

Hints from issue discussion:
Hey there @Florian-BACHO, if you already have some implementation you could open a PR with, feel free to do it! Otherwise, I'm happy to look into this :)

Commit message: Update postgres docstore dependencies

Patch to review:
diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/pyproject.toml b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/pyproject.toml
index 0c0f0ef3d3..36eeb193f4 100644
--- a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/pyproject.toml
+++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/pyproject.toml
@@ -26,14 +26,14 @@ dev = [
 
 [project]
 name = "llama-index-storage-docstore-postgres"
-version = "0.3.0"
+version = "0.4.0"
 description = "llama-index docstore postgres integration"
 authors = [{name = "Your Name", email = "you@example.com"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
 dependencies = [
-    "llama-index-storage-kvstore-postgres>=0.3.0,<0.4",
+    "llama-index-storage-kvstore-postgres>=0.4.0,<0.5",
     "llama-index-core>=0.12.0,<0.13",
 ]
 
diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/uv.lock b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/uv.lock
index a59192284d..9d65f26e0d 100644
--- a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/uv.lock
+++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-postgres/uv.lock
@@ -1609,7 +1609,7 @@ wheels = [
 
 [[package]]
 name = "llama-index-storage-docstore-postgres"
-version = "0.3.0"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "llama-index-core" },
diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/llama_index/storage/kvstore/postgres/base.py b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/llama_index/storage/kvstore/postgres/base.py
index 85066e0872..b73bc2583a 100644
--- a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/llama_index/storage/kvstore/postgres/base.py
+++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/llama_index/storage/kvstore/postgres/base.py
@@ -2,6 +2,14 @@
 from typing import Any, Dict, List, Optional, Tuple, Type
 from urllib.parse import urlparse
 
+try:
+    import sqlalchemy
+    import sqlalchemy.ext.asyncio  # noqa
+except ImportError:
+    raise ImportError(
+        "`sqlalchemy[asyncio]` package should be pre installed"
+    )
+
 from llama_index.core.storage.kvstore.types import (
     DEFAULT_BATCH_SIZE,
     DEFAULT_COLLECTION,
@@ -80,6 +88,8 @@ def __init__(
         async_connection_string: str,
         table_name: str,
         schema_name: str = "public",
+        engine: Optional[sqlalchemy.engine.Engine] = None,
+        async_engine: Optional[sqlalchemy.ext.asyncio.AsyncEngine] = None,
         perform_setup: bool = True,
         debug: bool = False,
         use_jsonb: bool = False,
@@ -87,11 +97,9 @@ def __init__(
         try:
             import asyncpg  # noqa
             import psycopg2  # noqa
-            import sqlalchemy
-            import sqlalchemy.ext.asyncio  # noqa
         except ImportError:
             raise ImportError(
-                "`sqlalchemy[asyncio]`, `psycopg2-binary` and `asyncpg` "
+                "`psycopg2-binary` and `asyncpg` "
                 "packages should be pre installed"
             )
 
@@ -104,6 +112,8 @@ def __init__(
         self.perform_setup = perform_setup
         self.debug = debug
         self.use_jsonb = use_jsonb
+        self._engine = engine
+        self._async_engine = async_engine
         self._is_initialized = False
 
         from sqlalchemy.orm import declarative_base
@@ -177,10 +187,10 @@ def _connect(self) -> Any:
         from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
         from sqlalchemy.orm import sessionmaker
 
-        self._engine = create_engine(self.connection_string, echo=self.debug)
+        self._engine = self._engine or create_engine(self.connection_string, echo=self.debug)
         self._session = sessionmaker(self._engine)
 
-        self._async_engine = create_async_engine(self.async_connection_string)
+        self._async_engine = self._async_engine or create_async_engine(self.async_connection_string)
         self._async_session = sessionmaker(self._async_engine, class_=AsyncSession)
 
     def _create_schema_if_not_exists(self) -> None:
diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml
index 84e8b31ab2..b52529f440 100644
--- a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml
+++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml
@@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-storage-kvstore-postgres"
-version = "0.3.0"
+version = "0.4.0"
 description = "llama-index kvstore postgres integration"
 authors = [{name = "Your Name", email = "you@example.com"}]
 requires-python = ">=3.9,<4.0"
diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/uv.lock b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/uv.lock
index 501d93736c..54f8e7639e 100644
--- a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/uv.lock
+++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/uv.lock
@@ -1623,7 +1623,7 @@ wheels = [
 
 [[package]]
 name = "llama-index-storage-kvstore-postgres"
-version = "0.3.0"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "llama-index-core" },


Code review comments:
diff_hunk: @@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-storage-kvstore-postgres"
-version = "0.3.0"
+version = "0.4.0"
path: llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml
text: @user1:
Hey, just one thing: since the changes are not particularly big, would it be ok if we bumped this to 0.3.1 instead of 0.4.0? In this way, you wouldn't have to change the upper limit for the docstore integration (you could leave it to <0.4 instead of bumping it to <0.5)

@author:
Sure !

@author:
I updated llama-index-storage-kvstore-postgres's version to 0.3.1 and reverted llama-index-storage-docstore-postgres to 0.3.0