[autotest] Add an option for label_cleaner

When -s option is enabled, label_cleaner checks if the server
where the tool is running on is 'primary' server.
If not, it just returns.

BUG=chromium:380365
TEST=Run the tool with/without -s option in local machine.
DEPLOY=apache

Change-Id: I0eb2c4463f94e51a96c81d3c02baae41fca025cd
Reviewed-on: https://chromium-review.googlesource.com/294878
Reviewed-by: Mungyung Ryu <mkryu@google.com>
Commit-Queue: Mungyung Ryu <mkryu@google.com>
Tested-by: Mungyung Ryu <mkryu@google.com>
diff --git a/frontend/afe/site_rpc_interface.py b/frontend/afe/site_rpc_interface.py
index 97ef017..a7cb19a 100644
--- a/frontend/afe/site_rpc_interface.py
+++ b/frontend/afe/site_rpc_interface.py
@@ -522,9 +522,10 @@
     shard.delete()
 
 
-def get_servers(role=None, status=None):
+def get_servers(hostname=None, role=None, status=None):
     """Get a list of servers with matching role and status.
 
+    @param hostname: FQDN of the server.
     @param role: Name of the server role, e.g., drone, scheduler. Default to
                  None to match any role.
     @param status: Status of the server, e.g., primary, backup, repair_required.
@@ -536,7 +537,7 @@
     if not server_manager_utils.use_server_db():
         raise error.RPCException('Server database is not enabled. Please try '
                                  'retrieve servers from global config.')
-    servers = server_manager_utils.get_servers(hostname=None, role=role,
+    servers = server_manager_utils.get_servers(hostname=hostname, role=role,
                                                status=status)
     return [s.get_details() for s in servers]
 
diff --git a/site_utils/label_cleaner.py b/site_utils/label_cleaner.py
index 921394f..991ebf9 100755
--- a/site_utils/label_cleaner.py
+++ b/site_utils/label_cleaner.py
@@ -5,7 +5,7 @@
 # found in the LICENSE file.
 
 
-"""Script to clean up labels that are not in use.
+"""Tool for cleaning up labels that are not in use.
 
 Delete given labels from database when they are not in use.
 Labels that match the query `SELECT_USED_LABELS_FORMAT` are considered in use.
@@ -24,12 +24,14 @@
 import argparse
 import logging
 import MySQLdb
+import socket
 import sys
 import traceback
 
 import common
 from autotest_lib.client.common_lib import global_config
 from autotest_lib.client.common_lib import logging_config
+from autotest_lib.server import frontend
 
 
 DB_SERVER = global_config.global_config.get_config_value('AUTOTEST_WEB', 'host')
@@ -130,6 +132,17 @@
         _delete_labels(conn, labels_to_del)
 
 
+def is_primary_server():
+    """Check if this server's status is primary
+
+    @return: True if primary, False otherwise.
+    """
+    server = frontend.AFE().run('get_servers', hostname=socket.getfqdn())
+    if server and server['status'] == 'primary':
+        return True
+    return False
+
+
 def main():
     parser = argparse.ArgumentParser(
             formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -145,6 +158,8 @@
     parser.add_argument('-n', dest='max_delete', type=int,
            help=('Max number of records to delete in each query.'),
            default=100)
+    parser.add_argument('-s', dest='check_status', action='store_true',
+           help=('Enforce to run only in a server that has primary status'))
     parser.add_argument('label', help='Label name to delete')
     options = parser.parse_args()
 
@@ -152,6 +167,10 @@
             datefmt='%Y-%m-%d %H:%M:%S')
 
     try:
+        if options.check_status and not is_primary_server():
+            logging.error('Cannot run in a non-primary server.')
+            return 1
+
         conn = MySQLdb.connect(host=options.db_server, user=USER,
                                passwd=PASSWD, db=DATABASE)
         used_labels = get_used_labels(conn)