From ad1a8cb6836cde005563d40450b85c824eddea4e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20NO=C3=8BL?= <tnoel@entrouvert.com>
Date: Mon, 8 Aug 2022 17:50:49 +0200
Subject: [PATCH] multitenant: show ETA on migrate_schemas (#68034)

---
 .../management/commands/migrate_schemas.py    | 11 +++++++++-
 tests_multitenant/test_create_tenant.py       | 20 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/hobo/multitenant/management/commands/migrate_schemas.py b/hobo/multitenant/management/commands/migrate_schemas.py
index 9c34f54..fe9f572 100644
--- a/hobo/multitenant/management/commands/migrate_schemas.py
+++ b/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,12 @@ class MigrateSchemasCommand(SyncCommon):
             all_migrations = set(
                 [(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()
+            if int(self.options.get('verbosity', 1)) >= 1 and len(tenants) > 1:
+                self._notice('=== Start migrate_schemas for %d tenants' % len(tenants))
+            for step, tenant in enumerate(tenants):
                 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'):
@@ -71,6 +77,9 @@ class MigrateSchemasCommand(SyncCommon):
                         self._notice("=== Skipping migrations of tenant %s" % tenant.domain_url)
                     continue
                 self.run_migrations(tenant, settings.TENANT_APPS)
+                if int(self.options.get('verbosity', 1)) >= 1:
+                    eta = start_datetime + len_tenants * (localtime() - start_datetime) / (step + 1)
+                    self._notice('=== migrate_schemas ETA: %s' % eta)
 
     def get_applied_migrations(self, app_labels):
         applied_migrations = []
diff --git a/tests_multitenant/test_create_tenant.py b/tests_multitenant/test_create_tenant.py
index 9bc99c2..8ae9f7b 100644
--- a/tests_multitenant/test_create_tenant.py
+++ b/tests_multitenant/test_create_tenant.py
@@ -141,3 +141,23 @@ 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 'Start migrate_schemas for 2 tenants' in captured.out
+    assert 'Skipping migrations of tenant host1.com' in captured.out
+    assert 'Skipping migrations of tenant host2.com' 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 'Start migrate_schemas for 2 tenants' in captured.out
+    assert 'Running migrate for tenant host1.com' in captured.out
+    assert 'Running migrate for tenant host2.com' in captured.out
+    assert 'migrate_schemas ETA: 2' in captured.out
-- 
2.35.1

