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 """Return raw ``report_content`` for a research row. 

62 

63 IMPORTANT: ``report_content`` is the answer body only — the 

64 assembled "## Sources" / "## Research Metrics" sections live in 

65 ``research_resources`` and the 

66 per-research metrics tables and are stitched in by 

67 ``web.services.report_assembly_service.assemble_full_report``. 

68 

69 Do **not** use this method for any user-facing display path. 

70 Call ``assemble_full_report(research, db_session)`` instead so 

71 legacy rows (which embed sources/metrics inline in 

72 ``report_content``) and new rows render identically. Current 

73 callers of this method use the truncated content for 

74 notification teasers / summary previews where answer-only is 

75 the desired shape. 

76 """ 

77 try: 

78 research = ( 

79 self.session.query(ResearchHistory) 

80 .filter_by(id=research_id) 

81 .first() 

82 ) 

83 

84 if not research or not research.report_content: 

85 return None 

86 

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

88 

89 except Exception: 

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

91 return None 

92 

93 def get_report_with_metadata( 

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

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

96 """Return ``report_content`` + research metadata for a row. 

97 

98 Same answer-only caveat as :meth:`get_report` — see that 

99 docstring before using ``["content"]`` for any user-facing 

100 rendering path; prefer ``assemble_full_report`` instead. 

101 """ 

102 try: 

103 research = ( 

104 self.session.query(ResearchHistory) 

105 .filter_by(id=research_id) 

106 .first() 

107 ) 

108 

109 if not research or not research.report_content: 

110 return None 

111 

112 return { 

113 "content": research.report_content, 

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

115 "query": research.query, 

116 "mode": research.mode, 

117 "created_at": research.created_at, 

118 "completed_at": research.completed_at, 

119 "duration_seconds": research.duration_seconds, 

120 } 

121 

122 except Exception: 

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

124 return None 

125 

126 def list_reports( 

127 self, username: Optional[str] = None 

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

129 """List reports from database.""" 

130 try: 

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

132 ResearchHistory.report_content.isnot(None) 

133 ) 

134 results = query.all() 

135 return [ 

136 { 

137 "id": r.id, 

138 "query": r.query, 

139 "mode": r.mode, 

140 "created_at": r.created_at, 

141 "completed_at": r.completed_at, 

142 } 

143 for r in results 

144 ] 

145 except Exception: 

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

147 return [] 

148 

149 def delete_report( 

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

151 ) -> bool: 

152 """Delete report from database.""" 

153 try: 

154 research = ( 

155 self.session.query(ResearchHistory) 

156 .filter_by(id=research_id) 

157 .first() 

158 ) 

159 

160 if not research: 

161 return False 

162 

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

164 self.session.commit() 

165 

166 return True 

167 

168 except Exception: 

169 logger.exception("Error deleting report") 

170 self.session.rollback() 

171 return False 

172 

173 def report_exists( 

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

175 ) -> bool: 

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

177 try: 

178 research = ( 

179 self.session.query(ResearchHistory) 

180 .filter_by(id=research_id) 

181 .first() 

182 ) 

183 

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

185 

186 except Exception: 

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

188 return False