fix: stream direct site videos
This commit is contained in:
@@ -286,6 +286,30 @@ class TelegramStorageUploader:
|
||||
except Exception:
|
||||
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:
|
||||
await self.pool.execute(
|
||||
"""
|
||||
@@ -435,7 +459,11 @@ class TelegramStorageUploader:
|
||||
await self.mark_media_link_only(media_id, "video too long")
|
||||
continue
|
||||
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:
|
||||
await self.mark_media_failed_attempt(media_id, "video download failed")
|
||||
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:
|
||||
@@ -9,5 +33,14 @@ def test_video_provider() -> 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__":
|
||||
test_video_provider()
|
||||
asyncio.run(test_http_video_download())
|
||||
|
||||
Reference in New Issue
Block a user