yangdx commited on
Commit
7c49266
·
1 Parent(s): d620bc2

Add drop support for AGE storage

Browse files
Files changed (1) hide show
  1. lightrag/kg/age_impl.py +21 -3
lightrag/kg/age_impl.py CHANGED
@@ -34,9 +34,9 @@ if not pm.is_installed("psycopg-pool"):
34
  if not pm.is_installed("asyncpg"):
35
  pm.install("asyncpg")
36
 
37
- import psycopg
38
- from psycopg.rows import namedtuple_row
39
- from psycopg_pool import AsyncConnectionPool, PoolTimeout
40
 
41
 
42
  class AGEQueryException(Exception):
@@ -871,3 +871,21 @@ class AGEStorage(BaseGraphStorage):
871
  async def index_done_callback(self) -> None:
872
  # AGES handles persistence automatically
873
  pass
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34
  if not pm.is_installed("asyncpg"):
35
  pm.install("asyncpg")
36
 
37
+ import psycopg # type: ignore
38
+ from psycopg.rows import namedtuple_row # type: ignore
39
+ from psycopg_pool import AsyncConnectionPool, PoolTimeout # type: ignore
40
 
41
 
42
  class AGEQueryException(Exception):
 
871
  async def index_done_callback(self) -> None:
872
  # AGES handles persistence automatically
873
  pass
874
+
875
+ async def drop(self) -> dict[str, str]:
876
+ """Drop the storage by removing all nodes and relationships in the graph.
877
+
878
+ Returns:
879
+ dict[str, str]: Status of the operation with keys 'status' and 'message'
880
+ """
881
+ try:
882
+ query = """
883
+ MATCH (n)
884
+ DETACH DELETE n
885
+ """
886
+ await self._query(query)
887
+ logger.info(f"Successfully dropped all data from graph {self.graph_name}")
888
+ return {"status": "success", "message": "graph data dropped"}
889
+ except Exception as e:
890
+ logger.error(f"Error dropping graph {self.graph_name}: {e}")
891
+ return {"status": "error", "message": str(e)}