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

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

2 

3from typing import Dict, Any, 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 

42 

43 if metadata: 

44 if research.research_meta: 

45 research.research_meta.update(metadata) 

46 else: 

47 research.research_meta = metadata 

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 

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 delete_report( 

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

108 ) -> bool: 

109 """Delete report from database.""" 

110 try: 

111 research = ( 

112 self.session.query(ResearchHistory) 

113 .filter_by(id=research_id) 

114 .first() 

115 ) 

116 

117 if not research: 

118 return False 

119 

120 research.report_content = None 

121 self.session.commit() 

122 

123 return True 

124 

125 except Exception: 

126 logger.exception("Error deleting report") 

127 self.session.rollback() 

128 return False 

129 

130 def report_exists( 

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

132 ) -> bool: 

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

134 try: 

135 research = ( 

136 self.session.query(ResearchHistory) 

137 .filter_by(id=research_id) 

138 .first() 

139 ) 

140 

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

142 

143 except Exception: 

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

145 return False