[change] Execute credentials auto_add in the background #479

Closes #479

openwisp_controller/connection/base/models.py

@@ -25,7 +25,7 @@
     get_command_schema,
 )
 from ..signals import is_working_changed
-from ..tasks import launch_command
+from ..tasks import auto_add_credentials_to_devices, launch_command
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +75,10 @@ class AbstractCredentials(ConnectorMixin, ShareableOrgMixinUniqueName, BaseModel
     Credentials for access
     """
 
+    # Controls the number of objects which can be stored in memory
+    # before commiting them to database during bulk auto add operation.
+    chunk_size = 1000
+
     connector = models.CharField(
         _('connection type'),
         choices=app_settings.CONNECTORS,
@@ -110,27 +114,41 @@ def __str__(self):
 
     def save(self, *args, **kwargs):
         super().save(*args, **kwargs)
-        self.auto_add_to_devices()
+        if self.auto_add:
+            transaction.on_commit(
+                lambda: auto_add_credentials_to_devices.delay(
+                    credential_id=self.pk, organization_id=self.organization_id
+                )
+            )
 
-    def auto_add_to_devices(self):
+    @classmethod
+    def auto_add_to_devices(cls, credential_id, organization_id):
         """
         When ``auto_add`` is ``True``, adds the credentials
         to each relevant ``Device`` and ``DeviceConnection`` objects
         """
-        if not self.auto_add:
-            return
-        device_model = load_model('config', 'Device')
-        devices = device_model.objects.exclude(config=None)
-        org = self.organization
-        if org:
-            devices = devices.filter(organization=org)
+        DeviceConnection = load_model('connection', 'DeviceConnection')
+        Device = load_model('config', 'Device')
+
+        devices = Device.objects.exclude(config=None)
+        if organization_id:
+            devices = devices.filter(organization_id=organization_id)
         # exclude devices which have been already added
-        devices = devices.exclude(deviceconnection__credentials=self)
-        for device in devices:
-            DeviceConnection = load_model('connection', 'DeviceConnection')
-            conn = DeviceConnection(device=device, credentials=self, enabled=True)
+        devices = devices.exclude(deviceconnection__credentials_id=credential_id)
+        device_connections = []
+        for device in devices.iterator():
+            conn = DeviceConnection(
+                device=device, credentials_id=credential_id, enabled=True
+            )
             conn.full_clean()
-            conn.save()
+            device_connections.append(conn)
+            # Send create query when chunk_size is reached
+            # and reset the device_connections list
+            if len(device_connections) >= cls.chunk_size:
+                DeviceConnection.objects.bulk_create(device_connections)
+                device_connections = []
+        if len(device_connections):
+            DeviceConnection.objects.bulk_create(device_connections)
 
     @classmethod
     def auto_add_credentials_to_device(cls, instance, created, **kwargs):

openwisp_controller/connection/tasks.py

@@ -64,3 +64,9 @@ def launch_command(command_id):
         command.status = 'failed'
         command._add_output(_(f'Internal system error: {e}'))
         command.save()
+
+
+@shared_task(soft_time_limit=180)
+def auto_add_credentials_to_devices(credential_id, organization_id):
+    Credentials = load_model('connection', 'Credentials')
+    Credentials.auto_add_to_devices(credential_id, organization_id)

openwisp_controller/connection/tests/test_models.py

@@ -272,44 +272,6 @@ def test_auto_add_to_new_device(self):
         self.assertEqual(d.deviceconnection_set.count(), 1)
         self.assertEqual(d.deviceconnection_set.first().credentials, c)
 
-    def test_auto_add_to_existing_device_on_creation(self):
-        d = self._create_device(organization=Organization.objects.first())
-        self._create_config(device=d)
-        self.assertEqual(d.deviceconnection_set.count(), 0)
-        c = self._create_credentials(auto_add=True, organization=None)
-        org2 = Organization.objects.create(name='org2', slug='org2')
-        self._create_credentials(name='cred2', auto_add=True, organization=org2)
-        d.refresh_from_db()
-        self.assertEqual(d.deviceconnection_set.count(), 1)
-        self.assertEqual(d.deviceconnection_set.first().credentials, c)
-        self._create_credentials(name='cred3', auto_add=False, organization=None)
-        d.refresh_from_db()
-        self.assertEqual(d.deviceconnection_set.count(), 1)
-        self.assertEqual(d.deviceconnection_set.first().credentials, c)
-
-    def test_auto_add_to_existing_device_on_edit(self):
-        d = self._create_device(organization=Organization.objects.first())
-        self._create_config(device=d)
-        self.assertEqual(d.deviceconnection_set.count(), 0)
-        c = self._create_credentials(auto_add=False, organization=None)
-        org2 = Organization.objects.create(name='org2', slug='org2')
-        self._create_credentials(name='cred2', auto_add=True, organization=org2)
-        d.refresh_from_db()
-        self.assertEqual(d.deviceconnection_set.count(), 0)
-        c.auto_add = True
-        c.full_clean()
-        c.save()
-        d.refresh_from_db()
-        self.assertEqual(d.deviceconnection_set.count(), 1)
-        self.assertEqual(d.deviceconnection_set.first().credentials, c)
-        # ensure further edits are idempotent
-        c.name = 'changed'
-        c.full_clean()
-        c.save()
-        d.refresh_from_db()
-        self.assertEqual(d.deviceconnection_set.count(), 1)
-        self.assertEqual(d.deviceconnection_set.first().credentials, c)
-
     def test_auto_add_device_missing_config(self):
         org = Organization.objects.first()
         self._create_device(organization=org)
@@ -906,3 +868,64 @@ def test_schedule_command_called(self, connect_mocked):
         command.refresh_from_db()
         self.assertEqual(command.status, 'success')
         self.assertEqual(command.output, 'mocked\n')
+
+    def test_auto_add_to_existing_device_on_edit(self):
+        d = self._create_device(organization=self._get_org())
+        self._create_config(device=d)
+        self.assertEqual(d.deviceconnection_set.count(), 0)
+        c = self._create_credentials(auto_add=False, organization=None)
+        org2 = Organization.objects.create(name='org2', slug='org2')
+        self._create_credentials(name='cred2', auto_add=True, organization=org2)
+        d.refresh_from_db()
+        self.assertEqual(d.deviceconnection_set.count(), 0)
+        c.auto_add = True
+        c.full_clean()
+        c.save()
+        d.refresh_from_db()
+        self.assertEqual(d.deviceconnection_set.count(), 1)
+        self.assertEqual(d.deviceconnection_set.first().credentials, c)
+        # ensure further edits are idempotent
+        c.name = 'changed'
+        c.full_clean()
+        c.save()
+        d.refresh_from_db()
+        self.assertEqual(d.deviceconnection_set.count(), 1)
+        self.assertEqual(d.deviceconnection_set.first().credentials, c)
+
+    def test_auto_add_to_existing_device_on_creation(self):
+        d = self._create_device(organization=self._get_org())
+        self._create_config(device=d)
+        self.assertEqual(d.deviceconnection_set.count(), 0)
+        c = self._create_credentials(auto_add=True, organization=None)
+        org2 = Organization.objects.create(name='org2', slug='org2')
+        self._create_credentials(name='cred2', auto_add=True, organization=org2)
+        d.refresh_from_db()
+        self.assertEqual(d.deviceconnection_set.count(), 1)
+        self.assertEqual(d.deviceconnection_set.first().credentials, c)
+        self._create_credentials(name='cred3', auto_add=False, organization=None)
+        d.refresh_from_db()
+        self.assertEqual(d.deviceconnection_set.count(), 1)
+        self.assertEqual(d.deviceconnection_set.first().credentials, c)
+
+    def test_chunk_size(self):
+        org = self._get_org()
+        self._create_config(device=self._create_device(organization=org))
+        self._create_config(
+            device=self._create_device(
+                organization=org, name='device2', mac_address='22:22:22:22:22:22'
+            )
+        )
+        self._create_config(
+            device=self._create_device(
+                organization=org, name='device3', mac_address='33:33:33:33:33:33'
+            )
+        )
+        with self.assertNumQueries(28):
+            credential = self._create_credentials(auto_add=True, organization=org)
+        self.assertEqual(credential.deviceconnection_set.count(), 3)
+
+        with mock.patch.object(Credentials, 'chunk_size', 2):
+            with self.assertNumQueries(30):
+                credential = self._create_credentials(
+                    name='Mocked Credential', auto_add=True, organization=org
+                )