Add drop support for Oracle

This commit is contained in:
yangdx
2025-03-31 03:19:06 +08:00
parent 6a51f38cae
commit 9959ea90a8

View File

@@ -27,7 +27,7 @@ if not pm.is_installed("oracledb"):
pm.install("oracledb") pm.install("oracledb")
from graspologic import embed from graspologic import embed
import oracledb import oracledb # type: ignore
class OracleDB: class OracleDB:
@@ -392,6 +392,21 @@ class OracleKVStorage(BaseKVStorage):
# Oracle handles persistence automatically # Oracle handles persistence automatically
pass pass
async def drop(self) -> dict[str, str]:
"""Drop the storage"""
try:
table_name = namespace_to_table_name(self.namespace)
if not table_name:
return {"status": "error", "message": f"Unknown namespace: {self.namespace}"}
drop_sql = SQL_TEMPLATES["drop_specifiy_table_workspace"].format(
table_name=table_name
)
await self.db.execute(drop_sql, {"workspace": self.db.workspace})
return {"status": "success", "message": "data dropped"}
except Exception as e:
return {"status": "error", "message": str(e)}
@final @final
@dataclass @dataclass
@@ -605,6 +620,21 @@ class OracleVectorDBStorage(BaseVectorStorage):
logger.error(f"Error retrieving vector data for IDs {ids}: {e}") logger.error(f"Error retrieving vector data for IDs {ids}: {e}")
return [] return []
async def drop(self) -> dict[str, str]:
"""Drop the storage"""
try:
table_name = namespace_to_table_name(self.namespace)
if not table_name:
return {"status": "error", "message": f"Unknown namespace: {self.namespace}"}
drop_sql = SQL_TEMPLATES["drop_specifiy_table_workspace"].format(
table_name=table_name
)
await self.db.execute(drop_sql, {"workspace": self.db.workspace})
return {"status": "success", "message": "data dropped"}
except Exception as e:
return {"status": "error", "message": str(e)}
@final @final
@dataclass @dataclass
@@ -933,6 +963,21 @@ class OracleGraphStorage(BaseGraphStorage):
logger.error(f"Error retrieving entity types: {e}") logger.error(f"Error retrieving entity types: {e}")
return [] return []
async def drop(self) -> dict[str, str]:
"""Drop the storage"""
try:
# 使用图形查询删除所有节点和关系
delete_edges_sql = """DELETE FROM LIGHTRAG_GRAPH_EDGES WHERE workspace=:workspace"""
await self.db.execute(delete_edges_sql, {"workspace": self.db.workspace})
delete_nodes_sql = """DELETE FROM LIGHTRAG_GRAPH_NODES WHERE workspace=:workspace"""
await self.db.execute(delete_nodes_sql, {"workspace": self.db.workspace})
return {"status": "success", "message": "graph data dropped"}
except Exception as e:
logger.error(f"Error dropping graph: {e}")
return {"status": "error", "message": str(e)}
async def get_knowledge_graph( async def get_knowledge_graph(
self, node_label: str, max_depth: int = 5 self, node_label: str, max_depth: int = 5
) -> KnowledgeGraph: ) -> KnowledgeGraph:
@@ -1343,4 +1388,6 @@ SQL_TEMPLATES = {
MATCH (a) MATCH (a)
WHERE a.workspace=:workspace AND a.name=:node_id WHERE a.workspace=:workspace AND a.name=:node_id
ACTION DELETE a)""", ACTION DELETE a)""",
# Drop tables
"drop_specifiy_table_workspace": "DELETE FROM {table_name} WHERE workspace=:workspace",
} }