session shutdown null check

datastax · Aug 27, 2024 · a37ffae · a37ffae
1 parent 766bf24
commit a37ffae
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/impl/astra_vector.py b/impl/astra_vector.py
@@ -17,7 +17,7 @@
 
 from cassandra import ConsistencyLevel, Unauthorized, ProtocolVersion
 from cassandra.auth import PlainTextAuthProvider
-from cassandra.cluster import Cluster, DriverException, NoHostAvailable
+from cassandra.cluster import Cluster, DriverException, NoHostAvailable, Session
 from cassandra.policies import RetryPolicy, ExponentialReconnectionPolicy
 from cassandra.query import (
     UNSET_VALUE,
@@ -140,7 +140,7 @@ class CassandraClient:
     def __init__(self, token, dbid=None) -> None:
         self.token = token
         self.dbid = dbid
-        self.session = None  # Initialize session to None
+        self.session: Optional[Session] = None  # Initialize session to None
 
     async def async_setup(self):
         if self.dbid is None:
@@ -149,7 +149,7 @@ async def async_setup(self):
         # Attempt to connect synchronously (assuming connect is a sync method)
         session = self.connect()
         if session:
-            self.session = session
+            self.session: Session = session
             # Perform async table creation
             await self.create_table()
         else:
@@ -322,7 +322,7 @@ def get_astra_bundle_url(self):
 
         return response.json()["downloadURL"]
 
-    def connect(self, retry=False):
+    def connect(self, retry=False) -> Session:
         dbid = self.dbid
         token = self.token
         if dbid is not None:
@@ -1500,7 +1500,8 @@ def upsert_assistant(
 
     def __del__(self):
         # close the connection when the client is destroyed
-        self.session.shutdown()
+        if self.session is not None:
+            self.session.shutdown()
 
     # TODO: make these async
     def selectAllFromTable(self, table):