improvement(db-scripts): improve import/export scripts

Import/export scripts missed graphs data which is needed for debugging
localy issues on prod. Added support for graphs data.

dev-db/export_data.py

@@ -1,33 +1,95 @@
 import json
 from pathlib import Path
+from collections import defaultdict
 
 from argus.backend.models.web import ArgusRelease, ArgusTest, ArgusGroup
 from argus.backend.plugins.sct.testrun import SCTTestRun
 from argus.backend.db import ScyllaCluster
 from argus.backend.util.encoders import ArgusJSONEncoder
+from argus.backend.models.result import (
+    ArgusGenericResultMetadata,
+    ArgusGenericResultData,
+    ArgusBestResultData,
+    ArgusGraphView,
+)
+
 db = ScyllaCluster.get()
 
-release_name = "perf-regression" # change name here
+release_name = "scylla-enterprise"  # change name here
 dest = Path(__file__).parent / "sample_data" / release_name
 dest.mkdir(parents=True, exist_ok=True)
 release = ArgusRelease.get(name=release_name)
 print(f"Saving release {release.id}")
 (dest / "release.json").write_text(json.dumps(release, cls=ArgusJSONEncoder))
 
-print('getting groups data')
+print("getting groups data")
 for group in ArgusGroup.filter(release_id=release.id).all():
     print(f"Saving group {group.id}")
-    (dest / f"group_{group.id}.json").write_text(json.dumps(group, cls=ArgusJSONEncoder))
+    (dest / f"group_{group.id}.json").write_text(
+        json.dumps(group, cls=ArgusJSONEncoder)
+    )
 
-print('getting test data')
+print("getting test data")
 for test in ArgusTest.filter(release_id=release.id).all():
     print(f"Saving test {test.id}")
     (dest / f"test_{test.id}.json").write_text(json.dumps(test, cls=ArgusJSONEncoder))
 
-print('getting runs data')
-total = db.session.execute(f"SELECT count(build_id) FROM sct_test_run WHERE release_id = {release.id}").one()["system.count(build_id)"]
+    # Export ArgusGenericResultMetadata for each test
+    for metadata in (
+        ArgusGenericResultMetadata.filter(test_id=test.id).allow_filtering().all()
+    ):
+        print(
+            f"Saving generic result metadata for test {test.id} and name {metadata.name}"
+        )
+        (dest / f"generic_result_metadata_{test.id}_{metadata.name}.json").write_text(
+            json.dumps(metadata, cls=ArgusJSONEncoder)
+        )
+
+    # Export ArgusGraphView for each test - group by test_id (partition key)
+    graph_views_by_test = defaultdict(list)
+    for graph_view in ArgusGraphView.filter(test_id=test.id).allow_filtering().all():
+        graph_views_by_test[str(test.id)].append(graph_view)
+
+    for test_id, views in graph_views_by_test.items():
+        print(f"Saving graph views for test {test_id}")
+        (dest / f"graph_view_{test_id}.json").write_text(
+            json.dumps(views, cls=ArgusJSONEncoder)
+        )
+
+    # Export ArgusBestResultData for each test - group by test_id and name (partition keys)
+    best_results_by_partition = defaultdict(list)
+    for best_result in (
+        ArgusBestResultData.filter(test_id=test.id).allow_filtering().all()
+    ):
+        partition_key = f"{test.id}_{best_result.name}"
+        best_results_by_partition[partition_key].append(best_result)
+
+    for partition_key, results in best_results_by_partition.items():
+        print(f"Saving best results for partition {partition_key}")
+        (dest / f"best_result_{partition_key}.json").write_text(
+            json.dumps(results, cls=ArgusJSONEncoder)
+        )
+
+print("getting runs data")
+total = db.session.execute(
+    f"SELECT count(build_id) FROM sct_test_run WHERE release_id = {release.id}"
+).one()["system.count(build_id)"]
 idx = 0
 for run in SCTTestRun.filter(release_id=release.id).all():
     idx += 1
     print(f"Saving {run.id} [{idx}/{total}]")
     (dest / f"run_{run.id}.json").write_text(json.dumps(run, cls=ArgusJSONEncoder))
+
+    # Export ArgusGenericResultData for each run - group by test_id and name (partition keys)
+    result_data_by_partition = defaultdict(list)
+    for result_data in (
+        ArgusGenericResultData.filter(run_id=run.id).allow_filtering().all()
+    ):
+        partition_key = f"{run.id}_{result_data.name}"
+        result_data_by_partition[partition_key].append(result_data)
+
+    for partition_key, results in result_data_by_partition.items():
+        print(f"Saving generic result data for partition {partition_key}")
+        (dest / f"generic_result_data_{partition_key}.json").write_text(
+            json.dumps(results, cls=ArgusJSONEncoder)
+        )

dev-db/import_data.py

@@ -3,21 +3,30 @@
 
 from argus.backend.plugins.sct.testrun import SCTTestRun
 from argus.backend.db import ScyllaCluster
+from argus.backend.models.result import (
+    ArgusGenericResultMetadata,
+    ArgusGenericResultData,
+    ArgusBestResultData,
+    ArgusGraphView,
+)
+from argus.backend.util.config import Config
 
 db = ScyllaCluster.get()
+if "127.0.0.10" in Config.CONFIG.get("SCYLLA_CONTACT_POINTS"):
+    raise Exception("This script should not be run on local DB!")
 
-release_name = "perf-regression" # change name here
+release_name = "scylla-enterprise"  # change name here
 dest = Path(__file__).parent / "sample_data" / release_name
 
-print('importing releases')
+print("importing releases")
 statement = db.prepare(f"INSERT INTO argus_release_v2 JSON ?")
 with (dest / "release.json").open(mode="rt", encoding="utf-8") as src:
     release_raw = src.read()
     release = json.loads(release_raw)
     db.session.execute(statement, parameters=(release_raw,))
     print(f"Saved {release['id']}")
 
-print('importing groups')
+print("importing groups")
 statement = db.prepare("INSERT INTO argus_group_v2 JSON ?")
 for file in dest.glob("group_*.json"):
     with file.open(mode="rt", encoding="utf-8") as src:
@@ -26,7 +35,7 @@
         db.session.execute(statement, parameters=(group_raw,))
         print(f"Saved {group['id']}")
 
-print('importing tests')
+print("importing tests")
 statement = db.prepare("INSERT INTO argus_test_v2 JSON ?")
 for file in dest.glob("test_*.json"):
     with file.open(mode="rt", encoding="utf-8") as src:
@@ -46,3 +55,50 @@
         run.assign_categories()
         run.save()
         print(f"Saved {run.id}")
+
+print("importing generic result metadata")
+statement = db.prepare("INSERT INTO generic_result_metadata_v1 JSON ?")
+for file in dest.glob("generic_result_metadata_*.json"):
+    with file.open(mode="rt", encoding="utf-8") as src:
+        metadata_raw = src.read()
+        metadata = json.loads(metadata_raw)
+        db.session.execute(statement, parameters=(metadata_raw,))
+        print(
+            f"Saved metadata for test {metadata['test_id']} and name {metadata['name']}"
+        )
+
+print("importing graph views")
+statement = db.prepare("INSERT INTO graph_view_v1 JSON ?")
+for file in dest.glob("graph_view_*.json"):
+    with file.open(mode="rt", encoding="utf-8") as src:
+        graph_views = json.loads(src.read())
+        for graph_view in graph_views:
+            graph_view_raw = json.dumps(graph_view)
+            db.session.execute(statement, parameters=(graph_view_raw,))
+            print(
+                f"Saved graph view {graph_view['id']} for test {graph_view['test_id']}"
+            )
+
+print("importing best results")
+statement = db.prepare("INSERT INTO generic_result_best_v2 JSON ?")
+for file in dest.glob("best_result_*.json"):
+    with file.open(mode="rt", encoding="utf-8") as src:
+        best_results = json.loads(src.read())
+        for best_result in best_results:
+            best_result_raw = json.dumps(best_result)
+            db.session.execute(statement, parameters=(best_result_raw,))
+            print(
+                f"Saved best result for test {best_result['test_id']} and name {best_result['name']}"
+            )
+
+print("importing generic result data")
+statement = db.prepare("INSERT INTO generic_result_data_v1 JSON ?")
+for file in dest.glob("generic_result_data_*.json"):
+    with file.open(mode="rt", encoding="utf-8") as src:
+        results = json.loads(src.read())
+        for result_data in results:
+            result_data_raw = json.dumps(result_data)
+            db.session.execute(statement, parameters=(result_data_raw,))
+            print(
+                f"Saved generic result data for run {result_data['run_id']}, test {result_data['test_id']}, name {result_data['name']}"
+            )