Coverage for lunar_toolkit/cli/build_db.py: 37%

1"""build-db - download CUMINDEX files and build local Parquet databases.""" 

2 

3import argparse 

4import sys 

5from pathlib import Path 

6 

7from ..constants import DATASETS, DEFAULT_CACHE_DIR 

8 

9COMMAND = "build-db" 

10 

11 

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

13 p = sub.add_parser( 

14 COMMAND, 

15 help="Build the local Parquet index database(s).", 

16 formatter_class=argparse.RawDescriptionHelpFormatter, 

17 epilog=""" 

18Examples 

19-------- 

20 lunar build-db # EDR only (default) 

21 lunar build-db edr cdr # EDR and CDR 

22 lunar build-db --all # EDR, CDR, and RDR 

23 lunar build-db edr --force # rebuild even if exists 

24 lunar build-db edr --update # append PDS volumes newer than the built release 

25 lunar build-db edr --release LROLRC_0066C --volume-only # one volume (~20k rows) 

26""", 

27 ) 

28 p.add_argument( 

29 "datasets", 

30 nargs="*", 

31 metavar="DATASET", 

32 help="Datasets to build: edr, cdr, rdr (default: edr if none given).", 

33 ) 

34 p.add_argument( 

35 "--all", 

36 action="store_true", 

37 help="Build all three datasets: edr, cdr, rdr.", 

38 ) 

39 p.add_argument( 

40 "--force", 

41 action="store_true", 

42 help="Rebuild even if the Parquet already exists.", 

43 ) 

44 p.add_argument( 

45 "--update", 

46 action="store_true", 

47 help="Append newer per-volume INDEX.TABs to an existing database.", 

48 ) 

49 p.add_argument( 

50 "--release", 

51 default=None, 

52 metavar="VOLUME", 

53 help="Build from this PDS volume, e.g. LROLRC_0066C (default: newest on the server).", 

54 ) 

55 p.add_argument( 

56 "--volume-only", 

57 action="store_true", 

58 help="Use the volume's own INDEX/INDEX.TAB instead of the cumulative CUMINDEX.TAB.", 

59 ) 

60 p.add_argument( 

61 "--wkt", 

62 action="store_true", 

63 help="Also store an FP_WKT footprint polygon column.", 

64 ) 

65 p.add_argument( 

66 "--memory-limit", 

67 default="4GB", 

68 metavar="SIZE", 

69 help="DuckDB memory limit; sorting spills to disk beyond it (default: 4GB).", 

70 ) 

71 p.add_argument( 

72 "--threads", 

73 type=int, 

74 default=None, 

75 metavar="N", 

76 help="DuckDB worker threads (default: all cores).", 

77 ) 

78 p.add_argument( 

79 "--data-dir", 

80 type=Path, 

81 default=None, 

82 metavar="DIR", 

83 help="Where to write Parquet files (default: LUNAR_DATA_DIR or the user data dir).", 

84 ) 

85 p.add_argument( 

86 "--cache-dir", 

87 default=str(DEFAULT_CACHE_DIR), 

88 metavar="DIR", 

89 help=f"Where to cache raw TAB/LBL downloads (default: {DEFAULT_CACHE_DIR}).", 

90 ) 

91 

92 

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

94 from ..lroc.cumindex import CumIndex 

95 

96 if args.all: 

97 datasets = list(DATASETS) 

98 elif args.datasets: 

99 invalid = [d for d in args.datasets if d not in DATASETS] 

100 if invalid: 

101 print( 

102 f"Unknown dataset(s): {', '.join(invalid)}. Choose from: {', '.join(DATASETS)}", 

103 file=sys.stderr, 

104 ) 

105 return 1 

106 datasets = args.datasets 

107 else: 

108 datasets = ["edr"] 

109 

110 cache_dir = Path(args.cache_dir) 

111 

112 for dataset in datasets: 

113 ds = DATASETS[dataset] 

114 print(f"\n{ds['label']}") 

115 ci = CumIndex(dataset=dataset, data_dir=args.data_dir, cache_dir=cache_dir) 

116 print(f" Parquet : {ci.parquet_path}") 

117 print(f" Cache : {cache_dir}") 

118 if args.update and ci.parquet_path.exists() and not args.force: 

119 ci.update_db(memory_limit=args.memory_limit, threads=args.threads) 

120 continue 

121 ci.build_db( 

122 force=args.force, 

123 release=args.release, 

124 volume_only=args.volume_only, 

125 memory_limit=args.memory_limit, 

126 threads=args.threads, 

127 wkt=args.wkt, 

128 ) 

129 

130 return 0