Coverage for mindsdb / integrations / handlers / zotero_handler / zotero_tables.py: 0%

1import pandas as pd 

2from mindsdb.utilities import log 

3from mindsdb_sql_parser import ast 

4from mindsdb.integrations.libs.api_handler import APITable 

5from mindsdb.integrations.utilities.sql_utils import extract_comparison_conditions 

6 

7logger = log.getLogger(__name__) 

8 

9 

10class AnnotationsTable(APITable): 

11 """Represents a table of annotations in Zotero.""" 

12 

13 def select(self, query: ast.Select): 

14 """Select annotations based on the provided query. 

15 

16 Parameters 

17 ---------- 

18 query : ast.Select 

19 AST (Abstract Syntax Tree) representation of the SQL query. 

20 

21 Returns 

22 ------- 

23 Response 

24 Response object containing the selected annotations as a DataFrame. 

25 """ 

26 if query.where is None: # Handle case for SELECT * FROM annotations 

27 df = self._get_items() 

28 return df[self.get_columns()] 

29 

30 conditions = extract_comparison_conditions(query.where) 

31 supported = False # Flag to check if the query is supported 

32 

33 for op, arg1, arg2 in conditions: 

34 if op in {'or', 'and'}: 

35 raise NotImplementedError('OR and AND are not supported') 

36 if arg1 == 'item_id' and op == '=': 

37 df = self._get_item(arg2) 

38 supported = True 

39 elif arg1 == 'parent_item_id' and op == '=': 

40 df = self._get_item_children(arg2) 

41 supported = True 

42 

43 if not supported: 

44 raise NotImplementedError('Only "item_id=" and "parent_item_id=" conditions are implemented') 

45 

46 return df[self.get_columns()] 

47 

48 def get_columns(self): 

49 """Get the columns of the annotations table. 

50 

51 Returns 

52 ------- 

53 list 

54 List of column names. 

55 """ 

56 return [ 

57 'annotationColor', 

58 'annotationComment', 

59 'annotationPageLabel', 

60 'annotationText', 

61 'annotationType', 

62 'dateAdded', 

63 'dateModified', 

64 'key', 

65 'parentItem', 

66 'relations', 

67 'tags', 

68 'version' 

69 ] 

70 

71 def _get_items(self) -> pd.DataFrame: 

72 """Get all annotations from the Zotero API. 

73 

74 Returns 

75 ------- 

76 pd.DataFrame 

77 DataFrame containing all annotations. 

78 """ 

79 if not self.handler.is_connected: 

80 self.handler.connect() 

81 

82 try: 

83 method = getattr(self.handler.api, 'items') 

84 result = method(itemType='annotation') 

85 

86 if isinstance(result, dict): 

87 return pd.DataFrame([result.get('data', {})]) 

88 if isinstance(result, list) and all(isinstance(item, dict) for item in result): 

89 data_list = [item.get('data', {}) for item in result] 

90 return pd.DataFrame(data_list) 

91 

92 except Exception as e: 

93 logger.error(f"Error fetching items: {e}") 

94 raise e 

95 

96 return pd.DataFrame() 

97 

98 def _get_item(self, item_id: str) -> pd.DataFrame: 

99 """Get a single annotation by item ID. 

100 

101 Parameters 

102 ---------- 

103 item_id : str 

104 The ID of the item to fetch. 

105 

106 Returns 

107 ------- 

108 pd.DataFrame 

109 DataFrame containing the annotation. 

110 """ 

111 if not self.handler.is_connected: 

112 self.handler.connect() 

113 

114 try: 

115 method = getattr(self.handler.api, 'item') 

116 result = method(item_id, itemType='annotation') 

117 

118 if isinstance(result, dict): 

119 return pd.DataFrame([result.get('data', {})]) 

120 

121 except Exception as e: 

122 logger.error(f"Error fetching item with ID {item_id}: {e}") 

123 raise e 

124 

125 return pd.DataFrame() 

126 

127 def _get_item_children(self, parent_item_id: str) -> pd.DataFrame: 

128 """Get annotations for a specific parent item ID. 

129 

130 Parameters 

131 ---------- 

132 parent_item_id : str 

133 The parent item ID to fetch annotations for. 

134 

135 Returns 

136 ------- 

137 pd.DataFrame 

138 DataFrame containing the annotations. 

139 """ 

140 if not self.handler.is_connected: 

141 self.handler.connect() 

142 

143 try: 

144 method = getattr(self.handler.api, 'children') 

145 result = method(parent_item_id, itemType='annotation') 

146 

147 if isinstance(result, dict): 

148 return pd.DataFrame([result.get('data', {})]) 

149 if isinstance(result, list) and all(isinstance(item, dict) for item in result): 

150 data_list = [item.get('data', {}) for item in result] 

151 return pd.DataFrame(data_list) 

152 

153 except Exception as e: 

154 logger.error(f"Error fetching children for parent item ID {parent_item_id}: {e}") 

155 raise e 

156 

157 return pd.DataFrame()