mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:48:03 +00:00
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import sqlalchemy
|
|
from llama_index.storage.kvstore.postgres import ( # ty:ignore[unresolved-import]
|
|
PostgresKVStore,
|
|
)
|
|
|
|
|
|
class PatchedPostgresKVStore(PostgresKVStore):
|
|
"""Patched PostgresKVStore that escapes the table name.
|
|
|
|
Our tables names contain "-" (they are uuids),
|
|
which are not supported by the default PostgresKVStore.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
connection_string: str,
|
|
async_connection_string: str,
|
|
table_name: str,
|
|
schema_name: str = "public",
|
|
engine: sqlalchemy.engine.Engine | None = None,
|
|
async_engine: sqlalchemy.ext.asyncio.AsyncEngine # ty:ignore[possibly-missing-submodule]
|
|
| None = None,
|
|
perform_setup: bool = True,
|
|
debug: bool = False,
|
|
use_jsonb: bool = False,
|
|
) -> None:
|
|
super().__init__(
|
|
connection_string=connection_string,
|
|
async_connection_string=async_connection_string,
|
|
table_name=table_name,
|
|
schema_name=schema_name,
|
|
engine=engine,
|
|
async_engine=async_engine,
|
|
perform_setup=perform_setup,
|
|
debug=debug,
|
|
use_jsonb=use_jsonb,
|
|
)
|
|
self._table_class.__tablename__ = f'"{self._table_class.__tablename__}"'
|