yangdx commited on
Commit
e957265
·
1 Parent(s): 6221908

Fix update status handling bugs in drop function of json kv storage

Browse files
lightrag/kg/json_doc_status_impl.py CHANGED
@@ -141,9 +141,14 @@ class JsonDocStatusStorage(DocStatusStorage):
141
  None
142
  """
143
  async with self._storage_lock:
 
144
  for doc_id in doc_ids:
145
- self._data.pop(doc_id, None)
146
- await set_all_update_flags(self.namespace)
 
 
 
 
147
 
148
  async def drop(self) -> dict[str, str]:
149
  """Drop all document status data from storage and clean up resources
@@ -160,7 +165,9 @@ class JsonDocStatusStorage(DocStatusStorage):
160
  """
161
  try:
162
  async with self._storage_lock:
163
- self._data.update({})
 
 
164
  await self.index_done_callback()
165
  logger.info(f"Process {os.getpid()} drop {self.namespace}")
166
  return {"status": "success", "message": "data dropped"}
 
141
  None
142
  """
143
  async with self._storage_lock:
144
+ any_deleted = False
145
  for doc_id in doc_ids:
146
+ result = self._data.pop(doc_id, None)
147
+ if result is not None:
148
+ any_deleted = True
149
+
150
+ if any_deleted:
151
+ await set_all_update_flags(self.namespace)
152
 
153
  async def drop(self) -> dict[str, str]:
154
  """Drop all document status data from storage and clean up resources
 
165
  """
166
  try:
167
  async with self._storage_lock:
168
+ self._data.clear()
169
+ await set_all_update_flags(self.namespace)
170
+
171
  await self.index_done_callback()
172
  logger.info(f"Process {os.getpid()} drop {self.namespace}")
173
  return {"status": "success", "message": "data dropped"}
lightrag/kg/json_kv_impl.py CHANGED
@@ -140,9 +140,14 @@ class JsonKVStorage(BaseKVStorage):
140
  None
141
  """
142
  async with self._storage_lock:
 
143
  for doc_id in ids:
144
- self._data.pop(doc_id, None)
145
- await set_all_update_flags(self.namespace)
 
 
 
 
146
 
147
  async def drop_cache_by_modes(self, modes: list[str] | None = None) -> bool:
148
  """Delete specific records from storage by by cache mode
@@ -183,7 +188,9 @@ class JsonKVStorage(BaseKVStorage):
183
  """
184
  try:
185
  async with self._storage_lock:
186
- self._data.update({})
 
 
187
  await self.index_done_callback()
188
  logger.info(f"Process {os.getpid()} drop {self.namespace}")
189
  return {"status": "success", "message": "data dropped"}
 
140
  None
141
  """
142
  async with self._storage_lock:
143
+ any_deleted = False
144
  for doc_id in ids:
145
+ result = self._data.pop(doc_id, None)
146
+ if result is not None:
147
+ any_deleted = True
148
+
149
+ if any_deleted:
150
+ await set_all_update_flags(self.namespace)
151
 
152
  async def drop_cache_by_modes(self, modes: list[str] | None = None) -> bool:
153
  """Delete specific records from storage by by cache mode
 
188
  """
189
  try:
190
  async with self._storage_lock:
191
+ self._data.clear()
192
+ await set_all_update_flags(self.namespace)
193
+
194
  await self.index_done_callback()
195
  logger.info(f"Process {os.getpid()} drop {self.namespace}")
196
  return {"status": "success", "message": "data dropped"}