yangdx
commited on
Commit
·
7c49266
1
Parent(s):
d620bc2
Add drop support for AGE storage
Browse files- 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)}
|