Bladeren bron

fix(core): Apply data source syncs atomically (#23480)

Apply data file changes, the completed status and last_synced in one
transaction, serialized per data source by an advisory lock taken
before its files are read. A sync that fails during its import now
leaves the previously synced files unchanged. The data source model
tests move into one DataSourceTestCase.

Fixes #23464
Martin Hauser 1 dag geleden
bovenliggende
commit
758af3335a
3 gewijzigde bestanden met toevoegingen van 162 en 21 verwijderingen
  1. 25 11
      netbox/core/models/data.py
  2. 134 10
      netbox/core/tests/test_models.py
  3. 3 0
      netbox/netbox/constants.py

+ 25 - 11
netbox/core/models/data.py

@@ -8,12 +8,12 @@ import yaml
 from django.conf import settings
 from django.conf import settings
 from django.core.exceptions import ValidationError
 from django.core.exceptions import ValidationError
 from django.core.validators import RegexValidator
 from django.core.validators import RegexValidator
-from django.db import models
+from django.db import connections, models, router, transaction
 from django.urls import reverse
 from django.urls import reverse
 from django.utils import timezone
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 
 
-from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
+from netbox.constants import ADVISORY_LOCK_KEYS, CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
 from netbox.models import PrimaryModel
 from netbox.models import PrimaryModel
 from netbox.models.features import JobsMixin
 from netbox.models.features import JobsMixin
 from netbox.registry import registry
 from netbox.registry import registry
@@ -183,6 +183,8 @@ class DataSource(JobsMixin, PrimaryModel):
         """
         """
         from core.signals import post_sync, pre_sync
         from core.signals import post_sync, pre_sync
 
 
+        using = router.db_for_write(self.__class__, instance=self)
+
         if self.status == DataSourceStatusChoices.SYNCING:
         if self.status == DataSourceStatusChoices.SYNCING:
             raise SyncError(_("Cannot initiate sync; syncing already in progress."))
             raise SyncError(_("Cannot initiate sync; syncing already in progress."))
 
 
@@ -190,7 +192,7 @@ class DataSource(JobsMixin, PrimaryModel):
         pre_sync.send(sender=self.__class__, instance=self)
         pre_sync.send(sender=self.__class__, instance=self)
 
 
         self.status = DataSourceStatusChoices.SYNCING
         self.status = DataSourceStatusChoices.SYNCING
-        DataSource.objects.filter(pk=self.pk).update(status=self.status)
+        DataSource.objects.using(using).filter(pk=self.pk).update(status=self.status)
 
 
         # Replicate source data locally
         # Replicate source data locally
         try:
         try:
@@ -199,10 +201,18 @@ class DataSource(JobsMixin, PrimaryModel):
             raise SyncError(
             raise SyncError(
                 _("There was an error initializing the backend. A dependency needs to be installed: ") + str(e)
                 _("There was an error initializing the backend. A dependency needs to be installed: ") + str(e)
             )
             )
-        with backend.fetch() as local_path:
+
+        # The backend's cleanup runs inside the transaction, so its failure rolls back the import too
+        with transaction.atomic(using=using), backend.fetch() as local_path:
+
+            # Serialize syncs of this source without locking its row, which edits and status updates still write
+            with connections[using].cursor() as cursor:
+                cursor.execute(
+                    'SELECT pg_advisory_xact_lock(%s, %s)', (ADVISORY_LOCK_KEYS['data-source-sync'], self.pk)
+                )
 
 
             logger.debug(f'Syncing files from source root {local_path}')
             logger.debug(f'Syncing files from source root {local_path}')
-            data_files = self.datafiles.all()
+            data_files = DataFile.objects.using(using).filter(source=self)
             known_paths = {df.path for df in data_files}
             known_paths = {df.path for df in data_files}
             logger.debug(f'Starting with {len(known_paths)} known files')
             logger.debug(f'Starting with {len(known_paths)} known files')
 
 
@@ -220,13 +230,13 @@ class DataSource(JobsMixin, PrimaryModel):
                     continue
                     continue
 
 
             # Bulk update modified files
             # Bulk update modified files
-            updated_count = DataFile.objects.bulk_update(
+            updated_count = DataFile.objects.using(using).bulk_update(
                 updated_files, ('last_updated', 'size', 'hash', 'data'), batch_size=settings.BULK_UPDATE_CHUNK_SIZE
                 updated_files, ('last_updated', 'size', 'hash', 'data'), batch_size=settings.BULK_UPDATE_CHUNK_SIZE
             )
             )
             logger.debug(f"Updated {updated_count} files")
             logger.debug(f"Updated {updated_count} files")
 
 
             # Bulk delete deleted files
             # Bulk delete deleted files
-            deleted_count, __ = DataFile.objects.filter(pk__in=deleted_file_ids).delete()
+            deleted_count, __ = DataFile.objects.using(using).filter(pk__in=deleted_file_ids).delete()
             logger.debug(f"Deleted {deleted_count} files")
             logger.debug(f"Deleted {deleted_count} files")
 
 
             # Walk the local replication to find new files
             # Walk the local replication to find new files
@@ -239,13 +249,17 @@ class DataSource(JobsMixin, PrimaryModel):
                 datafile.refresh_from_disk(source_root=local_path)
                 datafile.refresh_from_disk(source_root=local_path)
                 datafile.full_clean()
                 datafile.full_clean()
                 new_datafiles.append(datafile)
                 new_datafiles.append(datafile)
-            created_count = len(DataFile.objects.bulk_create(new_datafiles, batch_size=100))
+            created_count = len(DataFile.objects.using(using).bulk_create(new_datafiles, batch_size=100))
             logger.debug(f"Created {created_count} data files")
             logger.debug(f"Created {created_count} data files")
 
 
-        # Update status & last_synced time
+            # Update status & last_synced time
+            last_synced = timezone.now()
+            DataSource.objects.using(using).filter(pk=self.pk).update(
+                status=DataSourceStatusChoices.COMPLETED, last_synced=last_synced
+            )
+
         self.status = DataSourceStatusChoices.COMPLETED
         self.status = DataSourceStatusChoices.COMPLETED
-        self.last_synced = timezone.now()
-        DataSource.objects.filter(pk=self.pk).update(status=self.status, last_synced=self.last_synced)
+        self.last_synced = last_synced
 
 
         # Emit the post_sync signal
         # Emit the post_sync signal
         post_sync.send(sender=self.__class__, instance=self)
         post_sync.send(sender=self.__class__, instance=self)

+ 134 - 10
netbox/core/tests/test_models.py

@@ -1,21 +1,64 @@
+import tempfile
 import uuid
 import uuid
+from contextlib import contextmanager
 from datetime import timedelta
 from datetime import timedelta
+from io import StringIO
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 from unittest.mock import MagicMock, patch
 
 
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ObjectDoesNotExist
 from django.core.exceptions import ObjectDoesNotExist
-from django.test import TestCase
+from django.core.management import call_command
+from django.db import IntegrityError, connection
+from django.test import TestCase, override_settings
+from django.test.utils import CaptureQueriesContext
 from django.utils import timezone
 from django.utils import timezone
 
 
-from core.choices import JobNotificationChoices, JobStatusChoices, ObjectChangeActionChoices
-from core.models import DataSource, Job, ObjectType
+from core.choices import DataSourceStatusChoices, JobNotificationChoices, JobStatusChoices, ObjectChangeActionChoices
+from core.data_backends import LocalBackend
+from core.exceptions import SyncError
+from core.models import AutoSyncRecord, DataFile, DataSource, Job, ObjectType
+from core.signals import post_sync
 from dcim.models import Device, Location, Site
 from dcim.models import Device, Location, Site
-from extras.models import Notification
+from extras.models import ConfigTemplate, Notification
 from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
 from netbox.constants import CENSOR_TOKEN, CENSOR_TOKEN_CHANGED
 from users.models import User
 from users.models import User
-
-
-class DataSourceIgnoreRulesTestCase(TestCase):
+from utilities.querysets import RestrictedQuerySet
+from utilities.testing import PinnedConnectionRouter
+
+
+class DataSourceTestCase(TestCase):
+
+    def setUp(self):
+        # A synced local source, used by the sync tests
+        temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(temp_dir.cleanup)
+        self.root = Path(temp_dir.name)
+        (self.root / 'a.txt').write_text('old\n')
+        (self.root / 'obsolete.txt').write_text('old\n')
+        self.datasource = DataSource.objects.create(name='Data Source 2', type='local', source_url=self.root.as_uri())
+        self.datasource.sync()
+        self.original_files = list(self.datasource.datafiles.values())
+        self.last_synced = self.datasource.last_synced
+
+        # The next sync finds one file changed, one removed and one added
+        (self.root / 'a.txt').write_text('new\n')
+        (self.root / 'obsolete.txt').unlink()
+        (self.root / 'b.txt').write_text('new\n')
+
+    def assert_files_unchanged(self):
+        self.assertEqual(list(self.datasource.datafiles.values()), self.original_files)
+        self.datasource.refresh_from_db()
+        self.assertEqual(self.datasource.last_synced, self.last_synced)
+
+    def assert_files_synced(self):
+        self.assertEqual(
+            dict(self.datasource.datafiles.values_list('path', 'data')),
+            {'a.txt': b'new\n', 'b.txt': b'new\n'},
+        )
+        self.datasource.refresh_from_db()
+        self.assertEqual(self.datasource.status, DataSourceStatusChoices.COMPLETED)
+        self.assertGreater(self.datasource.last_synced, self.last_synced)
 
 
     def test_no_ignore_rules(self):
     def test_no_ignore_rules(self):
         ds = DataSource(ignore_rules='')
         ds = DataSource(ignore_rules='')
@@ -34,9 +77,6 @@ class DataSourceIgnoreRulesTestCase(TestCase):
         self.assertTrue(ds._ignore('dev/script.py'))
         self.assertTrue(ds._ignore('dev/script.py'))
         self.assertFalse(ds._ignore('prod/script.py'))
         self.assertFalse(ds._ignore('prod/script.py'))
 
 
-
-class DataSourceChangeLoggingTestCase(TestCase):
-
     def test_password_added_on_create(self):
     def test_password_added_on_create(self):
         datasource = DataSource.objects.create(
         datasource = DataSource.objects.create(
             name='Data Source 1',
             name='Data Source 1',
@@ -151,6 +191,90 @@ class DataSourceChangeLoggingTestCase(TestCase):
         self.assertEqual(objectchange.postchange_data['parameters']['username'], 'username2')
         self.assertEqual(objectchange.postchange_data['parameters']['username'], 'username2')
         self.assertEqual(objectchange.postchange_data['parameters']['password'], CENSOR_TOKEN)
         self.assertEqual(objectchange.postchange_data['parameters']['password'], CENSOR_TOKEN)
 
 
+    def test_sync_rolls_back_read_failure(self):
+        """A sync that fails reading a new file keeps the synced files and their links."""
+        obsolete = self.datasource.datafiles.get(path='obsolete.txt')
+        template = ConfigTemplate.objects.create(
+            name='Config Template 1',
+            data_source=self.datasource,
+            data_file=obsolete,
+            data_path=obsolete.path,
+            template_code='old\n',
+            auto_sync_enabled=True,
+        )
+        autosync = AutoSyncRecord.objects.get(datafile=obsolete)
+        (self.root / 'link.txt').symlink_to('missing.txt')
+
+        with patch.object(post_sync, 'send') as send:
+            with self.assertRaises(FileNotFoundError):
+                call_command('syncdatasource', self.datasource.name, stdout=StringIO())
+
+        self.assert_files_unchanged()
+        self.assertEqual(self.datasource.status, DataSourceStatusChoices.FAILED)
+        template.refresh_from_db()
+        self.assertEqual(template.data_file_id, obsolete.pk)
+        self.assertTrue(AutoSyncRecord.objects.filter(pk=autosync.pk, datafile=obsolete).exists())
+        send.assert_not_called()
+
+    def test_sync_rolls_back_create_failure(self):
+        """A sync that fails creating its new files keeps the synced files."""
+        with patch.object(RestrictedQuerySet, 'bulk_create', side_effect=IntegrityError('Simulated bulk create error')):
+            with self.assertRaisesMessage(IntegrityError, 'Simulated bulk create error'):
+                self.datasource.sync()
+
+        self.assert_files_unchanged()
+
+    def test_sync_rolls_back_backend_cleanup_failure(self):
+        """A sync whose backend fails during cleanup keeps the synced files."""
+
+        @contextmanager
+        def fetch(backend):
+            yield str(self.root)
+            raise OSError('Simulated cleanup error')
+
+        with patch.object(LocalBackend, 'fetch', fetch):
+            with self.assertRaisesMessage(OSError, 'Simulated cleanup error'):
+                self.datasource.sync()
+
+        self.assert_files_unchanged()
+
+    def test_sync_locks_datasource_before_reading_files(self):
+        """Take the sync lock inside the transaction, before the files are read and the sync completes."""
+        with CaptureQueriesContext(connection) as context:
+            self.datasource.sync()
+
+        queries = [query['sql'] for query in context.captured_queries]
+        begins = [i for i, sql in enumerate(queries) if sql.startswith('SAVEPOINT')]
+        locks = [i for i, sql in enumerate(queries) if 'pg_advisory_xact_lock' in sql]
+        reads = [i for i, sql in enumerate(queries) if sql.startswith('SELECT') and 'FROM "core_datafile"' in sql]
+        completions = [
+            i for i, sql in enumerate(queries) if sql.startswith('UPDATE "core_datasource"') and "'completed'" in sql
+        ]
+        commits = [i for i, sql in enumerate(queries) if sql.startswith('RELEASE SAVEPOINT')]
+        self.assertTrue(locks)
+        self.assertLess(begins[0], locks[0])
+        self.assertLess(locks[0], reads[0])
+        self.assertLess(reads[0], completions[0])
+        self.assertLess(completions[0], commits[0])
+
+    def test_post_sync_failure_preserves_synced_files(self):
+        """A failing post_sync receiver leaves the completed sync in place."""
+        with patch.object(post_sync, 'send', side_effect=SyncError('Simulated post_sync error')):
+            with self.assertRaisesMessage(SyncError, 'Simulated post_sync error'):
+                self.datasource.sync()
+
+        self.assert_files_synced()
+
+    def test_sync_pins_datafile_queries_to_transaction_connection(self):
+        """Run the data file reads, updates and deletes of a sync on its transaction's connection."""
+        # Django routes a new data file's construction and validation itself
+        (self.root / 'b.txt').unlink()
+
+        with override_settings(DATABASE_ROUTERS=[PinnedConnectionRouter(DataFile)]):
+            self.datasource.sync()
+
+        self.assertEqual(dict(self.datasource.datafiles.values_list('path', 'data')), {'a.txt': b'new\n'})
+
 
 
 class ObjectTypeTestCase(TestCase):
 class ObjectTypeTestCase(TestCase):
 
 

+ 3 - 0
netbox/netbox/constants.py

@@ -34,6 +34,9 @@ ADVISORY_LOCK_KEYS = {
 
 
     # Custom field data
     # Custom field data
     'custom-field-data': 115100,
     'custom-field-data': 115100,
+
+    # Data source synchronization
+    'data-source-sync': 120100,
 }
 }
 
 
 # General-purpose tokens
 # General-purpose tokens