Coverage for src/qdrant_loader/core/pipeline/source_processor.py: 62%
77 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Source processor for handling different source types."""
3import asyncio
4import inspect
5from collections.abc import AsyncIterator, Callable, Mapping
6from datetime import datetime
8from qdrant_loader.config.source_config import SourceConfig
9from qdrant_loader.connectors.base import BaseConnector, ConnectorConfigurationError
10from qdrant_loader.core.document import Document
11from qdrant_loader.core.file_conversion import FileConversionConfig
12from qdrant_loader.utils.logging import LoggingConfig
13from qdrant_loader.utils.sensitive import sanitize_exception_message
15logger = LoggingConfig.get_logger(__name__)
18class SourceProcessor:
19 """Handles processing of different source types."""
21 def __init__(
22 self,
23 shutdown_event: asyncio.Event | None = None,
24 file_conversion_config: FileConversionConfig | None = None,
25 ):
26 self.shutdown_event = shutdown_event or asyncio.Event()
27 self.file_conversion_config = file_conversion_config
29 async def process_source_type(
30 self,
31 source_configs: Mapping[str, SourceConfig],
32 connector_factory: Callable[[SourceConfig], BaseConnector],
33 source_type: str,
34 ) -> list[Document]:
35 """Process documents from a specific source type.
37 Args:
38 source_configs: Mapping of source name to source configuration
39 connector_factory: Factory function that creates a connector from a source config
40 source_type: The type of source being processed
42 Returns:
43 List of documents from all sources of this type
44 """
45 logger.debug(f"Processing {source_type} sources: {list(source_configs.keys())}")
47 all_documents = []
49 for source_name, source_config in source_configs.items():
50 if self.shutdown_event.is_set():
51 logger.info(
52 f"Shutdown requested, skipping {source_type} source: {source_name}"
53 )
54 break
56 try:
57 logger.debug(f"Processing {source_type} source: {source_name}")
59 # Create connector instance and use as async context manager
60 maybe_connector = connector_factory(source_config)
61 connector = (
62 await maybe_connector
63 if asyncio.iscoroutine(maybe_connector)
64 or inspect.isawaitable(maybe_connector)
65 else maybe_connector
66 )
68 # Set file conversion config if available and connector supports it
69 if (
70 self.file_conversion_config
71 and hasattr(connector, "set_file_conversion_config")
72 and hasattr(source_config, "enable_file_conversion")
73 and source_config.enable_file_conversion
74 ):
75 logger.debug(
76 f"Setting file conversion config for {source_type} source: {source_name}"
77 )
78 connector.set_file_conversion_config(self.file_conversion_config)
80 # Use the connector as an async context manager to ensure proper initialization
81 async with connector:
82 documents = []
83 if isinstance(connector, BaseConnector):
84 try:
85 async for document in connector.stream_documents():
86 documents.append(document)
87 except NotImplementedError:
88 # Connector does not implement streaming; fall back to eager fetch
89 documents = await connector.get_documents()
90 else:
91 documents = await connector.get_documents()
93 logger.debug(
94 f"Retrieved {len(documents)} documents from {source_type} source: {source_name}"
95 )
96 all_documents.extend(documents)
98 except ConnectorConfigurationError:
99 # Fatal configuration error – re-raise immediately so the
100 # pipeline stops with a clear message instead of silently
101 # producing 0 documents.
102 raise
103 except Exception as e:
104 safe_error = sanitize_exception_message(e)
105 logger.error(
106 f"Failed to process {source_type} source {source_name}: {safe_error}",
107 error_type=type(e).__name__,
108 )
109 # Continue processing other sources even if one fails
110 continue
112 if all_documents:
113 logger.info(
114 f"📥 {source_type}: {len(all_documents)} documents from {len(source_configs)} sources"
115 )
116 return all_documents
118 async def stream_source_documents(
119 self,
120 source_configs: Mapping[str, SourceConfig],
121 connector_factory: Callable[[SourceConfig], BaseConnector],
122 source_type: str,
123 since: datetime | None = None,
124 ) -> AsyncIterator[Document]:
125 """Stream documents from a specific source type (WS-1).
127 Yields documents one at a time as they are fetched from the source.
128 """
129 logger.debug(f"Streaming {source_type} sources: {list(source_configs.keys())}")
131 for source_name, source_config in source_configs.items():
132 if self.shutdown_event.is_set():
133 logger.info(
134 f"Shutdown requested, skipping {source_type} source: {source_name}"
135 )
136 break
138 try:
139 logger.debug(f"Streaming {source_type} source: {source_name}")
141 maybe_connector = connector_factory(source_config)
142 connector = (
143 await maybe_connector
144 if asyncio.iscoroutine(maybe_connector)
145 or inspect.isawaitable(maybe_connector)
146 else maybe_connector
147 )
149 if (
150 self.file_conversion_config
151 and hasattr(connector, "set_file_conversion_config")
152 and hasattr(source_config, "enable_file_conversion")
153 and source_config.enable_file_conversion
154 ):
155 logger.debug(
156 f"Setting file conversion config for {source_type} source: {source_name}"
157 )
158 connector.set_file_conversion_config(self.file_conversion_config)
160 async with connector:
161 document_count = 0
162 async for document in connector.stream_documents(since=since):
163 yield document
164 document_count += 1
165 if document_count % 100 == 0:
166 logger.debug(
167 f"Streamed {document_count} documents from {source_type} source: {source_name}"
168 )
170 if document_count > 0:
171 logger.info(
172 f"✅ Streamed {document_count} documents from {source_type} source: {source_name}"
173 )
175 except ConnectorConfigurationError:
176 raise
177 except Exception as e:
178 safe_error = sanitize_exception_message(e)
179 logger.error(
180 f"Failed to stream {source_type} source {source_name}: {safe_error}",
181 error_type=type(e).__name__,
182 )
183 continue