Update query killer to work with multi-instance replicas

helm-quarry/dev-config.yaml

@@ -16,7 +16,6 @@ task_track_started: True
 worker_prefetch_multiplier: 1 # Tasks can run for a long time
 # Just query the quarry database itself.
 REPLICA_DOMAIN: ''
-REPLICA_HOST: 'mywiki'
 REPLICA_DB: 'mywiki_p'
 REPLICA_USER: 'repl'
 REPLICA_PASSWORD: 'repl'

helm-quarry/values.yaml

@@ -1,7 +1,7 @@
 web:
   repository: 'quay.io/wikimedia-quarry/quarry'
-  tag: pr-47 # web tag managed by github actions
+  tag: pr-48 # web tag managed by github actions
 
 worker:
   repository: 'quay.io/wikimedia-quarry/quarry'
-  tag: pr-47 # worker tag managed by github actions
+  tag: pr-48 # worker tag managed by github actions

quarry/default_config.yaml

@@ -31,6 +31,7 @@ REPLICA_DOMAIN: ''  # Change to `analytics.db.svc.wikimedia.cloud` for live repl
 REPLICA_USER: 'repl'  # For live replicas, your replica.my.cnf username
 REPLICA_PASSWORD: 'repl'  # For live replicas, your replica.my.cnf password
 REPLICA_PORT: 3306
+REPLICA_SLICES: ''
 
 TOOLS_DB_HOST: 'tools-readonly.db.svc.wikimedia.cloud'
 TOOLS_DB_PORT: 3306

quarry/web/killer.py

@@ -3,7 +3,7 @@
 import yaml
 import logging
 import pymysql
-from connections import Connections
+from quarry.web.replica import Replica
 
 __dir__ = os.path.dirname(__file__)
 config = yaml.safe_load(open(os.path.join(__dir__, "../default_config.yaml")))
@@ -21,31 +21,36 @@
 logging.info(
     "Started killer process, with limit %s", config["QUERY_TIME_LIMIT"]
 )
-conn = Connections(config)
 
-cur = conn.replica.cursor()
-try:
-    cur.execute("SHOW PROCESSLIST")
-    queries = cur.fetchall()
-    logging.info("Found %s queries running", len(queries))
-    to_kill = [
-        q
-        for q in queries
-        if q[5] > config["QUERY_TIME_LIMIT"] and q[4] != "Sleep"
-    ]
-    logging.info("Found %s queries to kill", len(to_kill))
-    for q in to_kill:
-        try:
-            cur.execute("KILL QUERY %s", q[0])
-            logging.info("Killed query with thread_id:%s" % q[0])
-        except pymysql.InternalError as e:
-            if e.args[0] == 1094:  # Error code for 'no such thread'
-                logging.info(
-                    "Query with thread_id:%s dead before it could be killed"
-                )
-            else:
-                raise
-finally:
-    logging.info("Finished killer process")
-    cur.close()
-    conn.close_all()
+hosts = [db_slice + '.' + config.get['REPLICA_DOMAIN'] for db_slice in config.get['REPLICA_SLICES'].split(',')]
+hosts.append(config.get['TOOLS_DB_HOST'])
+
+for host in hosts:
+    conn = Replica(config)
+
+    cur = conn.connection.cursor()
+    try:
+        cur.execute("SHOW PROCESSLIST")
+        queries = cur.fetchall()
+        logging.info("Found %s queries running", len(queries))
+        to_kill = [
+            q
+            for q in queries
+            if q[5] > config["QUERY_TIME_LIMIT"] and q[4] != "Sleep"
+        ]
+        logging.info("Found %s queries to kill", len(to_kill))
+        for q in to_kill:
+            try:
+                cur.execute("KILL QUERY %s", q[0])
+                logging.info("Killed query with thread_id:%s" % q[0])
+            except pymysql.InternalError as e:
+                if e.args[0] == 1094:  # Error code for 'no such thread'
+                    logging.info(
+                        "Query with thread_id:%s dead before it could be killed"
+                    )
+                else:
+                    raise
+    finally:
+        logging.info("Finished killer process")
+        cur.close()
+        del conn.connection

quarry/web/replica.py

@@ -10,9 +10,9 @@ class Replica:
     def __init__(self, config):
         self.config = config
         self.dbname = ""
-        self.is_tools_db = False
 
     def _db_name_mangler(self):
+        self.is_tools_db = False
         if self.dbname == "":
             raise ReplicaConnectionException(
                 "Attempting connection before a database is selected"
@@ -21,15 +21,16 @@ def _db_name_mangler(self):
             self.is_tools_db = True
             self.database_p = self.dbname
         elif self.dbname == "meta" or self.dbname == "meta_p":
-            self.is_tools_db = False
             self.database_name = "s7"
             self.database_p = "meta_p"
         elif self.dbname == "centralauth" or self.dbname == "centralauth_p":
-            self.is_tools_db = False
             self.database_name = "s7"
             self.database_p = "centralauth_p"
+        elif len(self.dbname) == 2:
+            # slice, eg. s1, s2
+            self.database_name = self.dbname
+            self.database_p = self.dbname
         else:
-            self.is_tools_db = False
             self.database_name = (
                 self.dbname
                 if not self.dbname.endswith("_p")