Refactor tunnel transport migrations to include schema_editor parameter

server/src/uds/migrations/0044_notification_notifier_servicetokenalias_and_more.py

@@ -16,19 +16,15 @@ def remove_null_service_pools(apps: typing.Any, schema_editor: typing.Any):  # p
     ServicePool.objects.filter(service__isnull=True).delete()
 
 
-# No-Op backwards migration
-def nop(apps: typing.Any, schema_editor: typing.Any):  # pylint: disable=unused-argument
-    pass
-
-
 # Python update network fields to allow ipv6
 # We will
 def update_network_model(apps: typing.Any, schema_editor: typing.Any):  # pylint: disable=unused-argument
     import uds.models.network  # pylint: disable=import-outside-toplevel,redefined-outer-name
 
     Network = apps.get_model('uds', 'Network')
+    db_alias = schema_editor.connection.alias
     try:
-        for net in Network.objects.all():
+        for net in Network.objects.using(db_alias).all():
             # Store the net_start and net_end on new fields "start" and "end", that are strings
             # to allow us to store ipv6 addresses
             # pylint: disable=protected-access
@@ -50,7 +46,8 @@ def update_transport_net_filtering(
     apps: typing.Any, schema_editor: typing.Any
 ):  # pylint: disable=unused-argument
     Transport = apps.get_model('uds', 'Transport')
-    for transport in Transport.objects.all():
+    db_alias = schema_editor.connection.alias
+    for transport in Transport.objects.using(db_alias).all():
         if transport.networks.count() == 0:
             transport.net_filtering = auth.NO_FILTERING
         else:
@@ -59,12 +56,13 @@ def update_transport_net_filtering(
 
 
 class Migration(migrations.Migration):
+    atomic = False
     dependencies = [
         ("uds", "0043_auto_20220704_2120"),
     ]
 
     operations = [
-        migrations.RunPython(remove_null_service_pools, nop),
+        migrations.RunPython(remove_null_service_pools, elidable=True),
         migrations.CreateModel(
             name="Notification",
             fields=[
@@ -234,7 +232,7 @@ class Migration(migrations.Migration):
             field=models.IntegerField(default=4),
         ),
         # Run python code to update network model
-        migrations.RunPython(update_network_model, nop),
+        migrations.RunPython(update_network_model, elidable=True),
         migrations.RemoveField(
             model_name="network",
             name="net_end",
@@ -254,7 +252,7 @@ class Migration(migrations.Migration):
             field=models.CharField(db_index=True, default=auth.NO_FILTERING, max_length=1),
         ),
         # Update transport net_filtering in base to networks field of Transport having any element
-        migrations.RunPython(update_transport_net_filtering, nop),
+        migrations.RunPython(update_transport_net_filtering, elidable=True),
         # And now remove the nets_positive field
         migrations.RemoveField(
             model_name="transport",

server/src/uds/migrations/0046_registered_server_migration_and_more.py

@@ -18,21 +18,22 @@ if typing.TYPE_CHECKING:
 
 def migrate_old_data(apps: typing.Any, schema_editor: typing.Any) -> None:
     try:
+        db_alias = schema_editor.connection.alias
         Server: 'type[uds.models.Server]' = apps.get_model('uds', 'Server')
         # Not typed, disappeared on this migration
         ActorToken = apps.get_model('uds', 'ActorToken')
 
         # First, add uuid to existing registered servers
-        for server in Server.objects.all():
+        for server in Server.objects.using(db_alias).all():
             server.uuid = uds.core.util.model.generate_uuid()
             server.save(update_fields=['uuid'])
 
         # Current Registered servers are tunnel servers, and all tunnel servers are linux os, so update ip
-        Server.objects.all().update(os_type=uds.core.types.os.KnownOS.LINUX.os_name())
+        Server.objects.using(db_alias).all().update(os_type=uds.core.types.os.KnownOS.LINUX.os_name())
 
         # Now append actors to registered servers, with "unknown" os type (legacy)
-        for token in ActorToken.objects.all():
+        for token in ActorToken.objects.using(db_alias).all():
-            Server.objects.create(
+            Server.objects.using(db_alias).create(
                 register_username=token.username,
                 register_ip=token.ip_from,
                 ip=token.ip,
@@ -63,12 +64,13 @@ def migrate_old_data(apps: typing.Any, schema_editor: typing.Any) -> None:
 
 
 def rollback_old_data(apps: typing.Any, schema_editor: typing.Any) -> None:
+    db_alias = schema_editor.connection.alias
     Server: 'type[uds.models.Server]' = apps.get_model('uds', 'Server')
     ActorToken = apps.get_model('uds', 'ActorToken')
-    for server in Server.objects.filter(type=ACTOR_TYPE):
+    for server in Server.objects.using(db_alias).filter(type=ACTOR_TYPE):
         if not server.data:
             continue  # Skip servers without data, they are not actors!!
-        ActorToken.objects.create(
+        ActorToken.objects.using(db_alias).create(
             username=server.register_username,
             ip_from=server.register_ip,
             ip=server.ip,
@@ -136,7 +138,7 @@ class Migration(migrations.Migration):
             model_name="server",
             old_name="username",
             new_name="register_username",
-        ),        
+        ),
         migrations.CreateModel(
             name="ServerGroup",
             fields=[

server/src/uds/migrations/fixers/properties_v4.py

@@ -12,13 +12,14 @@ def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
     Migrates old properties to new ones
     """
     try:
+        db_alias = schema_editor.connection.alias
         UserServiceProperty = apps.get_model('uds', 'UserServiceProperty')
         Properties: type['uds.models.Properties'] = apps.get_model('uds', 'Properties')
         # For testing
         # from uds.models import UserServiceProperty, Properties
 
-        for prop in UserServiceProperty.objects.all():
+        for prop in UserServiceProperty.objects.using(db_alias).all():
-            Properties.objects.create(
+            Properties.objects.using(db_alias).create(
                 owner_id=prop.user_service.uuid, owner_type='userservice', key=prop.name, value=prop.value
             )
     except Exception:
@@ -30,14 +31,15 @@ def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
     rollback migration
     """
     try:
+        db_alias = schema_editor.connection.alias
         UserServiceProperty = apps.get_model('uds', 'UserServiceProperty')
         Properties: type['uds.models.Properties'] = apps.get_model('uds', 'Properties')
         UserService: type['uds.models.UserService'] = apps.get_model('uds', 'UserService')
         # For testing
         # from uds.models import UserServiceProperty, Properties, UserService
 
-        for prop in Properties.objects.filter(owner_type='userservice'):
+        for prop in Properties.objects.using(db_alias).filter(owner_type='userservice'):
-            userService = UserService.objects.get(uuid=prop.owner_id)
+            userService = UserService.objects.using(db_alias).get(uuid=prop.owner_id)
-            UserServiceProperty.objects.create(name=prop.key, value=prop.value, user_service=userService)
+            UserServiceProperty.objects.using(db_alias).create(name=prop.key, value=prop.value, user_service=userService)
     except Exception:
         logger.error('Error migrating properties', exc_info=True)

server/src/uds/migrations/fixers/transports_v4/_migrator.py

@@ -9,18 +9,26 @@ if typing.TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-def tunnel_transport(apps: typing.Any, TransportType: typing.Type[typing.Any], serverAttr: str, is_html_server: bool = False) -> None:
+
+def tunnel_transport(
+    apps: typing.Any,
+    schema_editor: typing.Any,
+    TransportType: typing.Type[typing.Any],
+    serverAttr: str,
+    is_html_server: bool = False,
+) -> None:
     """
     Migrates an old tunnel transport to a new one (with tunnelServer)
     """
     try:
+        db_alias = schema_editor.connection.alias
         Transport: 'type[uds.models.Transport]' = apps.get_model('uds', 'Transport')
         ServerGroup: 'type[uds.models.ServerGroup]' = apps.get_model('uds', 'ServerGroup')
         Server: 'type[uds.models.Server]' = apps.get_model('uds', 'Server')
         # For testing
         # from uds.models import Transport, ServerGroup, Server
 
-        for t in Transport.objects.filter(data_type=TransportType.type_type):
+        for t in Transport.objects.using(db_alias).filter(data_type=TransportType.type_type):
             # Extract data
             obj = TransportType(Environment(t.uuid), None)
             obj.deserialize(t.data)
@@ -30,23 +38,31 @@ def tunnel_transport(apps: typing.Any, TransportType: typing.Type[typing.Any], s
             if is_html_server:
                 if not server.startswith('https://'):
                     # Skip if not https found
-                    logger.error('Skipping %s transport %s as it does not starts with https://', TransportType.__name__, t.name)
+                    logger.error(
+                        'Skipping %s transport %s as it does not starts with https://',
+                        TransportType.__name__,
+                        t.name,
+                    )
                     continue
-                host, port = (server+':443').split('https://')[1].split(':')[:2]
+                host, port = (server + ':443').split('https://')[1].split(':')[:2]
             else:  # Other servers are <host>:<port>
-                host, port = (server+':443').split(':')[:2]
+                host, port = (server + ':443').split(':')[:2]
             # If no host or port, skip
             if not host or not port:
-                logger.error('Skipping %s transport %s as it does not have host or port', TransportType.__name__, t.name)
+                logger.error(
+                    'Skipping %s transport %s as it does not have host or port', TransportType.__name__, t.name
+                )
                 continue
             # Look for an existing tunnel server (ServerGroup)
-            tunnel = ServerGroup.objects.filter(
+            tunnel = (
-                host=host, port=port, type=servers.ServerType.TUNNEL
+                ServerGroup.objects.using(db_alias)
-            ).first()
+                .filter(host=host, port=port, type=servers.ServerType.TUNNEL)
+                .first()
+            )
             if tunnel is None:
-                logger.info('Creating new tunnel server for %s: %s:%s', TransportType.__name__,  host, port)
+                logger.info('Creating new tunnel server for %s: %s:%s', TransportType.__name__, host, port)
                 # Create a new one, adding all tunnel servers to it
-                tunnel = ServerGroup.objects.create(
+                tunnel = ServerGroup.objects.using(db_alias).create(
                     name=f'Tunnel on {host}:{port}',
                     comments=f'Migrated from {t.name}',
                     host=host,
@@ -57,35 +73,42 @@ def tunnel_transport(apps: typing.Any, TransportType: typing.Type[typing.Any], s
                 # Append transport name to comments
                 tunnel.comments = f'{tunnel.comments}, {t.name}'[:255]
                 tunnel.save(update_fields=['comments'])
-            tunnel.servers.set(Server.objects.filter(type=servers.ServerType.TUNNEL))
+            tunnel.servers.set(Server.objects.using(db_alias).filter(type=servers.ServerType.TUNNEL))
             # Set tunnel server on transport
             logger.info('Setting tunnel server %s on transport %s', tunnel.name, t.name)
             obj.tunnel.value = tunnel.uuid
             # Save transport
             t.data = obj.serialize()
             t.save(update_fields=['data'])
-    except Exception as e:  # nosec: ignore this
+    except Exception:
-        print(e)
+        logger.exception('Exception found while migrating %s transports', TransportType)
-        logger.exception('Exception found while migrating HTML5RDP transports')
 
-def tunnel_transport_back(apps: typing.Any, TransportType: typing.Type[typing.Any], serverAttr: str, is_html_server: bool) -> None:
+
+def tunnel_transport_back(
+    apps: typing.Any,
+    schema_editor: typing.Any,
+    TransportType: typing.Type[typing.Any],
+    serverAttr: str,
+    is_html_server: bool,
+) -> None:
     """
     "Un-Migrates" an new tunnel transport to an old one (without tunnelServer)
     """
     try:
+        db_alias = schema_editor.connection.alias
         Transport: 'type[uds.models.Transport]' = apps.get_model('uds', 'Transport')
         ServerGroup: 'type[uds.models.ServerGroup]' = apps.get_model('uds', 'ServerGroup')
         # For testing
         # from uds.models import Transport, ServerGroup
 
-        for t in Transport.objects.filter(data_type=TransportType.type_type):
+        for t in Transport.objects.using(db_alias).filter(data_type=TransportType.type_type):
             # Extranct data
             obj = TransportType(Environment(t.uuid), None)
             obj.deserialize(t.data)
             # Guacamole server is https://<host>:<port>
             # Other tunnels are <host>:<port>
             server = getattr(obj, serverAttr)
-            tunnelServer = ServerGroup.objects.get(uuid=obj.tunnel.value)
+            tunnelServer = ServerGroup.objects.using(db_alias).get(uuid=obj.tunnel.value)
             if is_html_server:
                 server.value = f'https://{tunnelServer.host}:{tunnelServer.port}'
             else:
@@ -93,6 +116,5 @@ def tunnel_transport_back(apps: typing.Any, TransportType: typing.Type[typing.An
             # Save transport
             t.data = obj.serialize()
             t.save(update_fields=['data'])
-    except Exception as e:  # nosec: ignore this
+    except Exception:
-        print(e)
+        logger.exception('Exception found while migrating %s transports', TransportType)
-        logger.error('Exception found while migrating HTML5RDP transports: %s', e)

server/src/uds/migrations/fixers/transports_v4/html5rdp.py

@@ -84,8 +84,8 @@ class HTML5RDPTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, HTML5RDPTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport(apps, schema_editor, HTML5RDPTransport, 'guacamoleServer', is_html_server=True)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, HTML5RDPTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport_back(apps, schema_editor, HTML5RDPTransport, 'guacamoleServer', is_html_server=True)

server/src/uds/migrations/fixers/transports_v4/html5rds.py

@@ -86,8 +86,8 @@ class HTML5RDSTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, HTML5RDSTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport(apps, schema_editor, HTML5RDSTransport, 'guacamoleServer', is_html_server=True)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, HTML5RDSTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport_back(apps, schema_editor, HTML5RDSTransport, 'guacamoleServer', is_html_server=True)

server/src/uds/migrations/fixers/transports_v4/html5ssh.py

@@ -63,8 +63,8 @@ class HTML5SSHTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, HTML5SSHTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport(apps, schema_editor, HTML5SSHTransport, 'guacamoleServer', is_html_server=True)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, HTML5SSHTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport_back(apps, schema_editor, HTML5SSHTransport, 'guacamoleServer', is_html_server=True)

server/src/uds/migrations/fixers/transports_v4/html5vnc.py

@@ -66,8 +66,8 @@ class HTML5VNCTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, HTML5VNCTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport(apps, schema_editor, HTML5VNCTransport, 'guacamoleServer', is_html_server=True)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, HTML5VNCTransport, 'guacamoleServer', is_html_server=True)
+    _migrator.tunnel_transport_back(apps, schema_editor, HTML5VNCTransport, 'guacamoleServer', is_html_server=True)

server/src/uds/migrations/fixers/transports_v4/nicedcv.py

@@ -67,8 +67,8 @@ class NICEDCVTunnelTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, NICEDCVTunnelTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, NICEDCVTunnelTransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, NICEDCVTunnelTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, NICEDCVTunnelTransport, 'tunnelServer', is_html_server=False)

server/src/uds/migrations/fixers/transports_v4/nomachine.py

@@ -63,8 +63,8 @@ class TSNoMachineTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, TSNoMachineTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, TSNoMachineTransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, TSNoMachineTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, TSNoMachineTransport, 'tunnelServer', is_html_server=False)

server/src/uds/migrations/fixers/transports_v4/rdp.py

@@ -90,8 +90,8 @@ class TRDPTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, TRDPTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, TRDPTransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, TRDPTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, TRDPTransport, 'tunnelServer', is_html_server=False)

server/src/uds/migrations/fixers/transports_v4/rds.py

@@ -79,8 +79,8 @@ class TRDSTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, TRDSTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, TRDSTransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, TRDSTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, TRDSTransport, 'tunnelServer', is_html_server=False)

server/src/uds/migrations/fixers/transports_v4/spice.py

@@ -65,8 +65,8 @@ class TSPICETransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, TSPICETransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, TSPICETransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, TSPICETransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, TSPICETransport, 'tunnelServer', is_html_server=False)

server/src/uds/migrations/fixers/transports_v4/x2go.py

@@ -70,8 +70,8 @@ class TX2GOTransport(transports.Transport):
 
 
 def migrate(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport(apps, TX2GOTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport(apps, schema_editor, TX2GOTransport, 'tunnelServer', is_html_server=False)
 
 
 def rollback(apps: typing.Any, schema_editor: typing.Any) -> None:
-    _migrator.tunnel_transport_back(apps, TX2GOTransport, 'tunnelServer', is_html_server=False)
+    _migrator.tunnel_transport_back(apps, schema_editor, TX2GOTransport, 'tunnelServer', is_html_server=False)