From 267f72b098f09b720f393e4cd1cf8c949e87ab49 Mon Sep 17 00:00:00 2001
From: Bhargav Dodla <bdodla@expediagroup.com>
Date: Fri, 10 Jan 2025 16:19:46 +0530
Subject: [PATCH] fix: Ensure proper shutdown of the cluster in
 CassandraOnlineStore

---
 .../contrib/cassandra_online_store/cassandra_online_store.py     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py b/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py
index 0128817e97..cd97ace0c2 100644
--- a/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py
+++ b/sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py
@@ -395,6 +395,7 @@ def online_write_batch(
                 print(f"Error writing a batch: {exc}")
                 raise Exception("Error writing a batch") from exc
         session.shutdown()
+        self._cluster.shutdown()
         # correction for the last missing call to `progress`:
         if progress:
             progress(1)