From 701e50e8f15c996ac41d18fdec648573f0580edd Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Sat, 10 Jun 2023 02:45:45 +0300
Subject: [PATCH] index: introduce delete_node

Needed for https://github.com/iterative/dvc/pull/9424
---
 setup.cfg                   |  2 +-
 src/dvc_data/index/index.py | 11 +++++++++++
 src/dvc_data/index/view.py  |  5 +++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/setup.cfg b/setup.cfg
index f7054ca9..424dd5f2 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -32,7 +32,7 @@ install_requires=
     diskcache>=5.2.1
     nanotime>=0.5.2
     attrs>=21.3.0
-    sqltrie>=0.3.1,<1
+    sqltrie>=0.5.0,<1
 
 [options.extras_require]
 cli =
diff --git a/src/dvc_data/index/index.py b/src/dvc_data/index/index.py
index c90ab604..600948e4 100644
--- a/src/dvc_data/index/index.py
+++ b/src/dvc_data/index/index.py
@@ -107,6 +107,10 @@ def __delitem__(self, key):
         self._cache.pop(key, None)
         super().__delitem__(key)
 
+    def delete_node(self, key):
+        self._cache.pop(key, None)
+        super().delete_node(key)
+
 
 class Storage(ABC):
     def __init__(self, key: "DataIndexKey"):
@@ -390,6 +394,10 @@ def traverse(self, node_factory: Callable, **kwargs) -> Any:
     def has_node(self, key: DataIndexKey) -> bool:
         pass
 
+    @abstractmethod
+    def delete_node(self, key: DataIndexKey) -> None:
+        pass
+
     @abstractmethod
     def longest_prefix(
         self, key: DataIndexKey
@@ -629,6 +637,9 @@ def load(self, **kwargs):
     def has_node(self, key: DataIndexKey) -> bool:
         return self._trie.has_node(key)
 
+    def delete_node(self, key: DataIndexKey) -> None:
+        return self._trie.delete_node(key)
+
     def shortest_prefix(self, *args, **kwargs):
         return self._trie.shortest_prefix(*args, **kwargs)
 
diff --git a/src/dvc_data/index/view.py b/src/dvc_data/index/view.py
index 406f7f41..aec2956d 100644
--- a/src/dvc_data/index/view.py
+++ b/src/dvc_data/index/view.py
@@ -137,6 +137,11 @@ def _filter_fn(entry):
     def has_node(self, key: DataIndexKey) -> bool:
         return self.filter_fn(key) and self._index.has_node(key)
 
+    def delete_node(self, key: DataIndexKey) -> None:
+        if not self.filter_fn(key):
+            raise KeyError
+        self._index.delete_node(key)
+
     def longest_prefix(
         self, key: DataIndexKey
     ) -> Tuple[Optional[DataIndexKey], Optional[DataIndexEntry]]: