Coverage for lunar_toolkit/cli/workspace.py: 91%

1"""workspace - ad-hoc SQL on a work dir's DuckDB workspace (lunar.duckdb).""" 

2 

3import argparse 

4import sys 

5from pathlib import Path 

6 

7COMMAND = "workspace" 

8 

9 

10def add_parser(sub: argparse._SubParsersAction) -> None: 

11 p = sub.add_parser( 

12 COMMAND, 

13 help="Query a work dir's DuckDB workspace (images, pairs, bundle-adjustment tables and views).", 

14 formatter_class=argparse.RawDescriptionHelpFormatter, 

15 epilog=""" 

16Examples 

17-------- 

18 lunar workspace work/malapert # tables and views with row counts 

19 lunar workspace work/malapert --sql "SELECT * FROM ba_run_summary" 

20 lunar workspace work/malapert --sql "SELECT run, rule, count(*) FROM image_rejections GROUP BY ALL" 

21 lunar workspace work/malapert --sql "SELECT * FROM image_rejections_at(1.5, 0.8, 1.0, 1)" 

22 lunar workspace work/malapert --export ba_image_stats ba_stats.parquet 

23 

24WORK_DIR is a lunar sfs work directory (its lunar.duckdb) or a .duckdb file; it is 

25opened read-only. Tables: images, pairs, ba_runs, ba_residual_stats, 

26ba_pixel_residuals, ba_match_offsets, ba_points, ba_convergence, ba_camera_offsets, 

27rejection_params. Views: ba_image_stats, ba_run_summary, image_rejections. 

28--export writes Parquet for .parquet paths, else CSV (geometries as WKT). 

29""", 

30 ) 

31 p.add_argument("work_dir", type=Path, metavar="WORK_DIR", help="Work directory or .duckdb file.") 

32 p.add_argument("--sql", action="append", metavar="SQL", help="Run SQL and print the result (repeatable).") 

33 p.add_argument("--tables", action="store_true", help="List tables and views with row counts.") 

34 p.add_argument( 

35 "--export", nargs=2, metavar=("TABLE", "PATH"), help="Write a table or view to Parquet / CSV." 

36 ) 

37 p.add_argument("--csv", metavar="FILE", help="Write the last --sql result to CSV.") 

38 p.add_argument("--max-rows", type=int, default=50, metavar="N", help="Rows to print (default: 50).") 

39 p.add_argument( 

40 "--memory-limit", default="1GB", metavar="SIZE", help="DuckDB memory_limit (default: 1GB)." 

41 ) 

42 p.add_argument("--threads", type=int, metavar="N", help="DuckDB threads (default: all cores).") 

43 

44 

45def run(args: argparse.Namespace) -> int: 

46 import duckdb 

47 

48 from ..workspace import Workspace, WorkspaceError, workspace_path 

49 

50 path = workspace_path(args.work_dir) 

51 if not path.exists(): 

52 print(f"Error: no workspace at {path}", file=sys.stderr) 

53 return 1 

54 try: 

55 ws = Workspace(path, read_only=True, memory_limit=args.memory_limit, threads=args.threads) 

56 except duckdb.Error as exc: 

57 print(f"Error: {exc}", file=sys.stderr) 

58 return 1 

59 with ws: 

60 try: 

61 if args.tables or not (args.sql or args.export): 

62 print(ws.tables().to_string(index=False)) 

63 res = None 

64 for sql in args.sql or []: 

65 res = ws.query(sql) 

66 geom = [ 

67 c for c in res.columns if res[c].map(lambda v: isinstance(v, bytes | bytearray)).any() 

68 ] 

69 shown = res.drop(columns=geom) 

70 if len(shown.columns): 70 ↛ 64line 70 didn't jump to line 64 because the condition on line 70 was always true

71 with_more = len(res) > args.max_rows 

72 print(shown.head(args.max_rows).to_string(index=False)) 

73 if with_more: 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true

74 print(f"... {len(res)} rows") 

75 if args.csv and res is not None: 

76 res.to_csv(args.csv, index=False) 

77 print(f"Wrote {args.csv}") 

78 if args.export: 

79 out = ws.export(args.export[0], args.export[1]) 

80 print(f"Wrote {out}") 

81 except (duckdb.Error, WorkspaceError, ValueError) as exc: 

82 print(f"Error: {exc}", file=sys.stderr) 

83 return 1 

84 return 0