From 19c1ec381876da3393e3e3e0592e33614a316380 Mon Sep 17 00:00:00 2001 From: TomSteenbergen Date: Mon, 10 Feb 2025 14:18:18 +0100 Subject: [PATCH 1/3] Enable autocommit for read-only operations Signed-off-by: TomSteenbergen --- .../postgres_online_store/postgres.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py index f43247a5457..f679f87812a 100644 --- a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py +++ b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py @@ -5,6 +5,7 @@ from typing import ( Any, AsyncGenerator, + Bool, Callable, Dict, Generator, @@ -58,7 +59,7 @@ class PostgreSQLOnlineStore(OnlineStore): _conn_pool_async: Optional[AsyncConnectionPool] = None @contextlib.contextmanager - def _get_conn(self, config: RepoConfig) -> Generator[Connection, Any, Any]: + def _get_conn(self, config: RepoConfig, autocommit: Bool = False) -> Generator[Connection, Any, Any]: assert config.online_store.type == "postgres" if config.online_store.conn_type == ConnectionType.pool: @@ -66,16 +67,18 @@ def _get_conn(self, config: RepoConfig) -> Generator[Connection, Any, Any]: self._conn_pool = _get_connection_pool(config.online_store) self._conn_pool.open() connection = self._conn_pool.getconn() + connection.set_autocommit(autocommit) yield connection self._conn_pool.putconn(connection) else: if not self._conn: self._conn = _get_conn(config.online_store) + self._conn .set_autocommit(autocommit) yield self._conn @contextlib.asynccontextmanager async def _get_conn_async( - self, config: RepoConfig + self, config: RepoConfig, autocommit: Bool = False ) -> AsyncGenerator[AsyncConnection, Any]: if config.online_store.conn_type == ConnectionType.pool: if not self._conn_pool_async: @@ -84,11 +87,13 @@ async def _get_conn_async( ) await self._conn_pool_async.open() connection = await self._conn_pool_async.getconn() + await connection.set_autocommit(autocommit) 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) + await self._conn_async.set_autocommit(autocommit) yield self._conn_async def online_write_batch( @@ -161,7 +166,7 @@ def online_read( config, table, keys, requested_features ) - with self._get_conn(config) as conn, conn.cursor() as cur: + with self._get_conn(config, autocommit=True) as conn, conn.cursor() as cur: cur.execute(query, params) rows = cur.fetchall() @@ -179,7 +184,7 @@ async def online_read_async( config, table, keys, requested_features ) - async with self._get_conn_async(config) as conn: + async with self._get_conn_async(config, autocommit=True) as conn: async with conn.cursor() as cur: await cur.execute(query, params) rows = await cur.fetchall() @@ -398,7 +403,7 @@ def retrieve_online_documents( Optional[ValueProto], ] ] = [] - with self._get_conn(config) as conn, conn.cursor() as cur: + with self._get_conn(config, autocommit=True) as conn, conn.cursor() as cur: table_name = _table_id(project, table) # Search query template to find the top k items that are closest to the given embedding From 128480dc6f97d2da7a3c2c923564c043f67abea1 Mon Sep 17 00:00:00 2001 From: TomSteenbergen Date: Mon, 10 Feb 2025 14:26:49 +0100 Subject: [PATCH 2/3] Run formatter and linter Signed-off-by: TomSteenbergen --- .../online_stores/postgres_online_store/postgres.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py index f679f87812a..8f1a244922d 100644 --- a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py +++ b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py @@ -5,7 +5,6 @@ from typing import ( Any, AsyncGenerator, - Bool, Callable, Dict, Generator, @@ -59,7 +58,9 @@ class PostgreSQLOnlineStore(OnlineStore): _conn_pool_async: Optional[AsyncConnectionPool] = None @contextlib.contextmanager - def _get_conn(self, config: RepoConfig, autocommit: Bool = False) -> Generator[Connection, Any, Any]: + def _get_conn( + self, config: RepoConfig, autocommit: bool = False + ) -> Generator[Connection, Any, Any]: assert config.online_store.type == "postgres" if config.online_store.conn_type == ConnectionType.pool: @@ -73,12 +74,12 @@ def _get_conn(self, config: RepoConfig, autocommit: Bool = False) -> Generator[C else: if not self._conn: self._conn = _get_conn(config.online_store) - self._conn .set_autocommit(autocommit) + self._conn.set_autocommit(autocommit) yield self._conn @contextlib.asynccontextmanager async def _get_conn_async( - self, config: RepoConfig, autocommit: Bool = False + self, config: RepoConfig, autocommit: bool = False ) -> AsyncGenerator[AsyncConnection, Any]: if config.online_store.conn_type == ConnectionType.pool: if not self._conn_pool_async: From d815bd7b881a9d749fd35c5d1aa7760816c7acf5 Mon Sep 17 00:00:00 2001 From: TomSteenbergen Date: Mon, 10 Feb 2025 14:31:12 +0100 Subject: [PATCH 3/3] Add missing commit Signed-off-by: TomSteenbergen --- .../feast/infra/online_stores/postgres_online_store/postgres.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py index 8f1a244922d..4f519003d61 100644 --- a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py +++ b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py @@ -345,6 +345,7 @@ def teardown( for table in tables: table_name = _table_id(project, table) cur.execute(_drop_table_and_index(table_name)) + conn.commit() except Exception: logging.exception("Teardown failed") raise