From ea7abae220ea1beb7a3ad4b7f26866b3401431f5 Mon Sep 17 00:00:00 2001
From: Tanmoy Sarkar <57363826+tanmoysrt@users.noreply.github.com>
Date: Fri, 20 Dec 2024 04:59:44 +0000
Subject: [PATCH] chore(database): use root user for db processlist and kill
 process

---
 agent/database.py | 16 +++++++++++++++-
 agent/site.py     | 13 ++++---------
 agent/web.py      | 19 ++++++++++++++-----
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/agent/database.py b/agent/database.py
index 36835845..c8596227 100644
--- a/agent/database.py
+++ b/agent/database.py
@@ -385,7 +385,7 @@ def fetch_process_list(self):
             return []
         return [
             {
-                "id": record["Id"],
+                "id": str(record["Id"]),
                 "command": record["Command"],
                 "query": record["Info"],
                 "state": record["State"],
@@ -394,10 +394,24 @@ def fetch_process_list(self):
                 "db_user_host": record["Host"].split(":")[0],
             }
             for record in result[0]["output"]
+            if record["db"] == self.database_name
         ]
 
     def kill_process(self, pid: str):
         with contextlib.suppress(Exception):
+            processes = self.fetch_process_list()
+            """
+            It's important to validate whether the pid belongs to the current database
+            As we are running it as `root` user, it can be possible to kill processes from other databases
+            by forging the request
+            """
+            isFound = False
+            for process in processes:
+                if process.get("id") == pid:
+                    isFound = True
+                    break
+            if not isFound:
+                return
             """
             The query can fail if the process is already killed.
             Anyway we need to reload the process list after killing to verify if the process is killed.
diff --git a/agent/site.py b/agent/site.py
index 392e2b8a..6731cf24 100644
--- a/agent/site.py
+++ b/agent/site.py
@@ -887,11 +887,6 @@ def run_sql_query(self, query: str, commit: bool = False, as_dict: bool = False)
             response["failed_query"] = db.last_executed_query
         return response
 
-    @job("Analyze Slow Queries")
-    def analyze_slow_queries_job(self, queries: list[dict], database_root_password: str):
-        return self.analyze_slow_queries(queries, database_root_password)
-
-    @step("Analyze Slow Queries")
     def analyze_slow_queries(self, queries: list[dict], database_root_password: str) -> list[dict]:
         """
         Args:
@@ -918,11 +913,11 @@ def fetch_summarized_database_performance_report(self, mariadb_root_password: st
         database = self.db_instance(username="root", password=mariadb_root_password)
         return database.fetch_summarized_performance_report()
 
-    def fetch_database_process_list(self):
-        return self.db_instance().fetch_process_list()
+    def fetch_database_process_list(self, mariadb_root_password: str):
+        return self.db_instance(username="root", password=mariadb_root_password).fetch_process_list()
 
-    def kill_database_process(self, pid: str):
-        return self.db_instance().kill_process(pid)
+    def kill_database_process(self, pid: str, mariadb_root_password: str):
+        return self.db_instance(username="root", password=mariadb_root_password).kill_process(pid)
 
     def db_instance(self, username: str | None = None, password: str | None = None) -> Database:
         if not username:
diff --git a/agent/web.py b/agent/web.py
index 0de8de31..c1c428a0 100644
--- a/agent/web.py
+++ b/agent/web.py
@@ -591,15 +591,20 @@ def run_sql(bench, site):
 
 
 @application.route(
-    "/benches/<string:bench>/sites/<string:site>/database/analyze-slow-queries", methods=["GET"]
+    "/benches/<string:bench>/sites/<string:site>/database/analyze-slow-queries", methods=["POST"]
 )
 @validate_bench_and_site
 def analyze_slow_queries(bench: str, site: str):
     queries = request.json["queries"]
     mariadb_root_password = request.json["mariadb_root_password"]
 
-    job = Server().benches[bench].sites[site].analyze_slow_queries_job(queries, mariadb_root_password)
-    return {"job": job}
+    return Response(
+        json.dumps(
+            Server().benches[bench].sites[site].analyze_slow_queries(queries, mariadb_root_password),
+            cls=JSONEncoderForSQLQueryResult,
+        ),
+        mimetype="application/json",
+    )
 
 
 @application.route(
@@ -618,14 +623,18 @@ def database_performance_report(bench, site):
 
 @application.route("/benches/<string:bench>/sites/<string:site>/database/processes", methods=["GET", "POST"])
 def database_process_list(bench, site):
-    return jsonify(Server().benches[bench].sites[site].fetch_database_process_list())
+    data = request.json
+    return jsonify(
+        Server().benches[bench].sites[site].fetch_database_process_list(data["mariadb_root_password"])
+    )
 
 
 @application.route(
     "/benches/<string:bench>/sites/<string:site>/database/kill-process/<string:pid>", methods=["GET", "POST"]
 )
 def database_kill_process(bench, site, pid):
-    Server().benches[bench].sites[site].kill_database_process(pid)
+    data = request.json
+    Server().benches[bench].sites[site].kill_database_process(pid, data["mariadb_root_password"])
     return "killed"