Skip to content

Commit

Permalink
Merge pull request #93 from NabilMostafa/stream-uploaded-files
Browse files Browse the repository at this point in the history
Stream uploaded files
  • Loading branch information
omarryhan authored Apr 19, 2022
2 parents 3dc48d7 + 9d31f97 commit 66b65c3
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 4 deletions.
9 changes: 9 additions & 0 deletions aiogoogle/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ class MediaUpload:
validate (bool): Whether or not a session should validate the upload size before sending
pipe_from (file object): class object to stream file content from
"""

def __init__(
Expand All @@ -64,6 +65,7 @@ def __init__(
chunk_size=None,
resumable=None,
validate=True,
pipe_from=None
):
if isinstance(file_path_or_bytes, bytes):
self.file_body = file_path_or_bytes
Expand All @@ -78,6 +80,7 @@ def __init__(
self.chunk_size = chunk_size or DEFAULT_UPLOAD_CHUNK_SIZE
self.resumable = resumable
self.validate = validate
self.pipe_from = pipe_from

async def run_validation(self, size_func):
if self.validate and self.max_size:
Expand All @@ -93,6 +96,8 @@ async def aiter_file(self, aiter_func):
if self.file_path:
async for chunk in aiter_func(self.file_path, self.chunk_size):
yield chunk
elif self.pipe_from:
yield self.pipe_from.read()
else:
async for chunk in self._aiter_body():
yield chunk
Expand Down Expand Up @@ -265,6 +270,8 @@ class Response:
upload_file (str): path of the upload file specified in the request
pipe_from (file object): class object to stream file content from
session_factory (aiogoogle.sessions.abc.AbstractSession): A callable implementation of aiogoogle's session interface
"""

Expand All @@ -280,6 +287,7 @@ def __init__(
download_file=None,
pipe_to=None,
upload_file=None,
pipe_from=None,
session_factory=None,
):
if json and data:
Expand All @@ -295,6 +303,7 @@ def __init__(
self.download_file = download_file
self.pipe_to = pipe_to
self.upload_file = upload_file
self.pipe_from = pipe_from
self.session_factory = session_factory

@staticmethod
Expand Down
20 changes: 16 additions & 4 deletions aiogoogle/resource.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
"data",
"json",
"upload_file",
"pipe_from",
"download_file",
"pipe_to",
"timeout",
Expand Down Expand Up @@ -60,6 +61,7 @@ def wrapper(
data=None,
json=None,
upload_file=None,
pipe_from=None,
download_file=None,
pipe_to=None,
timeout=None,
Expand All @@ -85,6 +87,7 @@ def wrapper(
data,
json,
upload_file,
pipe_from,
download_file,
pipe_to,
timeout,
Expand Down Expand Up @@ -389,6 +392,7 @@ def __call__(
data=None,
json=None,
upload_file=None,
pipe_from=None,
download_file=None,
pipe_to=None,
timeout=None,
Expand Down Expand Up @@ -579,13 +583,20 @@ def __call__(
)
media_download = MediaDownload(pipe_to=pipe_to)

if not upload_file and not pipe_from:
media_upload = None

# Process upload_file
if upload_file:
media_upload = self._build_upload_media(
upload_file, uri, validate, fallback_url=url
upload_file, uri, validate, fallback_url=url, pipe_from=None
)

# Process pipe_from
if pipe_from:
media_upload = self._build_upload_media(
upload_file, uri, validate, fallback_url=url, pipe_from=pipe_from
)
else:
media_upload = None

return Request(
method=self["httpMethod"],
Expand Down Expand Up @@ -631,7 +642,7 @@ def _build_url(self, base_url, uri_params, validate):
else:
return base_url + self["path"]

def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url):
def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url, pipe_from):
if not self["supportsMediaUpload"]:
if validate:
raise ValidationError(
Expand Down Expand Up @@ -660,6 +671,7 @@ def _build_upload_media(self, upload_file, qualified_url, validate, fallback_url

return MediaUpload(
upload_file,
pipe_from=pipe_from,
upload_path=media_upload_url,
max_size=max_size,
mime_range=mime_range,
Expand Down
2 changes: 2 additions & 0 deletions aiogoogle/sessions/aiohttp_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ async def resolve_response(request, response):
json = None
download_file = None
upload_file = None
pipe_from = None
pipe_to = None

# If downloading file:
Expand Down Expand Up @@ -89,6 +90,7 @@ async def resolve_response(request, response):
download_file=download_file,
pipe_to=pipe_to,
upload_file=upload_file,
pipe_from=pipe_from,
session_factory=session_factory,
)

Expand Down
2 changes: 2 additions & 0 deletions aiogoogle/sessions/curio_asks_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ async def resolve_response(request, response):
json = None
download_file = None
upload_file = None
pipe_from = None

# If downloading file:
if request.media_download:
Expand Down Expand Up @@ -67,6 +68,7 @@ async def resolve_response(request, response):
req=request,
download_file=download_file,
upload_file=upload_file,
pipe_from=pipe_from,
session_factory=session_factory,
)

Expand Down
40 changes: 40 additions & 0 deletions examples/stream_upload_file.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
#!/usr/bin/python3.7

'''
Scopes Required:
* https://www.googleapis.com/auth/drive
* https://www.googleapis.com/auth/drive.file
API explorer link:
* https://developers.google.com/apis-explorer/#p/drive/v3/drive.files.create
'''

import asyncio

from helpers import Aiogoogle, user_creds, client_creds


class MyFile:
@staticmethod
async def read():
return b'Hello World'


async def stream_upload_file():
async with Aiogoogle(user_creds=user_creds, client_creds=client_creds) as aiogoogle:
# Create API
drive_v3 = await aiogoogle.discover("drive", "v3")

req = drive_v3.files.create(
pipe_from=MyFile(),
fields="id"
)

# Upload file
upload_res = await aiogoogle.as_user(req)
print(f"Uploaded stream file successfully.\nFile ID: {upload_res['id']}")

if __name__ == "__main__":
asyncio.run(stream_upload_file())

0 comments on commit 66b65c3

Please sign in to comment.