@@ -106,6 +106,7 @@ def upload_to_azure_blob(self, bundle_uuid, bundle_location, is_dir=False):
106106 file_name = "contents.tar.gz" if is_dir else "contents.gz"
107107 target_location = f"{ self .target_store_url } /{ bundle_uuid } /{ file_name } "
108108
109+ # TODO: This step might cause repeated upload. Can not check by checking size (Azure blob storage is zipped).
109110 if FileSystems .exists (target_location ):
110111 path_util .remove (target_location )
111112
@@ -202,33 +203,35 @@ def delete_original_bundle(self, bundle_uuid, bundle_location):
202203 description = 'Manages your local CodaLab Worksheets service deployment'
203204 )
204205 parser .add_argument (
205- '-w ' , '--worksheet ' , type = str , help = 'The worksheet uuid that needs migration'
206+ '-a ' , '--all ' , help = 'Run migration on all worksheets and all bundles' , action = 'store_true' ,
206207 )
207- parser .add_argument ('--target_store_name' , type = str , help = 'The destination bundle store name' )
208208 parser .add_argument (
209- '-d' ,
210- '--dry-run' ,
211- help = 'Only upload the bundle to Azure, does not modify database' ,
212- action = 'store_true' ,
209+ '-w' , '--worksheet' , type = str , help = 'The worksheet uuid that needs migration'
213210 )
211+ parser .add_argument ('--target_store_name' , type = str , help = 'The destination bundle store name' )
214212 parser .add_argument (
215- '-k ' , '--keep ' , help = 'Keep bundle content in origin bundle store ' , action = 'store_true'
213+ '-c ' , '--change_db ' , help = 'Change the bundle location in the database ' , action = 'store_true' ,
216214 )
215+ parser .add_argument ('-d' , '--delete' , help = 'Delete the original database' , action = 'store_true' )
217216
218217 args = parser .parse_args ()
219218
220219 worksheet_uuid = args .worksheet
221220 target_store_name = (
222221 "azure-store-default" if args .target_store_name is None else args .target_store_name
223222 )
224- if worksheet_uuid is not None and not spec_util .UUID_REGEX .match (worksheet_uuid ):
225- raise Exception ("Input worksheet uuid has wrong format. " )
226223
227224 # TODO: write output to log / log files
228225 migration = Migration (target_store_name )
229226 migration .setUp ()
230227
231- bundle_uuids = migration .get_bundle_uuids (worksheet_uuid )
228+ if args .all :
229+ bundle_uuids = migration .get_bundle_uuids (worksheet_uuid = None )
230+ else :
231+ # Must specify worksheet uuid
232+ if worksheet_uuid is not None and not spec_util .UUID_REGEX .match (worksheet_uuid ):
233+ raise Exception ("Input worksheet uuid has wrong format. " )
234+ bundle_uuids = migration .get_bundle_uuids (worksheet_uuid )
232235
233236 for bundle_uuid in bundle_uuids :
234237 logging .info (bundle_uuid )
@@ -256,8 +259,8 @@ def delete_original_bundle(self, bundle_uuid, bundle_location):
256259 is_dir = bundle_info ['type' ] == 'directory'
257260 migration .upload_to_azure_blob (bundle_uuid , bundle_location , is_dir )
258261
259- if not args .dry_run : # If dry_run, only upload to new bundle location
262+ if args .change_db : # If need to change the database, continue to run
260263 migration .modify_bundle_data (bundle , bundle_uuid , is_dir )
261264 migration .sanity_check (bundle_uuid , bundle_location , bundle_info , is_dir )
262- if not args .keep :
265+ if args .delete :
263266 migration .delete_original_bundle (bundle_uuid , bundle_location )
0 commit comments