From 34cd4f05a64c2de1a056547bc427588c121a3065 Mon Sep 17 00:00:00 2001
From: Emmanuel Cazenave <ecazenave@entrouvert.com>
Date: Mon, 29 Nov 2021 17:32:16 +0100
Subject: [PATCH] multitenant: add --legacy-hostname option to create_tenant
 (#59094)

---
 .../management/commands/create_tenant.py      | 54 ++++++++++++++-----
 hobo/multitenant/models.py                    | 16 ++++++
 tests_multitenant/test_create_tenant.py       | 47 +++++++++++++++-
 3 files changed, 103 insertions(+), 14 deletions(-)

diff --git a/hobo/multitenant/management/commands/create_tenant.py b/hobo/multitenant/management/commands/create_tenant.py
index e8ed36f..aaae8d4 100644
--- a/hobo/multitenant/management/commands/create_tenant.py
+++ b/hobo/multitenant/management/commands/create_tenant.py
@@ -12,8 +12,9 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         parser.add_argument('hostnames', metavar='HOSTNAME', nargs='+')
+        parser.add_argument('--legacy-hostname', dest='legacy_hostname')
 
-    def handle(self, hostnames, **options):
+    def handle(self, hostnames, legacy_hostname, **options):
         verbosity = int(options.get('verbosity'))
         if not hostnames:
             raise CommandError("you must give at least one tenant hostname")
@@ -23,6 +24,9 @@ class Command(BaseCommand):
             hostnames.remove('-')
             hostnames.extend([x.strip() for x in sys.stdin.readlines()])
 
+        if legacy_hostname and len(hostnames) > 1:
+            raise CommandError("You must specify only hostname when using --legacy-hostname")
+
         for hostname in hostnames:
             try:
                 tenant_base = TenantMiddleware.base()
@@ -33,28 +37,52 @@ class Command(BaseCommand):
             tenant_dir = os.path.join(tenant_base, hostname)
             if os.path.exists(tenant_dir):
                 raise CommandError('tenant already exists')
+
             tenant_dir_tmp = os.path.join(tenant_base, hostname + '__CREATION_IN_PROGRESS.invalid')
-            try:
-                os.mkdir(tenant_dir_tmp, 0o755)
-            except OSError as e:
-                raise CommandError('cannot start tenant creation (%s)' % str(e))
+            if legacy_hostname:
+                legacy_tenant_dir = os.path.join(tenant_base, legacy_hostname)
+                if not os.path.exists(legacy_tenant_dir):
+                    raise CommandError('legacy tenant does not exists')
+                try:
+                    os.rename(legacy_tenant_dir, tenant_dir_tmp)
+                except OSError as e:
+                    raise CommandError('cannot change legacy tenant (%s)' % str(e))
+            else:
+                try:
+                    os.mkdir(tenant_dir_tmp, 0o755)
+                except OSError as e:
+                    raise CommandError('cannot start tenant creation (%s)' % str(e))
+
             # tenant creation in database
             try:
                 connection.set_schema_to_public()
                 schema = TenantMiddleware.hostname2schema(hostname)
+                legacy_schema = None
+                if legacy_hostname:
+                    legacy_schema = TenantMiddleware.hostname2schema(legacy_hostname)
                 tenant = get_tenant_model()(schema_name=schema, domain_url=hostname)
                 if verbosity >= 1:
+                    msg = self.style.NOTICE("=== Creating schema ") + self.style.SQL_TABLE(tenant.schema_name)
+                    if legacy_schema:
+                        msg = (
+                            self.style.NOTICE("=== Renamin schema ")
+                            + self.style.SQL_TABLE(legacy_schema)
+                            + " to "
+                            + self.style.SQL_TABLE(tenant.schema_name)
+                        )
                     print()
-                    print(
-                        self.style.NOTICE("=== Creating schema ") + self.style.SQL_TABLE(tenant.schema_name)
-                    )
+                    print(msg)
                 with transaction.atomic():
-                    tenant.create_schema(check_if_exists=True)
+                    tenant.create_schema(check_if_exists=True, legacy_schema_name=legacy_schema)
             except Exception as e:
-                os.rmdir(tenant_dir_tmp)
+                if legacy_hostname:
+                    os.rename(tenant_dir_tmp, legacy_tenant_dir)
+                else:
+                    os.rmdir(tenant_dir_tmp)
                 raise CommandError('tenant creation failed (%s)' % str(e))
-            for folder in ('media', 'static', 'templates'):
-                path = os.path.join(tenant_dir_tmp, folder)
-                os.mkdir(path, 0o755)
+            if not legacy_hostname:
+                for folder in ('media', 'static', 'templates'):
+                    path = os.path.join(tenant_dir_tmp, folder)
+                    os.mkdir(path, 0o755)
             # activate the tenant by creating its directory
             os.rename(tenant_dir_tmp, tenant_dir)
diff --git a/hobo/multitenant/models.py b/hobo/multitenant/models.py
index a682a4a..8ac561c 100644
--- a/hobo/multitenant/models.py
+++ b/hobo/multitenant/models.py
@@ -7,6 +7,7 @@ from django.db import connection
 from django.utils import timezone
 from django.utils.six.moves.urllib.parse import urljoin
 from tenant_schemas.models import TenantMixin
+from tenant_schemas.postgresql_backend.base import _check_schema_name
 from tenant_schemas.utils import django_is_in_test_mode, get_public_schema_name, schema_exists
 
 
@@ -57,6 +58,21 @@ class Tenant(TenantMixin):
     def build_absolute_uri(self, location):
         return urljoin(self.get_base_url(), location)
 
+    def create_schema(self, check_if_exists=False, sync_schema=True, verbosity=1, legacy_schema_name=None):
+        if not legacy_schema_name:
+            return super().create_schema(check_if_exists, sync_schema, verbosity)
+        # safety check
+        _check_schema_name(self.schema_name)
+        _check_schema_name(legacy_schema_name)
+        if check_if_exists and schema_exists(self.schema_name):
+            return False
+        if not schema_exists(legacy_schema_name):
+            return False
+        # rename schema
+        cursor = connection.cursor()
+        cursor.execute('ALTER SCHEMA %s RENAME TO %s' % (legacy_schema_name, self.schema_name))
+        connection.set_schema_to_public()
+
     def delete(self, force_drop=False, *args, **kwargs):
         """
         Deletes this row. Drops the tenant's schema if the attribute
diff --git a/tests_multitenant/test_create_tenant.py b/tests_multitenant/test_create_tenant.py
index 7170a5d..88bfe22 100644
--- a/tests_multitenant/test_create_tenant.py
+++ b/tests_multitenant/test_create_tenant.py
@@ -33,7 +33,7 @@ def configuration(settings, tmpdir):
 def schema_exists(schema_name):
     with connection.cursor() as cursor:
         cursor.execute('select schema_name from information_schema.schemata')
-        return 'www_example_com' in [row[0] for row in cursor.fetchall()]
+        return schema_name in [row[0] for row in cursor.fetchall()]
 
 
 def test_create_tenant(db):
@@ -84,3 +84,48 @@ def test_migrate_schemas_skip_applied(db, capsys):
     call_command('migrate_schemas', verbosity=1)
     captured = capsys.readouterr()
     assert 'Skipping migrations of tenant www.example.com' in captured.out
+
+
+def test_create_tenant_from_existing_tenant(db):
+    assert not schema_exists('www_example_com')
+    call_command('create_tenant', 'www.example.com')
+    assert schema_exists('www_example_com')
+    tenants = list(TenantMiddleware.get_tenants())
+    assert len(tenants) == 1
+    tenant = tenants[0]
+    assert tenant.domain_url == 'www.example.com'
+    assert tenant.schema_name == 'www_example_com'
+    with tenant_context(tenant):
+        User.objects.create(username='admin')
+
+    assert not schema_exists('www_newexample_com')
+    call_command('create_tenant', '--legacy-hostname', 'www.example.com', 'www.newexample.com')
+    assert not schema_exists('www_example_com')
+    assert schema_exists('www_newexample_com')
+    tenants = list(TenantMiddleware.get_tenants())
+    assert len(tenants) == 1
+    tenant = tenants[0]
+    assert tenant.domain_url == 'www.newexample.com'
+    assert tenant.schema_name == 'www_newexample_com'
+
+
+def test_create_tenant_from_existing_tenant_error_too_many_hosts(db):
+    assert not schema_exists('host1_com')
+    assert not schema_exists('host2_com')
+    assert not schema_exists('host3_com')
+    with pytest.raises(CommandError) as exc_info:
+        call_command('create_tenant', '--legacy-hostname', 'host1.com', 'host2.com', 'host3.com')
+    assert str(exc_info.value) == 'You must specify only hostname when using --legacy-hostname'
+    assert not schema_exists('host1_com')
+    assert not schema_exists('host2_com')
+    assert not schema_exists('host3_com')
+
+
+def test_create_tenant_from_existing_tenant_error_does_not_exists(db):
+    assert not schema_exists('host1_com')
+    assert not schema_exists('host2_com')
+    with pytest.raises(CommandError) as exc_info:
+        call_command('create_tenant', '--legacy-hostname', 'host1.com', 'host2.com')
+    assert str(exc_info.value) == 'legacy tenant does not exists'
+    assert not schema_exists('host1_com')
+    assert not schema_exists('host2_com')
-- 
2.30.2

