diff --git a/treeherder/model/management/commands/run_sql.py b/treeherder/model/management/commands/run_sql.py new file mode 100644 index 000000000..0860f13da --- /dev/null +++ b/treeherder/model/management/commands/run_sql.py @@ -0,0 +1,78 @@ +import MySQLdb +from optparse import make_option + +from django.core.management.base import BaseCommand + +from treeherder.model.models import Datasource +from django.conf import settings + + +class Command(BaseCommand): + help = ("Runs an arbitrary sql statement or file" + " on a number of databases.") + + option_list = BaseCommand.option_list + ( + + make_option( + '--datasources', + action='store', + dest='datasources', + default='all', + help='A comma separated list of datasources to execute the sql code on'), + + make_option( + '--data-type', + action='store', + dest='data_type', + default='jobs', + choices=['jobs', 'objectstore'], + help='The target data-type of the sql code'), + + make_option( + '-f', '--file', + dest='sql_file', + help='Sql source file', + metavar='FILE', + default="") + + ) + + def handle(self, *args, **options): + + if not options["sql_file"]: + self.stderr.write("No sql file provided!") + return + + datasources = Datasource.objects.filter(contenttype=options['data_type']) + if options['datasources'] != 'all': + if ',' in options['datasources']: + datasources = datasources.filter( + project__in=options['datasources'].split(',')) + else: + datasources = datasources.filter( + project=options['datasources']) + + with open(options["sql_file"]) as sql_file: + sql_code = sql_file.read() + + self.stdout.write("{0} datasource found".format( + len(datasources) + )) + for datasource in datasources: + self.stdout.write("--------------------------") + db = MySQLdb.connect( + host=datasource.host, + db=datasource.name, + user=settings.TREEHERDER_DATABASE_USER, + passwd=settings.TREEHERDER_DATABASE_PASSWORD) + try: + cursor = db.cursor() + cursor.execute(sql_code) + self.stdout.write("Sql code executed on {0}".format(datasource)) + except Exception as e: + error_string = "!!! Sql code execution failed on {0} !!!" + self.stderr.write(error_string.format(datasource)) + self.stderr.write("{0}".format(e)) + finally: + if cursor: + cursor.close()