Coverage for mindsdb / interfaces / tabs / tabs_controller.py: 12%

1import json 

2from json import JSONDecodeError 

3from typing import Dict, List 

4from pathlib import Path 

5 

6from mindsdb.utilities import log 

7from mindsdb.utilities.context import context as ctx 

8from mindsdb.utilities.exception import EntityNotExistsError 

9from mindsdb.interfaces.storage.fs import FileStorageFactory, RESOURCE_GROUP, FileStorage 

10 

11 

12logger = log.getLogger(__name__) 

13 

14 

15TABS_FILENAME = "tabs" 

16 

17 

18def get_storage(): 

19 # deprecated 

20 

21 storageFactory = FileStorageFactory(resource_group=RESOURCE_GROUP.TAB, sync=True) 

22 

23 # resource_id is useless for 'tabs' 

24 # use constant 

25 return storageFactory(0) 

26 

27 

28class TabsController: 

29 """Tool for adding, editing, and deleting user's tabs 

30 

31 Attributes: 

32 storage_factory (FileStorageFactory): callable object which returns tabs file storage 

33 """ 

34 

35 def __init__(self) -> None: 

36 self.storage_factory = FileStorageFactory(resource_group=RESOURCE_GROUP.TAB, sync=True) 

37 

38 def _get_file_storage(self) -> FileStorage: 

39 """Get user's tabs file storage 

40 NOTE: file storage depend is company_id sensitive, so need to recreate it each time 

41 

42 Returns: 

43 FileStorage 

44 """ 

45 return self.storage_factory(0) 

46 

47 def _get_next_tab_id(self) -> int: 

48 """Get next free tab id 

49 

50 Returns: 

51 int: id for next tab 

52 """ 

53 tabs_files = self._get_tabs_files() 

54 tabs_ids = list(tabs_files.keys()) 

55 if len(tabs_ids) == 0: 

56 return 1 

57 return max(tabs_ids) + 1 

58 

59 def _get_tabs_files(self) -> Dict[int, Path]: 

60 """Get list of paths to each tab file 

61 

62 Returns: 

63 Dict[int, Path] 

64 """ 

65 tabs = {} 

66 for child in self._get_file_storage().folder_path.iterdir(): 

67 if (child.is_file() and child.name.startswith("tab_")) is False: 

68 continue 

69 tab_id = child.name.replace("tab_", "") 

70 if tab_id.isnumeric() is False: 

71 continue 

72 tabs[int(tab_id)] = child 

73 return tabs 

74 

75 def _get_tabs_meta(self) -> List[Dict]: 

76 """Get tabs info without content 

77 

78 Returns: 

79 List[Dict] 

80 """ 

81 all_tabs = self.get_all() 

82 for tab in all_tabs: 

83 del tab["content"] 

84 return all_tabs 

85 

86 def _load_tab_data(self, tab_id: int, raw_data) -> Dict: 

87 """Load tab JSON while handling trailing garbage.""" 

88 if isinstance(raw_data, bytes): 

89 raw_data_str = raw_data.decode("utf-8") 

90 else: 

91 raw_data_str = raw_data 

92 

93 try: 

94 return json.loads(raw_data_str) 

95 except JSONDecodeError: 

96 decoder = json.JSONDecoder() 

97 stripped = raw_data_str.lstrip() 

98 data, idx = decoder.raw_decode(stripped) 

99 

100 trailing = stripped[idx:].strip() 

101 if trailing: 

102 logger.warning("Detected trailing data in tab %s/%s, attempting to sanitize", ctx.company_id, tab_id) 

103 try: 

104 sanitized_bytes = json.dumps(data).encode("utf-8") 

105 self._get_file_storage().file_set(f"tab_{tab_id}", sanitized_bytes) 

106 except Exception as rewrite_error: 

107 logger.warning("Failed to rewrite sanitized tab %s/%s: %s", ctx.company_id, tab_id, rewrite_error) 

108 return data 

109 

110 def _migrate_legacy(self) -> None: 

111 """Convert old single-file tabs storage to multiple files""" 

112 file_storage = self._get_file_storage() 

113 try: 

114 file_data = file_storage.file_get(TABS_FILENAME) 

115 except FileNotFoundError: 

116 return 

117 except Exception: 

118 file_storage.delete() 

119 return 

120 

121 try: 

122 data = json.loads(file_data) 

123 except Exception: 

124 file_storage.delete() 

125 return 

126 

127 if isinstance(data, dict) is False or isinstance(data.get("tabs"), str) is False: 

128 file_storage.delete() 

129 return 

130 

131 try: 

132 tabs_list = json.loads(data["tabs"]) 

133 except Exception: 

134 file_storage.delete() 

135 return 

136 

137 if isinstance(tabs_list, list) is False: 

138 file_storage.delete() 

139 return 

140 

141 for tab in tabs_list: 

142 tab_id = self._get_next_tab_id() 

143 

144 b_types = json.dumps( 

145 {"index": tab.get("index", 0), "name": tab.get("name", "undefined"), "content": tab.get("value", "")} 

146 ).encode("utf-8") 

147 file_storage.file_set(f"tab_{tab_id}", b_types) 

148 

149 file_storage.delete(TABS_FILENAME) 

150 

151 def get_all(self) -> List[Dict]: 

152 """Get list of all tabs 

153 

154 Returns: 

155 List[Dict]: all tabs data 

156 """ 

157 self._get_file_storage().pull() 

158 self._migrate_legacy() 

159 

160 tabs_files = self._get_tabs_files() 

161 tabs_list = [] 

162 for tab_id, tab_path in tabs_files.items(): 

163 try: 

164 data = self._load_tab_data(tab_id, tab_path.read_text()) 

165 except Exception as e: 

166 logger.error(f"Can't read data of tab {ctx.company_id}/{tab_id}: {e}") 

167 continue 

168 tabs_list.append({"id": tab_id, **data}) 

169 

170 tabs_list.sort(key=lambda x: x["index"]) 

171 return tabs_list 

172 

173 def get(self, tab_id: int) -> Dict: 

174 """Get data of single tab 

175 

176 Args: 

177 tab_id (int): id of the tab 

178 

179 Returns: 

180 dict: tabs data 

181 """ 

182 if isinstance(tab_id, int) is False: 

183 raise ValueError("Tab id must be integer") 

184 

185 try: 

186 raw_tab_data = self._get_file_storage().file_get(f"tab_{tab_id}") 

187 except FileNotFoundError as e: 

188 raise EntityNotExistsError(f"tab {tab_id}") from e 

189 

190 try: 

191 data = self._load_tab_data(tab_id, raw_tab_data) 

192 except Exception as e: 

193 logger.error(f"Can't read data of tab {ctx.company_id}/{tab_id}: {e}") 

194 raise Exception(f"Can't read data of tab: {e}") from e 

195 

196 return {"id": tab_id, **data} 

197 

198 def add(self, index: int = None, name: str = "undefined", content: str = "") -> Dict: 

199 """Add new tab 

200 

201 Args: 

202 index (int, optional): index of new tab 

203 name (str, optional): name of new tab 

204 content (str, optional): content of new tab 

205 

206 Returns: 

207 dict: new tab meta info: id, name and index 

208 """ 

209 file_storage = self._get_file_storage() 

210 tab_id = self._get_next_tab_id() 

211 

212 reorder_required = index is not None 

213 if index is None: 

214 all_tabs = self.get_all() 

215 if len(all_tabs) == 0: 

216 index = 0 

217 else: 

218 index = max([x.get("index", 0) for x in all_tabs]) + 1 

219 

220 data_bytes = json.dumps({"index": index, "name": name, "content": content}).encode("utf-8") 

221 file_storage.file_set(f"tab_{tab_id}", data_bytes) 

222 

223 if reorder_required: 

224 all_tabs = self.get_all() 

225 all_tabs.sort(key=lambda x: (x["index"], 0 if x["id"] == tab_id else 1)) 

226 file_storage.sync = False 

227 for tab_index, tab in enumerate(all_tabs): 

228 tab["index"] = tab_index 

229 data_bytes = json.dumps(tab).encode("utf-8") 

230 file_storage.file_set(f"tab_{tab['id']}", data_bytes) 

231 file_storage.sync = True 

232 file_storage.push() 

233 

234 return {"id": tab_id, "index": index, "name": name} 

235 

236 def modify(self, tab_id: int, index: int = None, name: str = None, content: str = None) -> Dict: 

237 """Modify the tab 

238 

239 Args: 

240 tab_id (int): if of the tab to modify 

241 index (int, optional): tab's new index 

242 name (str, optional): tab's new name 

243 content (str, optional): tab's new content 

244 

245 Returns: 

246 dict: new tab meta info: id, name and index 

247 """ 

248 file_storage = self._get_file_storage() 

249 current_data = self.get(tab_id) 

250 

251 # region modify index 

252 if index is not None and current_data["index"] != index: 

253 current_data["index"] = index 

254 all_tabs = [x for x in self.get_all() if x["id"] != tab_id] 

255 all_tabs.sort(key=lambda x: x["index"]) 

256 file_storage.sync = False 

257 for tab_index, tab in enumerate(all_tabs): 

258 if tab_index < index: 

259 tab["index"] = tab_index 

260 else: 

261 tab["index"] = tab_index + 1 

262 data_bytes = json.dumps(tab).encode("utf-8") 

263 file_storage.file_set(f"tab_{tab['id']}", data_bytes) 

264 file_storage.sync = True 

265 file_storage.push() 

266 # endregion 

267 

268 # region modify name 

269 if name is not None and current_data["name"] != name: 

270 current_data["name"] = name 

271 # endregion 

272 

273 # region modify content 

274 if content is not None and current_data["content"] != content: 

275 current_data["content"] = content 

276 # endregion 

277 

278 data_bytes = json.dumps(current_data).encode("utf-8") 

279 file_storage.file_set(f"tab_{tab_id}", data_bytes) 

280 

281 return {"id": current_data["id"], "index": current_data["index"], "name": current_data["name"]} 

282 

283 def delete(self, tab_id: int): 

284 file_storage = self._get_file_storage() 

285 try: 

286 file_storage.file_get(f"tab_{tab_id}") 

287 except FileNotFoundError as e: 

288 raise EntityNotExistsError(f"tab {tab_id}") from e 

289 

290 file_storage.delete(f"tab_{tab_id}") 

291 

292 

293tabs_controller = TabsController()