Coverage for core/management/commands/syncdatasource.py: 0%

1from django.core.management.base import BaseCommand, CommandError 

2 

3from core.choices import DataSourceStatusChoices 

4from core.models import DataSource 

5 

6 

7class Command(BaseCommand): 

8 help = "Synchronize a data source from its remote upstream" 

9 

10 def add_arguments(self, parser): 

11 parser.add_argument('name', nargs='*', help="Data source(s) to synchronize") 

12 parser.add_argument( 

13 "--all", action='store_true', dest='sync_all', 

14 help="Synchronize all data sources" 

15 ) 

16 

17 def handle(self, *args, **options): 

18 

19 # Find DataSources to sync 

20 if options['sync_all']: 

21 datasources = DataSource.objects.all() 

22 elif options['name']: 

23 datasources = DataSource.objects.filter(name__in=options['name']) 

24 # Check for invalid names 

25 found_names = {ds['name'] for ds in datasources.values('name')} 

26 if invalid_names := set(options['name']) - found_names: 

27 raise CommandError(f"Invalid data source names: {', '.join(invalid_names)}") 

28 else: 

29 raise CommandError("Must specify at least one data source, or set --all.") 

30 

31 if len(options['name']) > 1: 

32 self.stdout.write(f"Syncing {len(datasources)} data sources.") 

33 

34 for i, datasource in enumerate(datasources, start=1): 

35 self.stdout.write(f"[{i}] Syncing {datasource}... ", ending='') 

36 self.stdout.flush() 

37 try: 

38 datasource.sync() 

39 self.stdout.write(datasource.get_status_display()) 

40 self.stdout.flush() 

41 except Exception as e: 

42 DataSource.objects.filter(pk=datasource.pk).update(status=DataSourceStatusChoices.FAILED) 

43 raise e 

44 

45 if len(options['name']) > 1: 

46 self.stdout.write("Finished.")