fix: Update PostgresDatastore type as PgClient provideers.

GoogleCloudPlatform · Sep 20, 2024 · 5d1bced · 5d1bced
1 parent 2907d4a
commit 5d1bced
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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):