From 5d1bced596fd023f580ca1ab02d814a124a7dba0 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 21 Sep 2024 02:52:56 +0530 Subject: [PATCH] fix: Update PostgresDatastore type as `PgClient` provideers. --- retrieval_service/datastore/providers/alloydb.py | 4 +++- retrieval_service/datastore/providers/cloudsql_postgres.py | 4 +++- retrieval_service/datastore/providers/postgres.py | 3 ++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/retrieval_service/datastore/providers/alloydb.py b/retrieval_service/datastore/providers/alloydb.py index 5bb4c9b1..9f2f35ec 100644 --- a/retrieval_service/datastore/providers/alloydb.py +++ b/retrieval_service/datastore/providers/alloydb.py @@ -19,6 +19,8 @@ from pgvector.asyncpg import register_vector from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from postgres import Client as PgClient + import models @@ -41,7 +43,7 @@ class Config(BaseModel, datastore.AbstractConfig): class Client(datastore.Client[Config]): __connector: Optional[AsyncConnector] = None - __pg_ds: PostgresDatastore + __pg_ds: PgClient @datastore.classproperty def kind(cls): diff --git a/retrieval_service/datastore/providers/cloudsql_postgres.py b/retrieval_service/datastore/providers/cloudsql_postgres.py index be4bf11e..bcb8361c 100644 --- a/retrieval_service/datastore/providers/cloudsql_postgres.py +++ b/retrieval_service/datastore/providers/cloudsql_postgres.py @@ -20,6 +20,8 @@ from pgvector.asyncpg import register_vector from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from postgres import Client as PgClient + import models @@ -40,7 +42,7 @@ class Config(BaseModel, datastore.AbstractConfig): class Client(datastore.Client[Config]): - __pg_ds: PostgresDatastore + __pg_ds: PgClient __connector: Optional[Connector] = None @datastore.classproperty diff --git a/retrieval_service/datastore/providers/postgres.py b/retrieval_service/datastore/providers/postgres.py index c2eb68b4..d0062d39 100644 --- a/retrieval_service/datastore/providers/postgres.py +++ b/retrieval_service/datastore/providers/postgres.py @@ -19,6 +19,7 @@ from pgvector.asyncpg import register_vector from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from postgres import Client as PgClient import models @@ -38,7 +39,7 @@ class Config(BaseModel, datastore.AbstractConfig): class Client(datastore.Client[Config]): - __pg_ds: PostgresDatastore + __pg_ds: PgClient @datastore.classproperty def kind(cls):