Add support for deleting variables

damnit/backend/api.py

@@ -1,11 +1,12 @@
+import glob
 from pathlib import Path
 from contextlib import contextmanager
 
 import h5py
 import xarray as xr
 
 from .db import DamnitDB
-from ..ctxsupport.ctxrunner import DataType
+from ..ctxsupport.ctxrunner import DataType, add_to_h5_file
 
 
 class VariableData:
@@ -117,3 +118,14 @@ def __getitem__(self, name):
     @property
     def file(self):
         return self._h5_path
+
+def delete_variable(db, name):
+    # Remove from the database
+    db.delete_variable(name)
+
+    # And the HDF5 files
+    for h5_path in glob.glob(f"{db.path.parent}/extracted_data/*.h5"):
+        with add_to_h5_file(h5_path) as f:
+            if name in f:
+                del f[f".reduced/{name}"]
+                del f[name]

damnit/backend/db.py

@@ -73,6 +73,8 @@ def db_path(root_path: Path):
 
 class DamnitDB:
     def __init__(self, path=DB_NAME, allow_old=False):
+        self.path = path.absolute()
+
         db_existed = path.exists()
         log.debug("Opening database at %s", path)
         self.conn = sqlite3.connect(path, timeout=30)
@@ -187,19 +189,20 @@ def update_views(self):
         max_diff_cols = ", ".join([col_select_sql.format(var=var, col="max_diff")
                                for var in variables])
 
-        self.conn.executescript(f"""
-            DROP VIEW IF EXISTS runs;
-            CREATE VIEW runs
-            AS SELECT run_info.proposal, run_info.run, start_time, added_at, {runs_cols}
-               FROM run_variables INNER JOIN run_info ON run_variables.proposal = run_info.proposal AND run_variables.run = run_info.run
-               GROUP BY run_info.run;
-
-            DROP VIEW IF EXISTS max_diffs;
-            CREATE VIEW max_diffs
-            AS SELECT proposal, run, {max_diff_cols}
-               FROM run_variables
-               GROUP BY run;
-        """)
+        with self.conn:
+            self.conn.executescript(f"""
+                DROP VIEW IF EXISTS runs;
+                CREATE VIEW runs
+                AS SELECT run_info.proposal, run_info.run, start_time, added_at, {runs_cols}
+                   FROM run_variables INNER JOIN run_info ON run_variables.proposal = run_info.proposal AND run_variables.run = run_info.run
+                   GROUP BY run_info.run;
+
+                DROP VIEW IF EXISTS max_diffs;
+                CREATE VIEW max_diffs
+                AS SELECT proposal, run, {max_diff_cols}
+                   FROM run_variables
+                   GROUP BY run;
+            """)
 
     def set_variable(self, proposal: int, run: int, name: str, reduced):
         timestamp = datetime.now(tz=timezone.utc).timestamp()
@@ -245,6 +248,22 @@ def set_variable(self, proposal: int, run: int, name: str, reduced):
             if is_new:
                 self.update_views()
 
+    def delete_variable(self, name: str):
+        with self.conn:
+            # First delete from the `variables` table
+            self.conn.execute("""
+            DELETE FROM variables
+            WHERE name = ?
+            """, (name,))
+
+            # And then `run_variables`
+            self.conn.execute("""
+            DELETE FROM run_variables
+            WHERE name = ?
+            """, (name, ))
+
+            self.update_views()
+
 class MetametaMapping(MutableMapping):
     def __init__(self, conn):
         self.conn = conn

damnit/gui/main_window.py

@@ -402,9 +402,9 @@ def open_column_dialog(self):
             self._columns_dialog.setWindowTitle("Column settings")
             layout = QtWidgets.QVBoxLayout()
 
-            layout.addWidget(QtWidgets.QLabel("These columns can be hidden but not reordered:"))
+            layout.addWidget(QtWidgets.QLabel("These columns can be hidden but not reordered or deleted:"))
             layout.addWidget(self.table_view._static_columns_widget)
-            layout.addWidget(QtWidgets.QLabel("Drag these columns to reorder them:"))
+            layout.addWidget(QtWidgets.QLabel("Drag these columns to reorder them, right-click to delete:"))
             layout.addWidget(self.table_view._columns_widget)
             self._columns_dialog.setLayout(layout)
 
@@ -460,7 +460,7 @@ def _create_menu_bar(self) -> None:
         fileMenu.addAction(action_exit)
 
         # Table menu
-        action_columns = QtWidgets.QAction("Select && reorder columns", self)
+        action_columns = QtWidgets.QAction("Select, delete, && reorder columns", self)
         action_columns.triggered.connect(self.open_column_dialog)
         self.action_autoscroll = QtWidgets.QAction('Scroll to newly added runs', self)
         self.action_autoscroll.setCheckable(True)

damnit/gui/table.py

@@ -5,7 +5,9 @@
 
 from PyQt5 import QtCore, QtWidgets, QtGui
 from PyQt5.QtCore import Qt
+from PyQt5.QtWidgets import QMessageBox
 
+from ..backend.api import delete_variable
 from ..backend.db import BlobTypes
 from ..util import StatusbarStylesheet, timestamp2str
 
@@ -44,6 +46,9 @@ def __init__(self) -> None:
         self._columns_widget.itemChanged.connect(self.item_changed)
         self._columns_widget.model().rowsMoved.connect(self.item_moved)
 
+        self._columns_widget.setContextMenuPolicy(Qt.CustomContextMenu)
+        self._columns_widget.customContextMenuRequested.connect(self.show_delete_menu)
+
         self._static_columns_widget.itemChanged.connect(self.item_changed)
         self._static_columns_widget.setStyleSheet("QListWidget {padding: 0px;} QListWidget::item { margin: 5px; }")
         self._columns_widget.setStyleSheet("QListWidget {padding: 0px;} QListWidget::item { margin: 5px; }")
@@ -112,6 +117,38 @@ def item_moved(self, parent, start, end, destination, row):
 
         self.settings_changed.emit()
 
+    def show_delete_menu(self, pos):
+        item = self._columns_widget.itemAt(pos)
+        if item is None:
+            # This happens if the user clicks on blank space inside the widget
+            return
+
+        global_pos = self._columns_widget.mapToGlobal(pos)
+        menu = QtWidgets.QMenu()
+        menu.addAction("Delete")
+        action = menu.exec(global_pos)
+        if action is not None:
+            name = self.model()._main_window.col_title_to_name(item.text())
+            self.confirm_delete_variable(name)
+
+    def confirm_delete_variable(self, name):
+        button = QMessageBox.warning(self, "Confirm deletion",
+                                     f"You are about to permanently delete the variable <b>'{name}'</b> "
+                                     "from the database and HDF5 files. This cannot be undone. "
+                                     "Are you sure you want to continue?",
+                                     QMessageBox.Yes | QMessageBox.No,
+                                     defaultButton=QMessageBox.No)
+        if button == QMessageBox.Yes:
+            main_window = self.model()._main_window
+            delete_variable(main_window.db, name)
+
+            # TODO: refactor this into simply removing the column from the table
+            # if we fix the bugs around adding/removing columns
+            # on-the-fly. Currently there are some lingering off-by-one errors
+            # or something that cause the wrong columns to be moved when moving
+            # a column after the number of columns has changed.
+            main_window.autoconfigure(main_window.context_dir)
+
     def add_new_columns(self, columns, statuses, positions = None):
         if positions is None:
             rows_count = self._columns_widget.count()

tests/test_gui.py

@@ -5,6 +5,7 @@
 from unittest.mock import patch
 from types import SimpleNamespace
 
+import h5py
 import pytest
 import numpy as np
 import pandas as pd
@@ -810,3 +811,29 @@ def image(run):
     # Check that images are formatted nicely
     df = pd.read_excel(export_path) if extension == ".xlsx" else pd.read_csv(export_path)
     assert df["Image"][0] == "<image>"
+
+def test_delete_variable(mock_db_with_data, qtbot, monkeypatch):
+    db_dir, db = mock_db_with_data
+    monkeypatch.chdir(db_dir)
+
+    # We'll delete the 'array' variable
+    assert "array" in db.variable_names()
+    win = MainWindow(db_dir, connect_to_kafka=False)
+
+    # If the user clicks 'No' then we should do nothing
+    with patch.object(QMessageBox, "warning", return_value=QMessageBox.No) as warning:
+        win.table_view.confirm_delete_variable("array")
+        warning.assert_called_once()
+    assert "array" in db.variable_names()
+
+    # Otherwise it should be deleted from the database and HDF5 files
+    with patch.object(QMessageBox, "warning", return_value=QMessageBox.Yes) as warning:
+        win.table_view.confirm_delete_variable("array")
+        warning.assert_called_once()
+
+    assert "array" not in db.variable_names()
+
+    proposal = db.metameta['proposal']
+    with h5py.File(db_dir / f"extracted_data/p{proposal}_r1.h5") as f:
+        assert "array" not in f.keys()
+        assert "array" not in f[".reduced"].keys()