Add finalizer to Connection

Author: aversey

python/hopsworks/connection.py

@@ -16,14 +16,15 @@
 
 import os
 import re
-import warnings
 import sys
+import warnings
+import weakref
 
-from requests.exceptions import ConnectionError
-
-from hopsworks.decorators import connected, not_connected
 from hopsworks import client, version
 from hopsworks.core import project_api, secret_api, variable_api
+from hopsworks.decorators import connected, not_connected
+from requests.exceptions import ConnectionError
+
 
 HOPSWORKS_PORT_DEFAULT = 443
 HOSTNAME_VERIFICATION_DEFAULT = True
@@ -241,6 +242,7 @@ def connect(self):
         """
         client.stop()
         self._connected = True
+        finalizer = weakref.finalize(self, self.close)
         try:
             # init client
             if client.base.Client.REST_ENDPOINT not in os.environ:
@@ -263,6 +265,7 @@ def connect(self):
             self._variable_api = variable_api.VariableApi()
         except (TypeError, ConnectionError):
             self._connected = False
+            finalizer.detach()
             raise
         print(
             "Connected. Call `.close()` to terminate connection gracefully.",
@@ -278,7 +281,7 @@ def close(self):
         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.
         """
         from hsfs import client as hsfs_client
         from hsfs import engine as hsfs_engine
@@ -299,6 +302,9 @@ def close(self):
         except:  # noqa: E722
             pass
 
+        if not self._connected:
+            return  # the connection is already closed
+
         client.stop()
         self._connected = False