modernize

src/bfabric_scripts/bfabric_list_workunit_parameters.py

@@ -1,11 +1,12 @@
 import argparse
 import json
-import sys
 
 import polars as pl
 import rich
 
 from bfabric import Bfabric
+from bfabric.cli_formatting import setup_script_logging
+from bfabric.experimental import MultiQuery
 
 
 def bfabric_list_workunit_parameters(client: Bfabric, application_id: int, max_workunits: int, format: str) -> None:
@@ -74,16 +75,13 @@ def print_results(format: str, merged_result: pl.DataFrame) -> None:
 def get_parameter_table(client: Bfabric, workunits_table_explode: pl.DataFrame) -> pl.DataFrame:
     """Returns a wide format table for the specified parameters, with the key `workunit_id` indicating the source."""
     # load the parameters table
-    collect = []
-    for i_frame, frame in enumerate(workunits_table_explode.iter_slices(100)):
-        print(
-            f"-- Reading parameters chunk {i_frame + 1} of {len(workunits_table_explode) // 100 + 1}", file=sys.stderr
-        )
-        chunk = (
-            client.read("parameter", {"id": frame["parameter_id"].to_list()}).to_polars().rename({"id": "parameter_id"})
-        )
-        collect.append(chunk)
-    parameter_table_full = pl.concat(collect, how="align")[["parameter_id", "key", "value"]]
+    collect = MultiQuery(client=client).read_multi(
+        endpoint="parameter",
+        obj={},
+        multi_query_key="id",
+        multi_query_vals=workunits_table_explode["parameter_id"].to_list(),
+    )
+    parameter_table_full = collect.to_polars().rename({"id": "parameter_id"})[["parameter_id", "key", "value"]]
     # add workunit id to parameter table
     parameter_table_full = parameter_table_full.join(
         workunits_table_explode[["workunit_id", "parameter_id"]], on="parameter_id", how="left"
@@ -94,6 +92,7 @@ def get_parameter_table(client: Bfabric, workunits_table_explode: pl.DataFrame)
 
 def main() -> None:
     """Parses command line arguments and calls `bfabric_list_workunit_parameters`."""
+    setup_script_logging()
     client = Bfabric.from_config()
     parser = argparse.ArgumentParser()
     parser.add_argument("application_id", type=int, help="The application ID to list the workunit parameters for.")