Coverage for src/usaspending/cli/download_award.py: 0%

1# src/usaspending/cli/download_award.py 

2 

3import argparse 

4import sys 

5import os 

6from typing import List 

7 

8# Helper to ensure the package root is in sys.path if running the script directly during development 

9if __name__ == "__main__": 

10 # Adjust path relative to the script location (src/usaspending/cli) 

11 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) 

12 if project_root not in sys.path: 

13 sys.path.insert(0, project_root) 

14 

15# Import necessary components from the library 

16from usaspending.client import USASpending 

17from usaspending.exceptions import DownloadError 

18from usaspending.logging_config import USASpendingLogger 

19 

20def main(): 

21 parser = argparse.ArgumentParser( 

22 description="Download detailed award data from USASpending.gov.", 

23 epilog="This CLI tool queues a download job, waits for completion, downloads the zip file, and extracts the contents." 

24 ) 

25 

26 parser.add_argument("award_id", help="The unique award identifier (e.g., PIIN/FAIN/etc.)") 

27 parser.add_argument("-o", "--output-dir", help="The directory to save and extract the files (defaults to current directory).") 

28 parser.add_argument("-f", "--format", choices=["csv", "tsv", "pstxt"], default="csv", help="The format of the files (default: csv).") 

29 parser.add_argument("--timeout", type=int, default=1800, help="Maximum time in seconds to wait (default: 1800s/30min).") 

30 parser.add_argument("--poll-interval", type=int, default=30, help="Interval in seconds between status checks (default: 30s).") 

31 parser.add_argument("--no-cleanup", action="store_true", help="Keep the zip file after extraction.") 

32 parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose debug logging.") 

33 

34 args = parser.parse_args() 

35 

36 # Configure logging 

37 log_level = "DEBUG" if args.verbose else "INFO" 

38 USASpendingLogger.configure(level=log_level, debug_mode=args.verbose) 

39 logger = USASpendingLogger.get_logger("usaspending.cli.download") 

40 

41 try: 

42 # Initialize client (assuming default configuration handling) 

43 client = USASpending() 

44 

45 # Get award data 

46 award = client.awards.find_by_award_id(args.award_id) 

47 

48 logger.info("--- Starting USASpending Award Download CLI ---") 

49 logger.info(f"Parameters: Award ID={args.award_id}, Type={award.category}, Format={args.format}") 

50 

51 job = award.download(file_format=args.format, destination_dir=args.output_dir) 

52 

53 logger.info(f"Job successfully queued. Tracking File: {job.file_name}") 

54 

55 # 2. Wait for completion (blocking operation) 

56 extracted_files: List[str] = job.wait_for_completion( 

57 timeout=args.timeout, 

58 poll_interval=args.poll_interval, 

59 cleanup_zip=not args.no_cleanup 

60 ) 

61 

62 logger.info("--- Download and Extraction Complete ---") 

63 # Determine the final extraction path 

64 extract_subdir_name = os.path.splitext(job.file_name)[0] 

65 final_path = os.path.join(job.destination_dir, extract_subdir_name) 

66 logger.info(f"Data extracted to: {os.path.abspath(final_path)}") 

67 logger.info(f"Total files extracted: {len(extracted_files)}") 

68 

69 except DownloadError as e: 

70 logger.error(f"Download failed. Status: {e.status}. Message: {e}") 

71 sys.exit(1) 

72 except KeyboardInterrupt: 

73 logger.warning("Download interrupted by user.") 

74 sys.exit(1) 

75 except Exception as e: 

76 logger.exception(f"An unexpected error occurred: {e}") 

77 sys.exit(1) 

78 

79if __name__ == "__main__": 

80 main()