Berkeley migration test imprv

src/app/berkeley_migration/berkeley_migration.ml

@@ -184,7 +184,8 @@ let mainnet_protocol_version =
   Protocol_version.create ~transaction:1 ~network:0 ~patch:0
 
 let mainnet_block_to_extensional ~logger ~mainnet_pool
-    ~(genesis_block : Mina_block.t) (block : Sql.Mainnet.Block.t) =
+    ~(genesis_block : Mina_block.t) (block : Sql.Mainnet.Block.t) ~bucket
+    ~batch_size =
   let query_mainnet_db ~f = Mina_caqti.query ~f mainnet_pool in
   let is_genesis_block = Int64.equal block.height Int64.one in
   let genesis_consensus_state =
@@ -199,7 +200,7 @@ let mainnet_block_to_extensional ~logger ~mainnet_pool
     (* we may try to be fetching more blocks than exist
        gsutil seems to get the ones that do exist, in that exist
     *)
-    let batch_size = 1000L in
+    let batch_size = Int64.of_int batch_size in
     if is_genesis_block then Deferred.unit
     else if !first_batch then (
       let num_blocks = Int64.(batch_size - (block.height % batch_size)) in
@@ -209,7 +210,7 @@ let mainnet_block_to_extensional ~logger ~mainnet_pool
           ; ("num_blocks", `Int (Int64.to_int_exn num_blocks))
           ] ;
       let%bind () =
-        Precomputed_block.fetch_batch ~height:block.height ~num_blocks
+        Precomputed_block.fetch_batch ~height:block.height ~num_blocks ~bucket
       in
       [%log info] "Done fetching first batch of precomputed blocks" ;
       first_batch := false ;
@@ -224,7 +225,7 @@ let mainnet_block_to_extensional ~logger ~mainnet_pool
           ] ;
       let%bind () =
         Precomputed_block.fetch_batch ~height:block.height
-          ~num_blocks:batch_size
+          ~num_blocks:batch_size ~bucket
       in
       [%log info] "Done fetching batch of precomputed blocks" ;
       Deferred.unit )
@@ -403,7 +404,7 @@ let mainnet_block_to_extensional ~logger ~mainnet_pool
       : Archive_lib.Extensional.Block.t )
 
 let main ~mainnet_archive_uri ~migrated_archive_uri ~runtime_config_file
-    ~end_global_slot () =
+    ~end_global_slot ~mina_network_blocks_bucket ~batch_size () =
   let logger = Logger.create () in
   let mainnet_archive_uri = Uri.of_string mainnet_archive_uri in
   let migrated_archive_uri = Uri.of_string migrated_archive_uri in
@@ -500,7 +501,7 @@ let main ~mainnet_archive_uri ~migrated_archive_uri ~runtime_config_file
               block.height block.state_hash ;
             let%bind extensional_block =
               mainnet_block_to_extensional ~logger ~mainnet_pool ~genesis_block
-                block
+                block ~bucket:mina_network_blocks_bucket ~batch_size
             in
             query_migrated_db ~f:(fun db ->
                 match%map
@@ -546,6 +547,15 @@ let () =
              ~doc:
                "NN Last global slot since genesis to include in the migration \
                 (if omitted, only canonical blocks will be migrated)"
+         and mina_network_blocks_bucket =
+           Param.flag "--mainnet-blocks-bucket"
+             ~aliases:[ "-mainnet-blocks-bucket" ]
+             Param.(required string)
+             ~doc:"Bucket with precomputed mainnet blocks"
+         and batch_size =
+           Param.flag "--batch-size" ~aliases:[ "-batch-size" ]
+             Param.(required int)
+             ~doc:"Batch size used when downloading precomputed blocks"
          in
          main ~mainnet_archive_uri ~migrated_archive_uri ~runtime_config_file
-           ~end_global_slot )))
+           ~end_global_slot ~mina_network_blocks_bucket ~batch_size )))

src/app/berkeley_migration/precomputed_block.ml

@@ -23,11 +23,10 @@ let make_batch_args ~height ~num_blocks =
   List.init (Int64.to_int_exn actual_num_blocks) ~f:(fun n ->
       sprintf "mainnet-%Ld-*.json" Int64.(start_height + Int64.of_int n) )
 
-let fetch_batch ~height ~num_blocks =
+let fetch_batch ~height ~num_blocks ~bucket =
   let batch_args = make_batch_args ~height ~num_blocks in
   let block_uris =
-    List.map batch_args ~f:(fun arg ->
-        sprintf "gs://mina_network_block_data/%s" arg )
+    List.map batch_args ~f:(fun arg -> sprintf "gs://%s/%s" bucket arg)
   in
   match%map
     Process.run ~prog:"gsutil" ~args:([ "-m"; "cp" ] @ block_uris @ [ "." ]) ()
@@ -48,12 +47,14 @@ let delete_fetched () : unit Deferred.t =
     Array.filter files ~f:(fun file -> Str.string_match block_re file 0)
   in
   let args = Array.to_list block_files in
-  match%map Process.run ~prog:"rm" ~args () with
-  | Ok _ ->
-      ()
-  | Error err ->
-      failwithf "Could not delete fetched precomputed blocks, error %s"
-        (Error.to_string_hum err) ()
+  if List.length args > 0 then
+    match%map Process.run ~prog:"rm" ~args () with
+    | Ok _ ->
+        ()
+    | Error err ->
+        failwithf "Could not delete fetched precomputed blocks, error %s"
+          (Error.to_string_hum err) ()
+  else Deferred.unit
 
 let get_json_item filter ~state_hash ~height =
   let target = make_target ~state_hash ~height in