Online & Utils: Use _get_conn_async

Author: job-almekinders

sdk/python/feast/infra/online_stores/contrib/postgres.py

@@ -28,6 +28,7 @@
 from feast.infra.online_stores.online_store import OnlineStore
 from feast.infra.utils.postgres.connection_utils import (
     _get_conn,
+    _get_conn_async,
     _get_connection_pool,
     _get_connection_pool_async,
 )
@@ -57,6 +58,8 @@ class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
 class PostgreSQLOnlineStore(OnlineStore):
     _conn: Optional[Connection] = None
     _conn_pool: Optional[ConnectionPool] = None
+
+    _conn_async: Optional[AsyncConnection] = None
     _conn_pool_async: Optional[AsyncConnectionPool] = None
 
     @contextlib.contextmanager
@@ -79,14 +82,19 @@ def _get_conn(self, config: RepoConfig) -> Generator[Connection, Any, Any]:
     async def _get_conn_async(
         self, config: RepoConfig
     ) -> AsyncGenerator[AsyncConnection, Any]:
-        if not self._conn_pool_async:
-            self._conn_pool_async = await _get_connection_pool_async(
-                config.online_store
-            )
-            await self._conn_pool_async.open()
-        connection = await self._conn_pool_async.getconn()
-        yield connection
-        await self._conn_pool_async.putconn(connection)
+        if config.online_store.conn_type == ConnectionType.pool:
+            if not self._conn_pool_async:
+                self._conn_pool_async = await _get_connection_pool_async(
+                    config.online_store
+                )
+                await self._conn_pool_async.open()
+            connection = await self._conn_pool_async.getconn()
+            yield connection
+            await self._conn_pool_async.putconn(connection)
+        else:
+            if not self._conn_async:
+                self._conn_async = await _get_conn_async(config.online_store)
+            yield self._conn_async
 
     def online_write_batch(
         self,

sdk/python/feast/infra/utils/postgres/connection_utils.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import psycopg
 import pyarrow as pa
-from psycopg.connection import Connection
+from psycopg import AsyncConnection, Connection
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
 
 from feast.infra.utils.postgres.postgres_config import PostgreSQLConfig
@@ -21,6 +21,16 @@ def _get_conn(config: PostgreSQLConfig) -> Connection:
     return conn
 
 
+async def _get_conn_async(config: PostgreSQLConfig) -> AsyncConnection:
+    """Get a psycopg `AsyncConnection`."""
+    conn = await psycopg.AsyncConnection.connect(
+        conninfo=_get_conninfo(config),
+        keepalives_idle=config.keepalives_idle,
+        **_get_conn_kwargs(config),
+    )
+    return conn
+
+
 def _get_connection_pool(config: PostgreSQLConfig) -> ConnectionPool:
     """Get a psycopg `ConnectionPool`."""
     return ConnectionPool(