Coverage for src / local_deep_research / storage / database.py: 96%

1"""Database-based report storage implementation.""" 

2 

3from typing import Dict, Any, List, Optional 

4from loguru import logger 

5from sqlalchemy.orm import Session 

6 

7from .base import ReportStorage 

8from ..database.models import ResearchHistory 

9 

10 

11class DatabaseReportStorage(ReportStorage): 

12 """Store reports in the database with caching support.""" 

13 

14 def __init__(self, session: Session): 

15 """Initialize database storage. 

16 

17 Args: 

18 session: SQLAlchemy database session 

19 """ 

20 self.session = session 

21 

22 def save_report( 

23 self, 

24 research_id: str, 

25 content: str, 

26 metadata: Optional[Dict[str, Any]] = None, 

27 username: Optional[str] = None, 

28 ) -> bool: 

29 """Save report to database.""" 

30 try: 

31 research = ( 

32 self.session.query(ResearchHistory) 

33 .filter_by(id=research_id) 

34 .first() 

35 ) 

36 

37 if not research: 

38 logger.error(f"Research {research_id} not found") 

39 return False 

40 

41 research.report_content = content # type: ignore[assignment] 

42 

43 if metadata: 

44 if research.research_meta: 

45 research.research_meta.update(metadata) # type: ignore[union-attr] 

46 else: 

47 research.research_meta = metadata # type: ignore[assignment] 

48 

49 self.session.commit() 

50 logger.info(f"Saved report for research {research_id} to database") 

51 return True 

52 

53 except Exception: 

54 logger.exception("Error saving report to database") 

55 self.session.rollback() 

56 return False 

57 

58 def get_report( 

59 self, research_id: str, username: Optional[str] = None 

60 ) -> Optional[str]: 

61 """Get report from database.""" 

62 try: 

63 research = ( 

64 self.session.query(ResearchHistory) 

65 .filter_by(id=research_id) 

66 .first() 

67 ) 

68 

69 if not research or not research.report_content: 

70 return None 

71 

72 return research.report_content # type: ignore[return-value] 

73 

74 except Exception: 

75 logger.exception("Error getting report from database") 

76 return None 

77 

78 def get_report_with_metadata( 

79 self, research_id: str, username: Optional[str] = None 

80 ) -> Optional[Dict[str, Any]]: 

81 """Get report with metadata from database.""" 

82 try: 

83 research = ( 

84 self.session.query(ResearchHistory) 

85 .filter_by(id=research_id) 

86 .first() 

87 ) 

88 

89 if not research or not research.report_content: 

90 return None 

91 

92 return { 

93 "content": research.report_content, 

94 "metadata": research.research_meta or {}, 

95 "query": research.query, 

96 "mode": research.mode, 

97 "created_at": research.created_at, 

98 "completed_at": research.completed_at, 

99 "duration_seconds": research.duration_seconds, 

100 } 

101 

102 except Exception: 

103 logger.exception("Error getting report with metadata") 

104 return None 

105 

106 def list_reports( 

107 self, username: Optional[str] = None 

108 ) -> List[Dict[str, Any]]: 

109 """List reports from database.""" 

110 try: 

111 query = self.session.query(ResearchHistory).filter( 

112 ResearchHistory.report_content.isnot(None) 

113 ) 

114 results = query.all() 

115 return [ 

116 { 

117 "id": r.id, 

118 "query": r.query, 

119 "mode": r.mode, 

120 "created_at": r.created_at, 

121 "completed_at": r.completed_at, 

122 } 

123 for r in results 

124 ] 

125 except Exception: 

126 logger.exception("Error listing reports from database") 

127 return [] 

128 

129 def delete_report( 

130 self, research_id: str, username: Optional[str] = None 

131 ) -> bool: 

132 """Delete report from database.""" 

133 try: 

134 research = ( 

135 self.session.query(ResearchHistory) 

136 .filter_by(id=research_id) 

137 .first() 

138 ) 

139 

140 if not research: 

141 return False 

142 

143 research.report_content = None # type: ignore[assignment] 

144 self.session.commit() 

145 

146 return True 

147 

148 except Exception: 

149 logger.exception("Error deleting report") 

150 self.session.rollback() 

151 return False 

152 

153 def report_exists( 

154 self, research_id: str, username: Optional[str] = None 

155 ) -> bool: 

156 """Check if report exists in database.""" 

157 try: 

158 research = ( 

159 self.session.query(ResearchHistory) 

160 .filter_by(id=research_id) 

161 .first() 

162 ) 

163 

164 return research is not None and research.report_content is not None 

165 

166 except Exception: 

167 logger.exception("Error checking if report exists") 

168 return False