[resourcemodels] Add database sets modelization

beat/editor/backend/resourcemodels.py

@@ -29,6 +29,8 @@ from PyQt5.QtSql import QSqlDatabase
 from PyQt5.QtSql import QSqlQuery
 from PyQt5.QtSql import QSqlTableModel
 
+from beat.core.database import Database
+
 from .asset import Asset
 from .asset import AssetType
 from .assetmodel import AssetModel
@@ -74,14 +76,18 @@ class ExperimentResources:
         QUEUE_TABLE_CLEANUP = "DROP TABLE IF EXISTS queues"
         QUEUE_TABLE = "CREATE TABLE queues(name varchar, env_name varchar, env_version varchar, env_type varchar)"
         INSERT_QUEUE = "INSERT INTO queues(name, env_name, env_version, env_type) VALUES (?, ?, ?, ?)"
+        DATASET_TABLE_CLEANUP = "DROP TABLE IF EXISTS datasets"
+        DATASET_TABLE = "CREATE TABLE datasets(name varchar, outputs integer)"
+        INSERT_DATASET = "INSERT INTO datasets(name, outputs) VALUES(?, ?)"
+
 
         query = QSqlQuery()
 
-        for query_str in [ALGORITHM_TABLE_CLEANUP, QUEUE_TABLE_CLEANUP]:
+        for query_str in [ALGORITHM_TABLE_CLEANUP, QUEUE_TABLE_CLEANUP, DATASET_TABLE_CLEANUP]:
             if not query.exec_(query_str):
                 raise RuntimeError(f"Failed to drop table: {query.lastError().text()}")
 
-        for query_str in [ALGORITHM_TABLE, QUEUE_TABLE]:
+        for query_str in [ALGORITHM_TABLE, QUEUE_TABLE, DATASET_TABLE]:
             if not query.exec_(query_str):
                 raise RuntimeError(
                     f"Failed to create table: {query.lastError().text()}"
@@ -157,6 +163,31 @@ class ExperimentResources:
                     f"Failed to insert queue: {query.lastError().text()}"
                 )
 
+        model.asset_type = AssetType.DATABASE
+
+        if not query.prepare(INSERT_DATASET):
+            raise RuntimeError(f"Failed to prepare query: {query.lastError().text()}")
+
+        for database_name in model.stringList():
+            database = Database(prefix_path, database_name)
+            if not database.valid:
+                print("Skipping invalid database:", database_name)
+                continue
+            protocols = database.protocol_names
+            for protocol_name in protocols:
+                sets = database.set_names(protocol_name)
+                for set_name in sets:
+                    set_data = database.set(protocol_name, set_name)
+                    name = f"{database_name}/{protocol_name}/{set_name}"
+                    output_count = len(set_data["outputs"])
+                    query.addBindValue(name)
+                    query.addBindValue(output_count)
+
+                    if not query.exec_():
+                        raise RuntimeError(
+                            f"Failed to insert dataset: {query.lastError().text()}"
+                        )
+
 
 class AlgorithmResourceModel(QSqlTableModel):
     def __init__(self, parent=None):
@@ -264,3 +295,27 @@ class QueueResourceModel(QSqlTableModel):
             print([self.index(i, j).data() for j in range(4)])
 
 
+class DatasetResourceModel(QSqlTableModel):
+    def __init__(self, parent=None):
+        super().__init__(parent=parent)
+
+        self._analyzer_enabled = False
+        self._output_count = None
+
+        self.setTable("datasets")
+        self.select()
+
+    def update(self):
+        filter_str = f""
+
+        if self._output_count is not None:
+            filter_str += f"outputs={self._output_count}"
+
+        self.setFilter(filter_str)
+
+    def setOutputCount(self, count):
+        if self._output_count == count:
+            return
+
+        self._output_count = count
+        self.update()