fix: Implement shutdown method for CassandraOnlineStore and ensure pr…

…oper shutdown in SparkKafkaProcessor
ExpediaGroup · Jan 16, 2025 · c9b2c6d · c9b2c6d
1 parent 95bc2a6
commit c9b2c6d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/sdk/python/feast/infra/contrib/spark_kafka_processor.py b/sdk/python/feast/infra/contrib/spark_kafka_processor.py
@@ -313,6 +313,8 @@ def batch_write_pandas_df(iterator, spark_serialized_artifacts, join_keys):
                     lambda x: None,
                 )
 
+                online_store.shutdown()
+
             yield pd.DataFrame([pd.Series(range(1, 2))])  # dummy result
 
         def batch_write(

diff --git a/...python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py b/...python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py
@@ -209,6 +209,15 @@ class CassandraOnlineStore(OnlineStore):
     _keyspace: str = "feast_keyspace"
     _prepared_statements: Dict[str, PreparedStatement] = {}
 
+    def shutdown(self):
+        """
+        Shutdown the Cassandra cluster and session.
+        """
+        if not self._session.is_shutdown:
+            self._session.shutdown()
+        if not self._cluster.is_shutdown:
+            self._cluster.shutdown()
+
     def _get_cluster(self, config: RepoConfig):
         """
         Establish the database connection, if not yet created,