multitenant: show ETA on migrate_schemas (#68034)

hobo/multitenant/management/commands/__init__.py

@@ -188,9 +188,11 @@ class SyncCommon(BaseCommand):
         if hasattr(settings, 'SHARED_APPS'):
             self.shared_apps = settings.SHARED_APPS
 
-    def _notice(self, output):
+    def _notice(self, output, flush=False):
         if int(self.options.get('verbosity', 1)) >= 1:
             self.stdout.write(self.style.NOTICE(output))
+            if flush:
+                self.stdout.flush()
 
 
 def disable_global_logging():

hobo/multitenant/management/commands/migrate_schemas.py

@@ -21,6 +21,7 @@ from django.core.management.commands.migrate import Command as MigrateCommand
 from django.db import connection
 from django.db.migrations.loader import MigrationLoader
 from django.db.migrations.recorder import MigrationRecorder
+from django.utils.timezone import localtime
 from tenant_schemas.postgresql_backend.base import FakeTenant
 from tenant_schemas.utils import get_public_schema_name, schema_exists
 
@@ -59,7 +60,10 @@ class MigrateSchemasCommand(SyncCommon):
             all_migrations = {
                 (app, migration) for app, migration in loader.disk_migrations if app in app_labels
             }
-            for tenant in TenantMiddleware.get_tenants():
+            tenants = list(TenantMiddleware.get_tenants())
+            len_tenants = len(tenants)
+            start_datetime = localtime()
+            for step, tenant in enumerate(tenants, start=1):
                 connection.set_tenant(tenant, include_public=False)
                 applied_migrations = self.get_applied_migrations(app_labels)
                 if options.get('fake') or options.get('migration_name') or options.get('app_label'):
@@ -68,9 +72,15 @@ class MigrateSchemasCommand(SyncCommon):
                     applied_migrations = []
                 if all([x in applied_migrations for x in all_migrations]):
                     if int(self.options.get('verbosity', 1)) >= 1:
-                        self._notice("=== Skipping migrations of tenant %s" % tenant.domain_url)
+                        self._notice(
+                            "=== Skipping migrations of tenant %s (%s/%s)"
+                            % (tenant.domain_url, step, len_tenants)
+                        )
                     continue
-                self.run_migrations(tenant, settings.TENANT_APPS)
+                self.run_migrations(tenant, settings.TENANT_APPS, step, len_tenants)
+                if int(self.options.get('verbosity', 1)) >= 1:
+                    eta = start_datetime + len_tenants * (localtime() - start_datetime) / step
+                    self._notice('=== migrate_schemas ETA: %s' % eta, flush=True)
 
     def get_applied_migrations(self, app_labels):
         applied_migrations = []
@@ -81,9 +91,9 @@ class MigrateSchemasCommand(SyncCommon):
         applied_migrations = [x for x in applied_migrations if x[0] in app_labels]
         return applied_migrations
 
-    def run_migrations(self, tenant, included_apps):
+    def run_migrations(self, tenant, included_apps, step=1, steps=1):
         if int(self.options.get('verbosity', 1)) >= 1:
-            self._notice("=== Running migrate for tenant %s" % tenant.domain_url)
+            self._notice("=== Running migrate for tenant %s (%s/%s)" % (tenant.domain_url, step, steps))
         connection.set_tenant(tenant, include_public=False)
         command = MigrateCommand()
         command.requires_system_checks = False
@@ -101,8 +111,10 @@ class MigrateSchemasCommand(SyncCommon):
         command.execute(*self.args, **self.options)
         connection.set_schema_to_public()
 
-    def _notice(self, output):
+    def _notice(self, output, flush=False):
         self.stdout.write(self.style.NOTICE(output))
+        if flush:
+            self.stdout.flush()
 
 
 Command = MigrateSchemasCommand

tests_multitenant/test_create_tenant.py

@@ -142,3 +142,21 @@ def test_create_tenant_from_existing_tenant_but_already_exists(db):
     with pytest.raises(CommandError) as exc_info:
         call_command('create_tenant', '--legacy-hostname', 'host1.com', 'host2.com')
     assert str(exc_info.value) == 'tenant already exists'
+
+
+def test_migrate_schemas_eta(db, capsys):
+    call_command('create_tenant', 'host1.com')
+    call_command('create_tenant', 'host2.com')
+    # all skipped, no ETA is displayed
+    call_command('migrate_schemas', verbosity=1)
+    captured = capsys.readouterr()
+    assert 'Skipping migrations of tenant host1.com (1/2)' in captured.out
+    assert 'Skipping migrations of tenant host2.com (2/2)' in captured.out
+    assert 'migrate_schemas ETA: 2' not in captured.out
+    # force re-migration and re-migrate, ETA is displayed
+    call_command('migrate_schemas', 'common', '0001_initial', verbosity=1)
+    call_command('migrate_schemas', verbosity=1)
+    captured = capsys.readouterr()
+    assert 'Running migrate for tenant host1.com (1/2)' in captured.out
+    assert 'Running migrate for tenant host2.com (2/2)' in captured.out
+    assert 'migrate_schemas ETA: 2' in captured.out