Coverage for src/qdrant_loader/connectors/base.py: 95%

1import os 

2import warnings 

3from collections.abc import AsyncIterator 

4from datetime import datetime 

5 

6from qdrant_loader.config.source_config import SourceConfig 

7from qdrant_loader.core.document import Document 

8from qdrant_loader.core.file_conversion import FileConversionConfig 

9 

10 

11def resolve_safe_path(base_dir: str, entity_id: str) -> str | None: 

12 """Resolve ``entity_id`` (a ``/``-separated relative path) under ``base_dir``. 

13 

14 Returns the resolved absolute path, or ``None`` if it would escape 

15 ``base_dir`` (e.g. via ``../`` segments), guarding against path traversal. 

16 """ 

17 candidate = os.path.join(base_dir, *entity_id.split("/")) 

18 base_real = os.path.realpath(base_dir) 

19 candidate_real = os.path.realpath(candidate) 

20 if candidate_real != base_real and not candidate_real.startswith( 

21 base_real + os.sep 

22 ): 

23 return None 

24 return candidate 

25 

26 

27class ConnectorConfigurationError(Exception): 

28 """Raised when a connector's configuration is invalid or access is denied. 

29 

30 This is a *fatal* error: the pipeline should stop rather than silently 

31 continuing with 0 documents. 

32 """ 

33 

34 

35class BaseConnector: 

36 """Base class for all connectors.""" 

37 

38 def __init__(self, config: SourceConfig): 

39 self.config = config 

40 self._initialized = False 

41 

42 async def __aenter__(self): 

43 """Async context manager entry.""" 

44 self._initialized = True 

45 return self 

46 

47 async def __aexit__(self, exc_type, exc_val, _exc_tb): 

48 """Async context manager exit.""" 

49 self._initialized = False 

50 

51 def set_file_conversion_config( 

52 self, file_conversion_config: FileConversionConfig 

53 ) -> None: 

54 """Set file conversion configuration. 

55 

56 This default implementation stores the configuration for potential 

57 use by subclasses that choose to honor it. 

58 

59 Args: 

60 file_conversion_config: Global file conversion configuration 

61 """ 

62 # Store on the instance so connectors that opt-in can access it. 

63 self._file_conversion_config = file_conversion_config 

64 

65 async def stream_documents( 

66 self, since: datetime | None = None 

67 ) -> AsyncIterator[Document]: 

68 """Stream documents from the source (WS-1 connector contract). 

69 

70 Connectors must implement true streaming. This default raises 

71 NotImplementedError to prevent silently materializing the full 

72 document list via get_documents(). 

73 """ 

74 if False: # pragma: no cover - makes this function an async generator 

75 yield "" # type: ignore 

76 raise NotImplementedError( 

77 f"{type(self).__name__} does not implement stream_documents" 

78 ) 

79 

80 async def get_documents(self) -> list[Document]: 

81 """Get documents from the source (DEPRECATED - use stream_documents).""" 

82 warnings.warn( 

83 "BaseConnector.get_documents is deprecated. Implement stream_documents() " 

84 "or use connector.stream_documents() to avoid materializing the full " 

85 "document list in memory.", 

86 DeprecationWarning, 

87 stacklevel=2, 

88 ) 

89 documents: list[Document] = [] 

90 async for document in self.stream_documents(): 

91 documents.append(document) 

92 return documents 

93 

94 async def fetch_by_id(self, entity_id: str) -> Document | None: 

95 """Fetch a single entity by ID (WS-1 connector contract). 

96 

97 Connectors that support single-event ingestion must override this method. 

98 """ 

99 raise NotImplementedError(f"{type(self).__name__} does not support fetch_by_id") 

100 

101 async def list_entity_ids(self) -> AsyncIterator[str]: 

102 """Stream all entity IDs for reconciliation (WS-1 connector contract). 

103 

104 Yields: 

105 Entity IDs from the source. 

106 """ 

107 # Make this an async generator so callers can use `async for` and 

108 # receive a NotImplementedError during iteration rather than at call-time. 

109 if False: # pragma: no cover - makes this function an async generator 

110 yield "" 

111 raise NotImplementedError( 

112 f"{type(self).__name__} does not support list_entity_ids" 

113 )