diff --git a/docs/common_tasks.rst b/docs/common_tasks.rst index 87118b9af..58be85d6c 100644 --- a/docs/common_tasks.rst +++ b/docs/common_tasks.rst @@ -97,11 +97,12 @@ Executing arbitrary SQL As part of a larger change, you may want to execute some arbitrary SQL on the server. You can do this with the `run_sql` management command. -Example: +Examples: .. code-block:: bash - > ./manage.py run_sql -f + > ./manage.py run_sql -s + > ./manage.py run_sql -f By default, this will run the sql against the `jobs` database for each project. If you want to run against the object store or only against a diff --git a/treeherder/model/management/commands/run_sql.py b/treeherder/model/management/commands/run_sql.py index dab5cfffe..1260fa7d6 100644 --- a/treeherder/model/management/commands/run_sql.py +++ b/treeherder/model/management/commands/run_sql.py @@ -33,6 +33,13 @@ class Command(BaseCommand): help=('The target data-type of the sql code (jobs or objectstore, ' 'default jobs)')), + make_option( + '-s', '--sql-statement', + action='store', + dest='sql_statement', + help='Sql statement', + default=''), + make_option( '-f', '--file', dest='sql_file', @@ -44,9 +51,17 @@ class Command(BaseCommand): def handle(self, *args, **options): - if not options["sql_file"]: - self.stderr.write("No sql file provided!") - return + sql_code = options["sql_statement"] + + if not sql_code: + if options["sql_file"]: + with open(options["sql_file"]) as sql_file: + sql_code = sql_file.read() + else: + self.stderr.write("Either a SQL statement or file must be specified! See --help.") + return + + self.stdout.write("SQL command: {}".format(sql_code)) datasources = Datasource.objects.filter(contenttype=options['data_type']) if options['datasources'] != 'all': @@ -57,27 +72,25 @@ class Command(BaseCommand): 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) + )) - 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() + 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()