From cd576a13aa91256970450232a1627159845b4550 Mon Sep 17 00:00:00 2001
From: Gal Shubeli <galshubeli93@gmail.com>
Date: Thu, 14 Aug 2025 15:44:05 +0300
Subject: [PATCH] rev-driver-changes

---
 graphiti_core/driver/driver.py          | 11 ++++++++
 graphiti_core/driver/falkordb_driver.py | 35 +++----------------------
 2 files changed, 14 insertions(+), 32 deletions(-)

diff --git a/graphiti_core/driver/driver.py b/graphiti_core/driver/driver.py
index ac1a3146..0268bc42 100644
--- a/graphiti_core/driver/driver.py
+++ b/graphiti_core/driver/driver.py
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import copy
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Coroutine
@@ -74,6 +75,16 @@ class GraphDriver(ABC):
     def delete_all_indexes(self) -> Coroutine:
         raise NotImplementedError()
 
+    def with_database(self, database: str) -> 'GraphDriver':
+        """
+        Returns a shallow copy of this driver with a different default database.
+        Reuses the same connection (e.g. FalkorDB, Neo4j).
+        """
+        cloned = copy.copy(self)
+        cloned._database = database
+
+        return cloned
+
     @abstractmethod
     async def build_indices_and_constraints(self, delete_existing: bool = False):
         raise NotImplementedError()
diff --git a/graphiti_core/driver/falkordb_driver.py b/graphiti_core/driver/falkordb_driver.py
index 48194ca1..fd33dafb 100644
--- a/graphiti_core/driver/falkordb_driver.py
+++ b/graphiti_core/driver/falkordb_driver.py
@@ -170,38 +170,9 @@ class FalkorDriver(GraphDriver):
             await self.client.connection.close()
 
     async def delete_all_indexes(self) -> None:
-        result = await self.execute_query("CALL db.indexes()")
-        if not result:
-            return
-
-        records, _, _ = result
-        drop_tasks = []
-
-        for record in records:
-            label = record["label"]
-            entity_type = record["entitytype"]
-
-            for field_name, index_type in record["types"].items():
-                if "RANGE" in index_type:
-                    drop_tasks.append(
-                        self.execute_query(f"DROP INDEX ON :{label}({field_name})")
-                    )
-                elif "FULLTEXT" in index_type:
-                    if entity_type == "NODE":
-                        drop_tasks.append(
-                            self.execute_query(
-                                f"DROP FULLTEXT INDEX FOR (n:{label}) ON (n.{field_name})"
-                            )
-                        )
-                    elif entity_type == "RELATIONSHIP":
-                        drop_tasks.append(
-                            self.execute_query(
-                                f"DROP FULLTEXT INDEX FOR ()-[e:{label}]-() ON (e.{field_name})"
-                            )
-                        )
-
-        if drop_tasks:
-            await asyncio.gather(*drop_tasks)
+        await self.execute_query(
+            'CALL db.indexes() YIELD name DROP INDEX name',
+        )
 
     async def build_indices_and_constraints(self, delete_existing=False):
         if delete_existing: