From 4f28fcfcff3b895430228d7fecaf821dbc42a5af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Sun, 30 Jun 2019 13:26:24 +0200
Subject: [PATCH 2/2] ctl: add --all-tenants parameter to runscript (#34405)

---
 tests/test_ctl.py                        | 20 ++++++++++++++++++++
 wcs/ctl/management/commands/runscript.py | 23 ++++++++++++++++++-----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/tests/test_ctl.py b/tests/test_ctl.py
index 796a4bb43..d6b4298af 100644
--- a/tests/test_ctl.py
+++ b/tests/test_ctl.py
@@ -5,6 +5,8 @@ from email.mime.text import MIMEText
 from email.mime.multipart import MIMEMultipart
 import psycopg2
 
+from django.core.management import call_command, CommandError
+
 from wcs.formdef import FormDef
 from wcs.workflows import Workflow
 from wcs.wf.jump import JumpWorkflowStatusItem
@@ -17,6 +19,7 @@ from wcs.qommon.management.commands.makemessages import Command as CmdMakeMessag
 from wcs.ctl.process_bounce import CmdProcessBounce
 from wcs.ctl.rebuild_indexes import rebuild_vhost_indexes
 from wcs.ctl.wipe_data import CmdWipeData
+from wcs.ctl.management.commands.runscript import Command as CmdRunScript
 from wcs.ctl.management.commands.trigger_jumps import select_and_jump_formdata
 from wcs.ctl.delete_tenant import CmdDeleteTenant
 from wcs.sql import get_connection_and_cursor, cleanup_connection
@@ -361,3 +364,20 @@ def test_rebuild_indexes(two_pubs):
 
     rebuild_vhost_indexes(two_pubs, destroy=True)
     assert os.listdir(os.path.join(two_pubs.app_dir, 'formdefs-url_name')) == ['example']
+
+def test_runscript(pub):
+    with pytest.raises(CommandError):
+        call_command('runscript')
+    with pytest.raises(CommandError):
+        call_command('runscript', '--domain=a', '--all-tenants')
+    with open(os.path.join(pub.app_dir, 'test2.py'), 'w') as fd:
+        fd.write('''
+import os
+from quixote import get_publisher
+open(os.path.join(get_publisher().app_dir, 'runscript.test'), 'w').close()
+''')
+    call_command('runscript', '--domain=example.net', os.path.join(pub.app_dir, 'test2.py'))
+    assert os.path.exists(os.path.join(pub.app_dir, 'runscript.test'))
+    os.unlink(os.path.join(pub.app_dir, 'runscript.test'))
+    call_command('runscript', '--all-tenants', os.path.join(pub.app_dir, 'test2.py'))
+    assert os.path.exists(os.path.join(pub.app_dir, 'runscript.test'))
diff --git a/wcs/ctl/management/commands/runscript.py b/wcs/ctl/management/commands/runscript.py
index 51fc069d7..b9364c0fc 100644
--- a/wcs/ctl/management/commands/runscript.py
+++ b/wcs/ctl/management/commands/runscript.py
@@ -19,6 +19,8 @@ import os.path
 import runpy
 import sys
 
+from django.core.management.base import CommandError
+
 from qommon.publisher import get_publisher_class
 
 from . import TenantCommand
@@ -28,17 +30,28 @@ class Command(TenantCommand):
     '''Run a script within a given host publisher context'''
 
     def add_arguments(self, parser):
-        parser.add_argument('-d', '--domain', '--vhost', metavar='DOMAIN', required=True)
+        parser.add_argument('-d', '--domain', '--vhost', metavar='DOMAIN')
         parser.add_argument('--app-dir', metavar='DIR', action='store', dest='app_dir', default=None)
+        parser.add_argument('--all-tenants', action='store_true')
         parser.add_argument('args', nargs=argparse.REMAINDER)
 
     def handle(self, *args, **options):
         if options.get('app_dir'):
             get_publisher_class().APP_DIR = options.get('app_dir')
         domain = options.get('domain')
-        self.init_tenant_publisher(domain, register_tld_names=False)
+        all_tenants = options.get('all_tenants')
+        if domain and all_tenants:
+            raise CommandError('--domain and --all-tenants are exclusive')
+        if not (domain or all_tenants):
+            raise CommandError('--domain or --all-tenants are required')
+        if domain:
+            domains = [domain]
+        else:
+            domains = get_publisher_class().get_tenants()
         fullpath = os.path.dirname(os.path.abspath(args[0]))
         sys.path.insert(0, fullpath)
-        module_name = os.path.splitext(os.path.basename(args[0]))[0]
-        sys.argv = args
-        runpy.run_module(module_name)
+        module_name = os.path.splitext(os.path.basename(args[0]))[0].encode('utf-8')
+        for domain in domains:
+            sys.argv = args[:]
+            self.init_tenant_publisher(domain, register_tld_names=False)
+            runpy.run_module(module_name)
-- 
2.20.1

