Coverage for src/qdrant_loader/cli/commands/ingest_cmd.py: 72%

1from __future__ import annotations 

2 

3import asyncio 

4import signal 

5import time 

6import traceback 

7from pathlib import Path 

8 

9from click.exceptions import ClickException 

10 

11from qdrant_loader.cli.async_utils import cancel_all_tasks 

12from qdrant_loader.cli.config_loader import ( 

13 load_config_with_workspace, 

14 setup_workspace, 

15) 

16from qdrant_loader.config.workspace import validate_workspace_flags 

17from qdrant_loader.utils.logging import LoggingConfig 

18from qdrant_loader.utils.sensitive import sanitize_exception_message 

19 

20from . import run_pipeline_ingestion 

21 

22# Backward-compatibility aliases for tests expecting underscored names 

23_load_config_with_workspace = load_config_with_workspace 

24_setup_workspace_impl = setup_workspace 

25_run_ingest_pipeline = run_pipeline_ingestion 

26_cancel_all_tasks_helper = cancel_all_tasks 

27 

28 

29async def run_ingest_command( 

30 workspace: Path | None, 

31 config: Path | None, 

32 env: Path | None, 

33 project: str | None, 

34 source_type: str | None, 

35 source: str | None, 

36 log_level: str, 

37 profile: bool, 

38 force: bool, 

39 resume: bool = True, 

40) -> None: 

41 """Implementation for the `ingest` CLI command with checkpoint resume support (WS-2).""" 

42 

43 ingest_start_time = time.perf_counter() 

44 

45 try: 

46 # Validate flag combinations 

47 validate_workspace_flags(workspace, config, env) 

48 

49 # Setup workspace if provided 

50 workspace_config = None 

51 if workspace: 

52 workspace_config = _setup_workspace_impl(workspace) 

53 

54 # Setup/reconfigure logging with workspace support 

55 log_file = ( 

56 str(workspace_config.logs_path / "ingest.log") 

57 if workspace_config 

58 else "qdrant-loader.log" 

59 ) 

60 if getattr(LoggingConfig, "reconfigure", None): # type: ignore[attr-defined] 

61 if getattr(LoggingConfig, "_initialized", False): # type: ignore[attr-defined] 

62 LoggingConfig.reconfigure(file=log_file, level=log_level) # type: ignore[attr-defined] 

63 else: 

64 LoggingConfig.setup(level=log_level, format="console", file=log_file) 

65 else: 

66 import logging as _py_logging 

67 

68 _py_logging.getLogger().handlers = [] 

69 LoggingConfig.setup(level=log_level, format="console", file=log_file) 

70 

71 # Load configuration 

72 _load_config_with_workspace(workspace_config, config, env) 

73 from qdrant_loader.config import get_settings 

74 

75 settings = get_settings() 

76 if settings is None: 

77 LoggingConfig.get_logger(__name__).error("settings_not_available") 

78 raise ClickException("Settings not available") 

79 

80 # Lazy import to avoid slow startup 

81 from qdrant_loader.core.qdrant_manager import QdrantManager 

82 

83 qdrant_manager = QdrantManager(settings) 

84 

85 async def _do_run(): 

86 await _run_ingest_pipeline( 

87 settings, 

88 qdrant_manager, 

89 project=project, 

90 source_type=source_type, 

91 source=source, 

92 force=force, 

93 resume=resume, 

94 metrics_dir=( 

95 str(workspace_config.metrics_path) if workspace_config else None 

96 ), 

97 ) 

98 

99 loop = asyncio.get_running_loop() 

100 stop_event = asyncio.Event() 

101 

102 def _handle_sigint(): 

103 logger = LoggingConfig.get_logger(__name__) 

104 logger.debug(" SIGINT received, cancelling all tasks...") 

105 stop_event.set() 

106 # Schedule cancellation of all running tasks safely on the event loop thread 

107 loop.call_soon_threadsafe( 

108 lambda: loop.create_task(_cancel_all_tasks_helper()) 

109 ) 

110 

111 try: 

112 loop.add_signal_handler(signal.SIGINT, _handle_sigint) 

113 except NotImplementedError: 

114 

115 def _signal_handler(_signum, _frame): 

116 logger = LoggingConfig.get_logger(__name__) 

117 logger.debug(" SIGINT received on Windows, cancelling all tasks...") 

118 loop.call_soon_threadsafe(stop_event.set) 

119 # Ensure the coroutine runs on the correct loop without race conditions 

120 asyncio.run_coroutine_threadsafe(_cancel_all_tasks_helper(), loop) 

121 

122 signal.signal(signal.SIGINT, _signal_handler) 

123 

124 try: 

125 if profile: 

126 import cProfile 

127 

128 profiler = cProfile.Profile() 

129 profiler.enable() 

130 try: 

131 await _do_run() 

132 finally: 

133 profiler.disable() 

134 profiler.dump_stats("profile.out") 

135 LoggingConfig.get_logger(__name__).info( 

136 "Profile saved to profile.out" 

137 ) 

138 else: 

139 await _do_run() 

140 

141 logger = LoggingConfig.get_logger(__name__) 

142 logger.info("Pipeline finished, awaiting cleanup.") 

143 pending = [ 

144 t 

145 for t in asyncio.all_tasks() 

146 if t is not asyncio.current_task() and not t.done() 

147 ] 

148 if pending: 

149 logger.debug(f" Awaiting {len(pending)} pending tasks before exit...") 

150 results = await asyncio.gather(*pending, return_exceptions=True) 

151 for idx, result in enumerate(results): 

152 if isinstance(result, Exception): 

153 logger.error( 

154 "Pending task failed during shutdown", 

155 task_index=idx, 

156 error=sanitize_exception_message(result), 

157 error_type=type(result).__name__, 

158 ) 

159 await asyncio.sleep(0.1) 

160 end_to_end_duration = time.perf_counter() - ingest_start_time 

161 logger.info( 

162 f"Ingestion end-to-end completed in {end_to_end_duration:.2f} seconds" 

163 ) 

164 except asyncio.CancelledError: 

165 # Preserve cancellation semantics so Ctrl+C results in a normal exit 

166 raise 

167 except Exception as e: 

168 logger = LoggingConfig.get_logger(__name__) 

169 error_msg = ( 

170 sanitize_exception_message(e) 

171 or f"Empty exception of type: {type(e).__name__}" 

172 ) 

173 sanitized_traceback = sanitize_exception_message(traceback.format_exc()) 

174 end_to_end_duration = time.perf_counter() - ingest_start_time 

175 logger.error( 

176 "Document ingestion process failed during execution", 

177 error=error_msg, 

178 error_type=type(e).__name__, 

179 sanitized_traceback=sanitized_traceback, 

180 end_to_end_duration_seconds=round(end_to_end_duration, 2), 

181 suggestion=( 

182 "Check data sources, configuration, and system resources. " 

183 "Run 'qdrant-loader project validate' to verify setup" 

184 ), 

185 ) 

186 raise ClickException(f"Failed to run ingestion: {error_msg}") from e 

187 finally: 

188 if stop_event.is_set(): 

189 logger = LoggingConfig.get_logger(__name__) 

190 logger.debug( 

191 " Cancellation already initiated by SIGINT; exiting gracefully." 

192 ) 

193 

194 except asyncio.CancelledError: 

195 # Bubble up cancellation to the caller/CLI, do not convert to ClickException 

196 raise 

197 except ClickException: 

198 raise 

199 except Exception as e: 

200 logger = LoggingConfig.get_logger(__name__) 

201 error_msg = ( 

202 sanitize_exception_message(e) 

203 or f"Empty exception of type: {type(e).__name__}" 

204 ) 

205 sanitized_traceback = sanitize_exception_message(traceback.format_exc()) 

206 end_to_end_duration = time.perf_counter() - ingest_start_time 

207 logger.error( 

208 "Unexpected error during ingestion command execution", 

209 error=error_msg, 

210 error_type=type(e).__name__, 

211 sanitized_traceback=sanitized_traceback, 

212 end_to_end_duration_seconds=round(end_to_end_duration, 2), 

213 suggestion="Check logs above for specific error details and verify system configuration", 

214 ) 

215 raise ClickException(f"Failed to run ingestion: {error_msg}") from e