Coverage for src/qdrant_loader/core/state/utils.py: 63%

1""" 

2Utilities for StateManager: database URL construction and common query builders. 

3""" 

4 

5from __future__ import annotations 

6 

7import os 

8from pathlib import Path 

9from typing import TYPE_CHECKING 

10 

11if TYPE_CHECKING: 

12 from qdrant_loader.config.state import StateManagementConfig 

13 

14from sqlalchemy import select 

15 

16from qdrant_loader.core.state.exceptions import DatabaseError 

17from qdrant_loader.core.state.models import DocumentStateRecord, IngestionHistory 

18 

19 

20def ensure_parent_directory(db_path: Path) -> None: 

21 """Ensure the parent directory of the database file exists and is writable.""" 

22 parent_dir = db_path.parent 

23 if not parent_dir.exists(): 

24 try: 

25 parent_dir.mkdir(parents=True, exist_ok=True) 

26 except Exception as e: # pragma: no cover - safety net 

27 raise DatabaseError( 

28 f"Cannot create database directory {parent_dir}: {e}" 

29 ) from e 

30 if not os.access(parent_dir, os.W_OK): 

31 raise DatabaseError(f"No write permission for database directory: {parent_dir}") 

32 

33 

34def generate_sqlite_aiosqlite_url(database_path: str) -> str: 

35 """Generate an aiosqlite URL from a configured database path string. 

36 

37 Supports special values like ":memory:" and already-prefixed sqlite URLs. 

38 Ensures parent directory exists for file-backed databases. 

39 """ 

40 if database_path in (":memory:", "sqlite:///:memory:", "sqlite://:memory:"): 

41 return "sqlite+aiosqlite:///:memory:" 

42 

43 if database_path.startswith("sqlite://"): 

44 # Convert to aiosqlite dialect 

45 return database_path.replace("sqlite://", "sqlite+aiosqlite://") 

46 

47 # Treat as filesystem path 

48 db_path = Path(database_path) 

49 if not db_path.is_absolute(): 

50 db_path = db_path.resolve() 

51 

52 ensure_parent_directory(db_path) 

53 

54 # Normalize to POSIX path for SQLAlchemy URL 

55 db_url_path = db_path.as_posix() 

56 # Absolute and relative are handled similarly here (three slashes) 

57 return f"sqlite+aiosqlite:///{db_url_path}" 

58 

59 

60def _normalize_async_url(url: str) -> str: 

61 """Ensure a database URL uses an async driver SQLAlchemy can use. 

62 

63 - postgres:// / postgresql:// -> postgresql+asyncpg:// (so a standard RDS/psql 

64 URL works unchanged) 

65 - sqlite:// -> sqlite+aiosqlite:// 

66 - anything already carrying a driver (e.g. postgresql+asyncpg://) passes through. 

67 """ 

68 if url.startswith(("postgresql://", "postgres://")): 

69 return "postgresql+asyncpg://" + url.split("://", 1)[1] 

70 if url.startswith("sqlite://") and "+aiosqlite" not in url: 

71 return url.replace("sqlite://", "sqlite+aiosqlite://", 1) 

72 return url 

73 

74 

75def generate_database_url(config: StateManagementConfig) -> str: 

76 """Build the async SQLAlchemy URL for the configured state backend. 

77 

78 Precedence: ``config.database_url`` (Postgres or any full SQLAlchemy URL) wins 

79 and selects the backend by dialect; otherwise fall back to the SQLite 

80 ``database_path`` (community default). 

81 """ 

82 database_url = getattr(config, "database_url", None) 

83 if database_url: 

84 return _normalize_async_url(database_url) 

85 return generate_sqlite_aiosqlite_url(config.database_path) 

86 

87 

88def build_ingestion_history_select( 

89 source_type: str, 

90 source: str, 

91 project_id: str | None = None, 

92 order_by_last_successful_desc: bool = False, 

93): 

94 """Create a select() for IngestionHistory with optional project filter and ordering.""" 

95 query = select(IngestionHistory).filter( 

96 IngestionHistory.source_type == source_type, IngestionHistory.source == source 

97 ) 

98 if project_id is not None: 

99 query = query.filter(IngestionHistory.project_id == project_id) 

100 if order_by_last_successful_desc: 

101 query = query.order_by(IngestionHistory.last_successful_ingestion.desc()) 

102 return query 

103 

104 

105def build_document_state_select( 

106 source_type: str, 

107 source: str, 

108 document_id: str | None = None, 

109 project_id: str | None = None, 

110): 

111 """Create a select() for DocumentStateRecord with optional project/doc filters.""" 

112 conditions = [ 

113 DocumentStateRecord.source_type == source_type, 

114 DocumentStateRecord.source == source, 

115 ] 

116 if document_id is not None: 

117 conditions.append(DocumentStateRecord.document_id == document_id) 

118 query = select(DocumentStateRecord).filter(*conditions) 

119 if project_id is not None: 

120 query = query.filter(DocumentStateRecord.project_id == project_id) 

121 return query