nmreval/src/gui_qt/lib/backup.py

import os
import sqlite3
from datetime import datetime
from pathlib import Path

from nmreval.configs import config_paths
from ..Qt import QtCore, QtWidgets

DB_FILE = '/tmp/nmreval.db'


class BackupManager(QtCore.QObject):
    def __init__(self):
        super().__init__()
        self.create_table()
        self._pid = None

    def create_table(self):
        con = sqlite3.connect(DB_FILE)
        con.execute(
            "CREATE TABLE IF NOT EXISTS sessions "
            "(pid INTEGER NOT NULL, backup_file TEXT NOT NULL, last_save TEXT);"
        )

    def create_entry(self, pid: int):
        backup_path = config_paths() / f'autosave_{datetime.now().strftime("%Y-%m-%d_%H%M%S")}_{pid}.nmr'

        con = sqlite3.connect(DB_FILE)
        con.execute('INSERT INTO sessions VALUES(?, ?, ?);',
                    (pid, str(backup_path), None))
        con.commit()
        con.close()
        self._pid = pid

        return backup_path

    def update_last_save(self):
        con = sqlite3.connect(DB_FILE)
        con.execute(
            'UPDATE sessions SET last_save = ? WHERE pid = ?',
            (datetime.now().strftime("%Y-%m-%d %H:%M:%S"), self._pid)
        )
        con.commit()
        con.close()

    def search_unsaved(self):
        con = sqlite3.connect(DB_FILE)
        cursor = con.cursor()
        res = cursor.execute('SELECT sessions.* FROM sessions;')
        con.commit()
        data = res.fetchall()
        con.close()

        missing_processes = []

        for pid, fname, save_date in data:
            try:
                os.kill(pid, 0)
            except ProcessLookupError:
                if Path(fname).exists():
                    missing_processes.append((pid, fname, save_date))
                else:
                    # remove entries without valid file
                    self.remove_row(pid)

        if missing_processes:
            msg = QLonelyBackupWindow()
            msg.add_files(missing_processes)
            msg.exec()
            toberead = msg.recover

            for pid in toberead[0]:
                self.remove_file(pid)

            for pid in toberead[1]:
                self.remove_row(pid)

            return list(toberead[1].values())

        return []

    def remove_row(self, pid):
        con = sqlite3.connect(DB_FILE)
        con.execute('DELETE FROM sessions WHERE pid = ?', (pid,))
        con.commit()
        con.close()

    def remove_file(self, pid: int = None):
        if pid is None:
            pid = self._pid

        con = sqlite3.connect(DB_FILE)
        cursor = con.cursor()

        # remove backup file
        res = cursor.execute('SELECT sessions.backup_file FROM sessions WHERE pid = ?', (pid,))
        con.commit()

        fname = Path(res.fetchone()[0])
        con.close()

        fname.unlink(missing_ok=True)

        # because autosave backups in a *.nmr.0 file and then moves it to *.nmr we look also for this one.
        fname.with_suffix('.nmr.0').unlink(missing_ok=True)

        # after removal of file also remove entry
        self.remove_row(pid)

    def delete_db_if_empty(self):
        con = sqlite3.connect(DB_FILE)
        cursor = con.cursor()
        res = cursor.execute('SELECT COUNT(sessions.pid) FROM sessions GROUP BY sessions.pid;')
        con.commit()

        remaining_processes = res.fetchone()
        con.close()
        if remaining_processes is None:
            Path(DB_FILE).unlink()

    def close(self):
        self.remove_file()
        self.delete_db_if_empty()


class QLonelyBackupWindow(QtWidgets.QDialog):
    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.recover = [[], {}]    # list of pid to delete, dict of files to read {pid: file}

        self.setWindowTitle('Adopt a file!')
        self.resize(720, 320)

        layout = QtWidgets.QVBoxLayout(self)
        self.label = QtWidgets.QLabel(self)
        self.label.setText('Abandoned backup file(s) looking for a loving home!\n'
                           '(Files will all be loaded to same instance)')
        layout.addWidget(self.label)

        self.table = QtWidgets.QTableWidget(self)
        self.table.setColumnCount(4)
        self.table.setHorizontalHeaderLabels(['File name', 'Last saved', 'File size', 'Action'])
        self.table.setGridStyle(QtCore.Qt.PenStyle.DashLine)
        # self.table.horizontalHeader().setStretchLastSection(True)
        self.table.verticalHeader().setVisible(False)
        self.table.horizontalHeader().setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
        layout.addWidget(self.table)

        self.buttons = QtWidgets.QDialogButtonBox(self)
        self.buttons.setOrientation(QtCore.Qt.Orientation.Horizontal)
        self.buttons.setStandardButtons(QtWidgets.QDialogButtonBox.Ok)
        layout.addWidget(self.buttons)

        self.buttons.accepted.connect(self.accept)

    def add_files(self, entries):
        self.table.setRowCount(len(entries))
        for i, (pid, path, date) in enumerate(entries):
            path = Path(path)
            item1 = QtWidgets.QTableWidgetItem(path.name)
            item1.setFlags(QtCore.Qt.ItemFlag.ItemIsEnabled)
            item1.setData(QtCore.Qt.ItemDataRole.UserRole, pid)
            item1.setData(QtCore.Qt.ItemDataRole.UserRole+1, path)
            self.table.setItem(i, 0, item1)

            item2 = QtWidgets.QTableWidgetItem(date)
            item2.setFlags(QtCore.Qt.ItemFlag.ItemIsEnabled)
            self.table.setItem(i, 1, item2)

            size = path.stat().st_size
            size_cnt = 0
            while size > 1024:
                # make file size human-readable
                size /= 1024
                size_cnt += 1
                if size_cnt == 5:
                    break

            byte = ['bytes', 'kB', 'MiB', 'GiB', 'TiB', 'PiB'][size_cnt]

            item3 = QtWidgets.QTableWidgetItem(f'{size:.2f} {byte}')
            item3.setFlags(QtCore.Qt.ItemFlag.ItemIsEnabled)
            self.table.setItem(i, 2, item3)

            cw = QtWidgets.QComboBox(self)
            cw.addItems(['Load', 'Delete', 'Keep for later'])
            self.table.setCellWidget(i, 3, cw)

        self.table.resizeColumnsToContents()

    def accept(self):
        for i in range(self.table.rowCount()):
            decision = self.table.cellWidget(i, 3).currentIndex()
            item = self.table.item(i, 0)
            pid = item.data(QtCore.Qt.ItemDataRole.UserRole)
            if decision == 0:
                # load file
                self.recover[1][pid] = item.data(QtCore.Qt.ItemDataRole.UserRole+1)
            elif decision == 1:
                # delete
                self.recover[0].append(pid)
            else:
                # do nothing
                pass

        super().accept()