2020class Command (BaseCommand ):
2121 def add_arguments (self , parser ):
2222 parser .add_argument ('sources' , nargs = '*' , type = str , help = 'Names of the sources to load (if omitted, load all)' )
23+ parser .add_argument ('--overwrite' , action = 'store_true' , help = 'Overwrite existing sources and source configs' )
2324
2425 def handle (self , * args , ** options ):
2526 sources = options .get ('sources' )
@@ -32,7 +33,7 @@ def handle(self, *args, **options):
3233 with transaction .atomic ():
3334 self .known_harvesters = self .sync_drivers ('share.harvesters' , apps .get_model ('share.Harvester' ))
3435 self .known_transformers = self .sync_drivers ('share.transformers' , apps .get_model ('share.Transformer' ))
35- self .update_sources (source_dirs )
36+ self .update_sources (source_dirs , overwrite = options . get ( 'overwrite' ) )
3637
3738 def sync_drivers (self , namespace , model ):
3839 names = set (extension .ExtensionManager (namespace ).entry_points_names ())
@@ -43,7 +44,7 @@ def sync_drivers(self, namespace, model):
4344 print ('Warning: Missing {} drivers: {}' .format (model ._meta .model_name , missing ))
4445 return names
4546
46- def update_sources (self , source_dirs ):
47+ def update_sources (self , source_dirs , overwrite ):
4748 loaded_sources = set ()
4849 loaded_configs = set ()
4950 for source_dir in source_dirs :
@@ -55,36 +56,42 @@ def update_sources(self, source_dirs):
5556 loaded_sources .add (name )
5657
5758 user = self .get_or_create_user (serialized .pop ('user' ))
58- source , _ = Source .objects .update_or_create (
59- name = name ,
60- defaults = {
61- 'user' : user ,
62- ** self .process_defaults (Source , serialized )
63- }
64- )
59+ source_defaults = {
60+ 'user' : user ,
61+ ** self .process_defaults (Source , serialized )
62+ }
63+ if overwrite :
64+ source , _ = Source .objects .update_or_create (name = name , defaults = source_defaults )
65+ else :
66+ source , _ = Source .objects .get_or_create (name = name , defaults = source_defaults )
67+
6568 with open (os .path .join (source_dir , 'icon.ico' ), 'rb' ) as fobj :
6669 source .icon .save (name , File (fobj ))
6770 for config in configs :
6871 assert config ['label' ] not in loaded_configs
6972 loaded_configs .add (config ['label' ])
70- self .update_source_config (source , config )
73+ self .update_source_config (source , config , overwrite )
7174
72- def update_source_config (self , source , serialized ):
75+ def update_source_config (self , source , serialized , overwrite ):
7376 label = serialized .pop ('label' )
7477 if serialized ['harvester' ] and serialized ['harvester' ] not in self .known_harvesters :
7578 print ('Unknown harvester {}! Skipping source config {}' .format (serialized ['harvester' ], label ))
7679 return
7780 if serialized ['transformer' ] and serialized ['transformer' ] not in self .known_transformers :
7881 print ('Unknown transformer {}! Skipping source config {}' .format (serialized ['transformer' ], label ))
7982 return
80- source_config , _ = apps .get_model ('share.SourceConfig' ).objects .update_or_create (
81- label = label ,
82- defaults = {
83- 'source' : source ,
84- ** self .process_defaults (apps .get_model ('share.SourceConfig' ), serialized )
85- }
86- )
87- self .schedule_harvest_task (source_config .label , source_config .disabled )
83+
84+ SourceConfig = apps .get_model ('share.SourceConfig' )
85+ config_defaults = {
86+ 'source' : source ,
87+ ** self .process_defaults (SourceConfig , serialized )
88+ }
89+ if overwrite :
90+ source_config , created = SourceConfig .objects .update_or_create (label = label , defaults = config_defaults )
91+ else :
92+ source_config , created = SourceConfig .objects .get_or_create (label = label , defaults = config_defaults )
93+ if overwrite or created :
94+ self .schedule_harvest_task (source_config .label , source_config .disabled )
8895
8996 def get_or_create_user (self , username ):
9097 try :
0 commit comments