diff --git a/yt_dlp_stream_to_s3/s3_client.py b/yt_dlp_stream_to_s3/client.py similarity index 65% rename from yt_dlp_stream_to_s3/s3_client.py rename to yt_dlp_stream_to_s3/client.py index 42fc0b9..ce9e405 100644 --- a/yt_dlp_stream_to_s3/s3_client.py +++ b/yt_dlp_stream_to_s3/client.py @@ -1,4 +1,5 @@ import boto3 +from botocore.exceptions import ClientError from yt_dlp_stream_to_s3.config import ( AWS_ENDPOINT_URL, AWS_SECRET_ACCESS_KEY, @@ -6,6 +7,7 @@ from yt_dlp_stream_to_s3.config import ( ) from yt_dlp_stream_to_s3.errors import YtDlpStreamToS3ConnectionError + try: s3 = boto3.client( "s3", @@ -13,7 +15,16 @@ try: aws_secret_access_key=AWS_SECRET_ACCESS_KEY, aws_access_key_id=AWS_ACCESS_KEY_ID, ) + except Exception: raise YtDlpStreamToS3ConnectionError( f"Could not connect to S3 service {AWS_ENDPOINT_URL}" ) + + +def create_presigned_url(bucket_name, object_name, expiration=3600): + return s3.generate_presigned_url( + "get_object", + Params={"Bucket": bucket_name, "Key": object_name}, + ExpiresIn=expiration, + ) diff --git a/yt_dlp_stream_to_s3/main.py b/yt_dlp_stream_to_s3/main.py index 74f4799..b26e7b7 100644 --- a/yt_dlp_stream_to_s3/main.py +++ b/yt_dlp_stream_to_s3/main.py @@ -1,14 +1,19 @@ import subprocess import yt_dlp +import boto3 from urllib.error import HTTPError -from pathlib import Path -from yt_dlp_stream_to_s3.s3_client import s3 +from botocore.exceptions import ClientError +from yt_dlp_stream_to_s3.client import s3, create_presigned_url from yt_dlp_stream_to_s3.utils import remove_queries_from_url from yt_dlp_stream_to_s3.config import AWS_S3_BUCKET_NAME from yt_dlp_stream_to_s3.errors import YtDlpStreamToS3Error -def yt_dlp_stream_to_s3(url: str, s3_extra_args: dict = ()) -> None: +def yt_dlp_stream_to_s3( + url: str, + expiration=86400 * 7, # 7 days + s3_extra_args: dict = (), +) -> None: media_url = remove_queries_from_url(url) ydl_opts = { @@ -47,6 +52,8 @@ def yt_dlp_stream_to_s3(url: str, s3_extra_args: dict = ()) -> None: ) yt_dlp_process.wait() + + return create_presigned_url(AWS_S3_BUCKET_NAME, filename, expiration) except HTTPError: raise YtDlpStreamToS3Error("Could not connect to the media resource.") @@ -55,3 +62,6 @@ def yt_dlp_stream_to_s3(url: str, s3_extra_args: dict = ()) -> None: except yt_dlp.utils.DownloadError: raise YtDlpStreamToS3Error("Could not download media. Check logs.") + + except ClientError: + raise YtDlpStreamToS3Error("Error communicating with the S3 service.")