From 6b89bd12577755d52b3f9821505992528ac9bf6a Mon Sep 17 00:00:00 2001
From: Miles Wells <k1o0@5tk.co>
Date: Tue, 15 Oct 2024 12:22:27 +0300
Subject: [PATCH] added remove_cache_table_files function

---
 CHANGELOG.md          |  1 +
 one/alf/cache.py      | 31 ++++++++++++++++++++++++++++++-
 one/api.py            | 21 ++++++++++++++++++++-
 one/tests/test_one.py | 14 ++++++++++++++
 4 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c9a48895..4af786ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -41,6 +41,7 @@ This version fixes issues with Alyx authentication in silent mode, and improves
 ### Added
 
 - one.alf.exceptions.ALFWarning category allows users to filter warnings relating to mixed revisions
+- one.alf.cache.remove_cache_table_files and One._remove_cache_table_files for deleting cache table files
 
 ## [2.9.1]
 
diff --git a/one/alf/cache.py b/one/alf/cache.py
index 8fccaa95..572ba210 100644
--- a/one/alf/cache.py
+++ b/one/alf/cache.py
@@ -32,7 +32,9 @@
 from one.converters import session_record2path
 from one.util import QC_TYPE, patch_cache
 
-__all__ = ['make_parquet_db', 'remove_missing_datasets', 'DATASETS_COLUMNS', 'SESSIONS_COLUMNS']
+__all__ = [
+    'make_parquet_db', 'remove_missing_datasets', 'remove_cache_table_files',
+    'DATASETS_COLUMNS', 'SESSIONS_COLUMNS']
 _logger = logging.getLogger(__name__)
 
 # -------------------------------------------------------------------------------------------------
@@ -338,3 +340,30 @@ def remove_missing_datasets(cache_dir, tables=None, remove_empty_sessions=True,
                 path = path.parent
 
     return sorted(to_delete)
+
+
+def remove_cache_table_files(folder, tables=('sessions', 'datasets')):
+    """Delete cache tables on disk.
+
+    Parameters
+    ----------
+    folder : pathlib.Path
+        The directory path containing cache tables to remove.
+    tables : list of str
+        A list of table names to remove, e.g. ['sessions', 'datasets'].
+        NB: This will also delete the cache_info.json metadata file.
+
+    Returns
+    -------
+    list of pathlib.Path
+        A list of the removed files.
+    """
+    filenames = ('cache_info.json', *(f'{t}.pqt' for t in tables))
+    removed = []
+    for file in map(folder.joinpath, filenames):
+        if file.exists():
+            file.unlink()
+            removed.append(file)
+        else:
+            _logger.warning('%s not found', file)
+    return removed
diff --git a/one/api.py b/one/api.py
index 483b9b93..6f1aa8d9 100644
--- a/one/api.py
+++ b/one/api.py
@@ -28,7 +28,8 @@
 import one.alf.io as alfio
 import one.alf.files as alfiles
 import one.alf.exceptions as alferr
-from .alf.cache import make_parquet_db, DATASETS_COLUMNS, SESSIONS_COLUMNS
+from .alf.cache import (
+    make_parquet_db, remove_cache_table_files, DATASETS_COLUMNS, SESSIONS_COLUMNS)
 from .alf.spec import is_uuid_string, QC, to_alf
 from . import __version__
 from one.converters import ConversionMixin, session_record2path
@@ -111,6 +112,24 @@ def _reset_cache(self):
                 'raw': {}}  # map of original table metadata
         })
 
+    def _remove_cache_table_files(self, tables=None):
+        """Delete cache tables on disk.
+
+        Parameters
+        ----------
+        tables : list of str
+            A list of table names to removes, e.g. ['sessions', 'datasets'].
+            If None, the currently loaded table names are removed. NB: This
+            will also delete the cache_info.json metadata file.
+
+        Returns
+        -------
+        list of pathlib.Path
+            A list of the removed files.
+        """
+        tables = tables or filter(lambda x: x[0] != '_', self._cache)
+        return remove_cache_table_files(self._tables_dir, tables)
+
     def load_cache(self, tables_dir=None, **kwargs):
         """
         Load parquet cache files from a local directory.
diff --git a/one/tests/test_one.py b/one/tests/test_one.py
index 1a2fc55e..e0142f5d 100644
--- a/one/tests/test_one.py
+++ b/one/tests/test_one.py
@@ -60,6 +60,7 @@
 
 class TestONECache(unittest.TestCase):
     """Test methods that use sessions and datasets tables.
+
     This class loads the parquet tables from the fixtures and builds a file tree in a temp folder
     """
     tempdir = None
@@ -1025,6 +1026,19 @@ def test_save_loaded_ids(self):
         self.assertEqual(ids, [])
         self.assertIsNone(filename)
 
+    def test_remove_cache_table_files(self):
+        """Test One._remove_cache_table_files method."""
+        root = self.one._tables_dir
+        for name in ('cache_info.json', 'foo.pqt'):
+            root.joinpath(name).touch()
+        removed = self.one._remove_cache_table_files()
+        expected = ['sessions.pqt', 'datasets.pqt', 'cache_info.json']
+        self.assertCountEqual(expected, [str(x.relative_to(root)) for x in removed])
+        with self.assertLogs('one.alf.cache', 30) as cm:
+            removed = self.one._remove_cache_table_files(tables=('foo',))
+        self.assertEqual([root / 'foo.pqt'], removed)
+        self.assertIn('cache_info.json not found', cm.records[0].message)
+
 
 @unittest.skipIf(OFFLINE_ONLY, 'online only test')
 class TestOneAlyx(unittest.TestCase):