From 8ba9f722ec7f9ee50f6ebb328cc531185037f21a Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@netboxlabs.com>
Date: Fri, 27 Jan 2023 15:23:36 -0500
Subject: [PATCH] Add management command to sync data sources

---
 netbox/core/management/__init__.py            |  0
 netbox/core/management/commands/__init__.py   |  0
 .../management/commands/syncdatasource.py     | 41 +++++++++++++++++++
 3 files changed, 41 insertions(+)
 create mode 100644 netbox/core/management/__init__.py
 create mode 100644 netbox/core/management/commands/__init__.py
 create mode 100644 netbox/core/management/commands/syncdatasource.py

diff --git a/netbox/core/management/__init__.py b/netbox/core/management/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/netbox/core/management/commands/__init__.py b/netbox/core/management/commands/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/netbox/core/management/commands/syncdatasource.py b/netbox/core/management/commands/syncdatasource.py
new file mode 100644
index 000000000..3d73f70ab
--- /dev/null
+++ b/netbox/core/management/commands/syncdatasource.py
@@ -0,0 +1,41 @@
+from django.core.management.base import BaseCommand, CommandError
+
+from core.models import DataSource
+
+
+class Command(BaseCommand):
+    help = "Synchronize a data source from its remote upstream"
+
+    def add_arguments(self, parser):
+        parser.add_argument('name', nargs='*', help="Data source(s) to synchronize")
+        parser.add_argument(
+            "--all", action='store_true', dest='sync_all',
+            help="Synchronize all data sources"
+        )
+
+    def handle(self, *args, **options):
+
+        # Find DataSources to sync
+        if options['sync_all']:
+            datasources = DataSource.objects.all()
+        elif options['name']:
+            datasources = DataSource.objects.filter(name__in=options['name'])
+            # Check for invalid names
+            found_names = {ds['name'] for ds in datasources.values('name')}
+            if invalid_names := set(options['name']) - found_names:
+                raise CommandError(f"Invalid data source names: {', '.join(invalid_names)}")
+        else:
+            raise CommandError(f"Must specify at least one data source, or set --all.")
+
+        if len(options['name']) > 1:
+            self.stdout.write(f"Syncing {len(datasources)} data sources.")
+
+        for i, datasource in enumerate(datasources, start=1):
+            self.stdout.write(f"[{i}] Syncing {datasource}... ", ending='')
+            self.stdout.flush()
+            datasource.sync()
+            self.stdout.write(datasource.get_status_display())
+            self.stdout.flush()
+
+        if len(options['name']) > 1:
+            self.stdout.write(f"Finished.")