Coverage for streamwise/file_manager.py: 60%

1""" 

2File manager functions for the Cluster Manager. 

3""" 

4 

5import sys 

6import logging 

7import aiofiles 

8import aiofiles.os 

9 

10from aiohttp import ClientTimeout 

11from aiohttp import ClientError 

12 

13from typing import Optional 

14from typing import Dict 

15from typing import Any 

16from typing import Union 

17 

18from io import BytesIO 

19 

20from quart import jsonify 

21from quart import send_file 

22from quart import send_from_directory 

23from quart import render_template 

24 

25from http import HTTPStatus 

26 

27from http_session_manager import get_global_session 

28import http_session_manager 

29 

30sys.path.append("..") 

31import quart_utils 

32from quart_utils import QuartReturn 

33from quart_utils import get_mime_type 

34 

35from tts_utils import generate_waveform_plt 

36 

37 

38# HTTP clients 

39CLIENT_TIMEOUT = ClientTimeout(total=5.0, connect=1.0) 

40 

41 

42async def get_audio_waveform(container_ip: str, container_port: int, file_name: str) -> QuartReturn: 

43 """Generate and return a waveform PNG image for a WAV audio file.""" 

44 if not container_ip or not container_port or not file_name: 

45 return jsonify({"error": "Container IP, port and file name are required"}), HTTPStatus.BAD_REQUEST 

46 if not file_name or not file_name.endswith((".wav")): 

47 return jsonify({"error": f"Invalid file name: {file_name}"}), HTTPStatus.BAD_REQUEST 

48 

49 # Download the WAV file into a local temp file 

50 temp_wav_path = f"/tmp/waveform_{file_name}" 

51 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file/{file_name}" 

52 try: 

53 session = await get_global_session() 

54 async with session.get(url, timeout=CLIENT_TIMEOUT) as response: 

55 if response.status == HTTPStatus.OK: 

56 content_type = "application/octet-stream" 

57 if "Content-Type" in response.headers: 

58 content_type = response.headers["Content-Type"] 

59 if content_type not in ("audio/wav", "audio/x-wav"): 

60 return jsonify({ 

61 "error": f"File {file_name} is not a WAV audio file: {content_type}" 

62 }), HTTPStatus.BAD_REQUEST 

63 audio_binary = await response.read() 

64 async with aiofiles.open(temp_wav_path, mode="wb") as wav_file: 

65 await wav_file.write(audio_binary) 

66 except Exception as ex: 

67 logging.error(f"Error downloading file from {url}: {ex}.") 

68 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR 

69 

70 waveform_png_path = generate_waveform_plt(temp_wav_path) 

71 

72 return await send_file( 

73 waveform_png_path, 

74 as_attachment=True, 

75 attachment_filename=f"waveform_{file_name}.png", 

76 mimetype="image/png") 

77 

78 

79async def get_video_info(container_ip: str, container_port: int, file_name: str) -> QuartReturn: 

80 """Get video file information from a container.""" 

81 if not container_ip or not container_port or not file_name: 

82 return jsonify({"error": "Container IP, port and file name are required"}), HTTPStatus.BAD_REQUEST 

83 if not file_name or not file_name.endswith((".mp4", ".avi", ".mkv", ".webm")): 

84 return jsonify({"error": f"Invalid file name: {file_name}"}), HTTPStatus.BAD_REQUEST 

85 

86 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file/{file_name}" 

87 try: 

88 session = await get_global_session() 

89 async with session.get(url, timeout=CLIENT_TIMEOUT) as response: 

90 if response.status != HTTPStatus.OK: 

91 return jsonify({"error": "Failed to get video info"}), HTTPStatus.INTERNAL_SERVER_ERROR 

92 

93 content_type = response.headers.get("Content-Type", "video/mp4") 

94 content = await response.read() 

95 content_length = len(content) 

96 return jsonify({ 

97 "file_name": file_name, 

98 "content_type": content_type, 

99 "content_length": content_length 

100 }), HTTPStatus.OK 

101 except Exception as ex: 

102 logging.error(f"Error getting video info from {url}: {ex}.") 

103 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR 

104 

105 

106async def list_files(path: str) -> QuartReturn: 

107 """List files in a directory.""" 

108 try: 

109 files = await quart_utils.list_files(path) 

110 return jsonify({"files": files}) 

111 except Exception as ex: 

112 logging.error(f"Error listing files in {path}: {ex}.") 

113 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR 

114 

115 

116async def download_local_file(path: str, file_name: str) -> QuartReturn: 

117 """Download a file.""" 

118 filepath = f"{path}/{file_name}" 

119 if not await aiofiles.os.path.exists(filepath): 

120 return jsonify({"error": "File not found"}), HTTPStatus.NOT_FOUND 

121 

122 if await aiofiles.os.path.isdir(filepath): 

123 files = await aiofiles.os.listdir(filepath) 

124 return jsonify({ 

125 "files": files 

126 }) 

127 

128 mimetype = get_mime_type(file_name) 

129 return await send_from_directory( 

130 path, 

131 file_name, 

132 mimetype=mimetype, 

133 as_attachment=True) 

134 

135 

136async def download_service_file( 

137 container_ip: str, 

138 container_port: int, 

139 file_name: str 

140) -> QuartReturn: 

141 """Download a file from a container.""" 

142 if not container_ip or not container_port or not file_name: 

143 return jsonify({"error": "Container IP, port and file name are required"}), HTTPStatus.BAD_REQUEST 

144 

145 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file/{file_name}" 

146 try: 

147 session = await get_global_session() 

148 timeout = ClientTimeout(total=1.0, connect=1.0) 

149 async with session.get(url, timeout=timeout) as response: 

150 if response.status == HTTPStatus.OK: 

151 content_type = "application/octet-stream" 

152 if "Content-Type" in response.headers: 

153 content_type = response.headers["Content-Type"] 

154 data = await response.read() 

155 return await send_file( 

156 BytesIO(data), 

157 as_attachment=True, 

158 attachment_filename=file_name, 

159 mimetype=content_type) 

160 return jsonify({"error": "Failed to download"}), HTTPStatus.INTERNAL_SERVER_ERROR 

161 except ClientError as client_err: 

162 logging.error(f"Client error downloading file from {url}: {client_err}.") 

163 return jsonify({"error": f"Client error: {client_err}"}), HTTPStatus.SERVICE_UNAVAILABLE 

164 except Exception as ex: 

165 logging.error(f"Error downloading file from {url} [{type(ex)}]: {ex}.") 

166 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR 

167 

168 

169async def get_file_info( 

170 container_ip: str, 

171 container_port: int, 

172 file_name: str 

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

174 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file_info/{file_name}" 

175 try: 

176 session = await get_global_session() 

177 async with session.get(url, timeout=CLIENT_TIMEOUT) as response: 

178 content_type = response.headers.get("Content-Type") 

179 if response.status == HTTPStatus.OK and content_type == "application/json": 

180 return await response.json() 

181 except ClientError as client_err: 

182 logging.error(f"Client error getting file info from {url}: {client_err}.") 

183 except Exception as ex: 

184 logging.error(f"Error getting file info from {url} [{type(ex)}]: {ex}.") 

185 return None 

186 

187 

188async def file_view( 

189 service_name: str, 

190 container_ip: str, 

191 container_port: int, 

192 file_name: str 

193) -> QuartReturn: 

194 """ 

195 View the contents of a file from a container. 

196 Returns HTML rendered template with file content. 

197 """ 

198 if not container_ip or not container_port or not file_name: 

199 return jsonify({"error": "Container IP, port and file name are required"}), HTTPStatus.BAD_REQUEST 

200 

201 content: Optional[Union[str, bytes]] = None 

202 content_type: Optional[str] = None 

203 content_length: Optional[int] = None 

204 file_info = None 

205 error: Optional[str] = None 

206 try: 

207 url_data = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file/{file_name}" 

208 session = await get_global_session() 

209 timeout = ClientTimeout(total=1.0, connect=1.0) 

210 async with session.get(url_data, timeout=timeout) as file_response: 

211 if file_response.status != HTTPStatus.OK: 

212 return jsonify({ 

213 "error": f"Failed to get file data: {file_response.status}" 

214 }), HTTPStatus.INTERNAL_SERVER_ERROR 

215 content_type = file_response.headers.get("Content-Type", "text/plain") 

216 raw_content = await file_response.read() 

217 content_length = len(raw_content) 

218 if content_type.startswith("text/") or content_type in ("application/json", "application/x-ndjson"): 

219 text_content: str = raw_content.decode("utf-8", errors="replace") 

220 content = text_content 

221 else: 

222 binary_content: bytes = raw_content 

223 content = binary_content 

224 

225 file_info = await get_file_info(container_ip, container_port, file_name) 

226 except ClientError as client_err: 

227 error = f"Client error getting file from {url_data}: {client_err}" 

228 except Exception as ex: 

229 error = str(ex) 

230 

231 return await render_template( 

232 "file_view.html", 

233 service_name=service_name, 

234 file_name=file_name, 

235 container_ip=container_ip, 

236 container_port=container_port, 

237 content=content, 

238 content_type=content_type, 

239 content_length=content_length, 

240 file_info=file_info, 

241 error=error) 

242 

243 

244async def file_stream( 

245 container_ip: str, 

246 container_port: int, 

247 file_name: str 

248) -> QuartReturn: 

249 """Stream a file from a container.""" 

250 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/file/{file_name}" 

251 try: 

252 session = await get_global_session() 

253 async with session.get(url, timeout=CLIENT_TIMEOUT) as response: 

254 if response.status == HTTPStatus.OK: 

255 content_type = response.headers.get("Content-Type", "application/octet-stream") 

256 data = await response.read() 

257 return await send_file( 

258 BytesIO(data), 

259 mimetype=content_type, 

260 attachment_filename=file_name) 

261 return jsonify({ 

262 "error": f"Failed to stream file {file_name} from {url}: {response.status}" 

263 }), response.status 

264 except ClientError as client_err: 

265 return jsonify({ 

266 "error": f"Client error while trying to reach {url}: {client_err}" 

267 }), HTTPStatus.SERVICE_UNAVAILABLE 

268 except TimeoutError as timeout_err: 

269 return jsonify({ 

270 "error": f"Timeout error while trying to reach {url}: {timeout_err}" 

271 }), HTTPStatus.GATEWAY_TIMEOUT 

272 except Exception as ex: 

273 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR