add run_sql shell command

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()