From f2ae6e5f4686de2f7deb43988aa82c52049a0f9d Mon Sep 17 00:00:00 2001 From: Aleksey Veresov Date: Fri, 18 Oct 2024 16:59:48 +0200 Subject: [PATCH] Add finalizer to Connection --- python/hopsworks_common/connection.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/hopsworks_common/connection.py b/python/hopsworks_common/connection.py index 5a4903a0c..6972b6a85 100644 --- a/python/hopsworks_common/connection.py +++ b/python/hopsworks_common/connection.py @@ -21,6 +21,7 @@ import re import sys import warnings +import weakref from typing import Any, Optional from hopsworks_common import client, usage, util, version @@ -351,6 +352,7 @@ def connect(self) -> None: """ client.stop() self._connected = True + finalizer = weakref.finalize(self, self.close) try: # determine engine, needed to init client if (self._engine is not None and self._engine.lower() == "spark") or ( @@ -413,6 +415,7 @@ def connect(self) -> None: self._provide_project() except (TypeError, ConnectionError): self._connected = False + finalizer.detach() raise self._check_compatibility() @@ -446,7 +449,7 @@ def close(self) -> None: This will clean up any materialized certificates on the local file system of external environments such as AWS SageMaker. - Usage is recommended but optional. + Usage is optional. !!! example ```python @@ -455,6 +458,9 @@ def close(self) -> None: conn.close() ``` """ + if not self._connected: + return # the connection is already closed + from hsfs import engine if OpenSearchClientSingleton._instance: