Unable to create multiple lsh indices each one in its own keyspace #171…

… (#172) Co-authored-by: Ron Assa <ron@analytika.io>
ekzhu · Feb 4, 2022 · 3d90696 · 3d90696
1 parent 14a116f
commit 3d90696
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/datasketch/storage.py b/datasketch/storage.py
@@ -265,22 +265,24 @@ class CassandraSharedSession(object):
         @classmethod
         def get_session(cls, seeds, **kwargs):
             _ = kwargs
+            keyspace = kwargs["keyspace"]
+            replication = kwargs["replication"]
+
             if cls.__session is None:
                 # Allow dependency injection
                 session = kwargs.get("session")
                 if session is None:
                     cluster = c_cluster.Cluster(seeds)
                     session = cluster.connect()
-                keyspace = kwargs["keyspace"]
-                replication = kwargs["replication"]
+                cls.__session = session
+            if cls.__session.keyspace != keyspace:
                 if kwargs.get("drop_keyspace", False):
-                    session.execute(cls.QUERY_DROP_KEYSPACE.format(keyspace))
-                session.execute(cls.QUERY_CREATE_KEYSPACE.format(
+                    cls.__session.execute(cls.QUERY_DROP_KEYSPACE.format(keyspace))
+                cls.__session.execute(cls.QUERY_CREATE_KEYSPACE.format(
                     keyspace=keyspace,
                     replication=str(replication),
                 ))
-                session.set_keyspace(keyspace)
-                cls.__session = session
+                cls.__session.set_keyspace(keyspace)
             return cls.__session
 
         @classmethod