fix: stream direct site videos
This commit is contained in:
@@ -286,6 +286,30 @@ class TelegramStorageUploader:
|
|||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
async def download_http_video(
|
||||||
|
self,
|
||||||
|
session: aiohttp.ClientSession,
|
||||||
|
url: str,
|
||||||
|
output_path: str,
|
||||||
|
) -> dict:
|
||||||
|
max_size = self.video_max_size_mb * 1024 * 1024
|
||||||
|
try:
|
||||||
|
async with session.get(url, timeout=aiohttp.ClientTimeout(total=self.download_timeout_sec)) as response:
|
||||||
|
if response.status != 200:
|
||||||
|
return {"error": f"video returned HTTP {response.status}", "permanent": False}
|
||||||
|
size = 0
|
||||||
|
with open(output_path, "wb") as fh:
|
||||||
|
async for chunk in response.content.iter_chunked(1024 * 1024):
|
||||||
|
size += len(chunk)
|
||||||
|
if size > max_size:
|
||||||
|
remove_download_files(output_path)
|
||||||
|
return {"error": "video too large", "permanent": True}
|
||||||
|
fh.write(chunk)
|
||||||
|
except Exception:
|
||||||
|
remove_download_files(output_path)
|
||||||
|
return {"error": "video download failed", "permanent": False}
|
||||||
|
return {"path": output_path, "size_bytes": size}
|
||||||
|
|
||||||
async def mark_media_uploaded(self, media_id: int, file_id: str, unique_id: str | None) -> None:
|
async def mark_media_uploaded(self, media_id: int, file_id: str, unique_id: str | None) -> None:
|
||||||
await self.pool.execute(
|
await self.pool.execute(
|
||||||
"""
|
"""
|
||||||
@@ -435,7 +459,11 @@ class TelegramStorageUploader:
|
|||||||
await self.mark_media_link_only(media_id, "video too long")
|
await self.mark_media_link_only(media_id, "video too long")
|
||||||
continue
|
continue
|
||||||
temp_path = str(TMP_DIR / f"{TMP_PREFIX}{raw_post_id}_{media_id}.mp4")
|
temp_path = str(TMP_DIR / f"{TMP_PREFIX}{raw_post_id}_{media_id}.mp4")
|
||||||
info = await self.download_video(url, temp_path)
|
info = (
|
||||||
|
await self.download_video(url, temp_path)
|
||||||
|
if video_provider(url)
|
||||||
|
else await self.download_http_video(session, url, temp_path)
|
||||||
|
)
|
||||||
if not info:
|
if not info:
|
||||||
await self.mark_media_failed_attempt(media_id, "video download failed")
|
await self.mark_media_failed_attempt(media_id, "video download failed")
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -1,4 +1,28 @@
|
|||||||
from vk_parser_app.workers.vk_storage_uploader import video_provider
|
import asyncio
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from vk_parser_app.workers.vk_storage_uploader import TelegramStorageUploader, video_provider
|
||||||
|
|
||||||
|
|
||||||
|
class FakeContent:
|
||||||
|
async def iter_chunked(self, _size):
|
||||||
|
yield b"video"
|
||||||
|
|
||||||
|
|
||||||
|
class FakeResponse:
|
||||||
|
status = 200
|
||||||
|
content = FakeContent()
|
||||||
|
|
||||||
|
async def __aenter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, *_args):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class FakeSession:
|
||||||
|
def get(self, *_args, **_kwargs):
|
||||||
|
return FakeResponse()
|
||||||
|
|
||||||
|
|
||||||
def test_video_provider() -> None:
|
def test_video_provider() -> None:
|
||||||
@@ -9,5 +33,14 @@ def test_video_provider() -> None:
|
|||||||
assert video_provider("http://127.0.0.1/video.mp4") is None
|
assert video_provider("http://127.0.0.1/video.mp4") is None
|
||||||
|
|
||||||
|
|
||||||
|
async def test_http_video_download() -> None:
|
||||||
|
path = "/tmp/vkparser_tg_media_http_test.mp4"
|
||||||
|
result = await TelegramStorageUploader().download_http_video(FakeSession(), "https://example.test/a.mp4", path)
|
||||||
|
assert result["size_bytes"] == 5
|
||||||
|
assert Path(path).read_bytes() == b"video"
|
||||||
|
Path(path).unlink()
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
test_video_provider()
|
test_video_provider()
|
||||||
|
asyncio.run(test_http_video_download())
|
||||||
|
|||||||
Reference in New Issue
Block a user