From 03bbfe0d2ef5d71291f999db303cc58c84956dd4 Mon Sep 17 00:00:00 2001
From: James Mathews <mathewj2@mskcc.org>
Date: Wed, 23 Aug 2023 15:45:51 -0400
Subject: [PATCH] Add db managements script to drop pending computatations
 (#196)

* Adding computed feature dropper.

* Add logging statements to dropper.

* Correct drop order.
---
 pyproject.toml.unversioned                    |  3 +-
 .../db/ondemand_dropper.py                    | 58 +++++++++++++++++++
 .../db/scripts/drop_ondemand_computations.py  | 40 +++++++++++++
 3 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 spatialprofilingtoolbox/db/ondemand_dropper.py
 create mode 100644 spatialprofilingtoolbox/db/scripts/drop_ondemand_computations.py

diff --git a/pyproject.toml.unversioned b/pyproject.toml.unversioned
index c832a1974..688409c6c 100644
--- a/pyproject.toml.unversioned
+++ b/pyproject.toml.unversioned
@@ -165,7 +165,8 @@ packages = [
     "status.py",
     "retrieve_feature_matrices.py",
     "drop.py",
-    "index_expressions_table.py"
+    "index_expressions_table.py",
+    "drop_ondemand_computations.py"
 ]
 "spatialprofilingtoolbox.db.data_model" = [
     "create_roles.sql",
diff --git a/spatialprofilingtoolbox/db/ondemand_dropper.py b/spatialprofilingtoolbox/db/ondemand_dropper.py
new file mode 100644
index 000000000..f35ef61ba
--- /dev/null
+++ b/spatialprofilingtoolbox/db/ondemand_dropper.py
@@ -0,0 +1,58 @@
+"""Drop ondemand-computed feature values, specifications, etc."""
+
+from typing import cast
+import re
+
+from psycopg2.extensions import cursor as Psycopg2Cursor
+
+from spatialprofilingtoolbox.standalone_utilities.log_formats import colorized_logger
+
+logger = colorized_logger(__name__)
+
+
+class OnDemandComputationsDropper:
+    """Drop ondemand-computed feature values, specifications, etc."""
+
+    @staticmethod
+    def drop(cursor: Psycopg2Cursor, pending_only: bool = False, drop_all: bool = False):
+        specifications = cast(list[str], OnDemandComputationsDropper.get_droppable(
+            cursor,
+            pending_only=pending_only,
+            drop_all=drop_all,
+        ))
+        OnDemandComputationsDropper.drop_features(cursor, specifications)
+
+    @staticmethod
+    def get_droppable(
+        cursor: Psycopg2Cursor,
+        pending_only: bool = False,
+        drop_all: bool = False,
+    ) -> list[str] | None:
+        if pending_only:
+            cursor.execute('SELECT feature_specification FROM pending_feature_computation;')
+            return [row[0] for row in cursor.fetchall()]
+        if drop_all:
+            cursor.execute('SELECT DISTINCT study FROM feature_specification;')
+            studies = [row[0] for row in cursor.fetchall()]
+            studies = [s for s in studies if re.search(r' ondemand computed features$', s)]
+            specifications: list[str] = []
+            for study in studies:
+                query = 'SELECT identifier FROM feature_specification WHERE study=%s ;'
+                cursor.execute(query, (study,))
+                _specifications = [row[0] for row in cursor.fetchall()]
+                specifications = specifications + _specifications
+            return specifications
+        return None
+
+    @staticmethod
+    def drop_features(cursor: Psycopg2Cursor, specifications: list[str]):
+        for specification in specifications:
+            queries = [
+                'DELETE FROM pending_feature_computation WHERE feature_specification=%s ;',
+                'DELETE FROM quantitative_feature_value WHERE feature=%s ;',
+                'DELETE FROM feature_specifier WHERE feature_specification=%s ;',
+                'DELETE FROM feature_specification WHERE identifier=%s ;',
+            ]
+            for query in queries:
+                logger.debug(query, specification)
+                cursor.execute(query, (specification,))
diff --git a/spatialprofilingtoolbox/db/scripts/drop_ondemand_computations.py b/spatialprofilingtoolbox/db/scripts/drop_ondemand_computations.py
new file mode 100644
index 000000000..6433bcbcc
--- /dev/null
+++ b/spatialprofilingtoolbox/db/scripts/drop_ondemand_computations.py
@@ -0,0 +1,40 @@
+"""Utility to managed features computed on demand."""
+import argparse
+
+from spatialprofilingtoolbox.db.ondemand_dropper import OnDemandComputationsDropper
+from spatialprofilingtoolbox.db.database_connection import get_and_validate_database_config
+from spatialprofilingtoolbox.workflow.common.cli_arguments import add_argument
+from spatialprofilingtoolbox import DBCursor
+
+from spatialprofilingtoolbox.standalone_utilities.log_formats import colorized_logger
+
+logger = colorized_logger('spt db status')
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog='spt db drop-ondemand-computations',
+        description='Report basic health status of the given scstudies database.'
+    )
+    add_argument(parser, 'database config')
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        '--pending-only',
+        action='store_true',
+        default=False,
+    )
+    group.add_argument(
+        '--all',
+        action='store_true',
+        default=False,
+    )
+    args = parser.parse_args()
+
+    config_file = get_and_validate_database_config(args)
+    with DBCursor(database_config_file=config_file) as cursor:
+        if args.pending_only:
+            OnDemandComputationsDropper.drop(cursor, pending_only=True)
+        if args.all:
+            OnDemandComputationsDropper.drop(cursor, drop_all=True)
+
+if __name__ == '__main__':
+    main()