From 095d2d8b0b150dbc491cd0854dd7da510194b97f Mon Sep 17 00:00:00 2001
From: Nabil Khalil <n.khalil@greyrook.com>
Date: Tue, 1 Mar 2022 12:56:19 +0100
Subject: [PATCH] Add streaming files option when downloading files

---
 aiogoogle/models.py                   |  9 ++++++++-
 aiogoogle/resource.py                 | 26 +++++++++++++++++++++++---
 aiogoogle/sessions/aiohttp_session.py | 13 +++++++++++--
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/aiogoogle/models.py b/aiogoogle/models.py
index f620a35..cdb3a83 100644
--- a/aiogoogle/models.py
+++ b/aiogoogle/models.py
@@ -118,10 +118,13 @@ class MediaDownload:
 
         chunksize (int): Size of a chunk of bytes that a session should write at a time when downloading.
 
+        pipe_to (object): class object to stream file content to
+
     """
 
-    def __init__(self, file_path, chunk_size=None):
+    def __init__(self, file_path=None, chunk_size=None, pipe_to=None):
         self.file_path = file_path
+        self.pipe_to = pipe_to
         self.chunk_size = chunk_size or DEFAULT_DOWNLOAD_CHUNK_SIZE
 
 
@@ -258,6 +261,8 @@ class Response:
 
         download_file (str): path of the download file specified in the request
 
+        pipe_to (object): class object to stream file content to specified in the request.
+
         upload_file (str): path of the upload file specified in the request
 
         session_factory (aiogoogle.sessions.abc.AbstractSession): A callable implementation of aiogoogle's session interface
@@ -273,6 +278,7 @@ def __init__(
         reason=None,
         req=None,
         download_file=None,
+        pipe_to=None,
         upload_file=None,
         session_factory=None,
     ):
@@ -287,6 +293,7 @@ def __init__(
         self.reason = reason
         self.req = req
         self.download_file = download_file
+        self.pipe_to = pipe_to
         self.upload_file = upload_file
         self.session_factory = session_factory
 
diff --git a/aiogoogle/resource.py b/aiogoogle/resource.py
index 92efb6d..480f195 100644
--- a/aiogoogle/resource.py
+++ b/aiogoogle/resource.py
@@ -23,6 +23,7 @@
     "json",
     "upload_file",
     "download_file",
+    "pipe_to",
     "timeout",
 ]
 
@@ -60,6 +61,7 @@ def wrapper(
         json=None,
         upload_file=None,
         download_file=None,
+        pipe_to=None,
         timeout=None,
         **uri_params,
     ):
@@ -84,6 +86,7 @@ def wrapper(
             json,
             upload_file,
             download_file,
+            pipe_to,
             timeout,
             **uri_params,
         )
@@ -387,6 +390,7 @@ def __call__(
         json=None,
         upload_file=None,
         download_file=None,
+        pipe_to=None,
         timeout=None,
         **uri_params,
     ) -> Request:
@@ -548,6 +552,15 @@ def __call__(
             if body is not None:
                 self._validate_body(body)
 
+        # Validate selected options
+        if download_file and pipe_to:
+            raise ValueError(
+                "can't have both (download_file) and (pipe_to) options"
+            )
+
+        if not download_file and not pipe_to:
+            media_download = None
+
         # Process download_file
         if download_file:
             if validate is True:
@@ -555,9 +568,16 @@ def __call__(
                     raise ValidationError(
                         "download_file was provided while method doesn't support media download"
                     )
-            media_download = MediaDownload(download_file)
-        else:
-            media_download = None
+            media_download = MediaDownload(file_path=download_file)
+
+        # Process pipe_to
+        if pipe_to:
+            if validate is True:
+                if self.__getitem__("supportsMediaDownload") is not True:
+                    raise ValidationError(
+                        "pipe_to was provided while method doesn't support media download"
+                    )
+            media_download = MediaDownload(pipe_to=pipe_to)
 
         # Process upload_file
         if upload_file:
diff --git a/aiogoogle/sessions/aiohttp_session.py b/aiogoogle/sessions/aiohttp_session.py
index 4706658..a36624c 100644
--- a/aiogoogle/sessions/aiohttp_session.py
+++ b/aiogoogle/sessions/aiohttp_session.py
@@ -46,14 +46,22 @@ async def resolve_response(request, response):
             json = None
             download_file = None
             upload_file = None
+            pipe_to = None
 
             # If downloading file:
             if request.media_download:
                 chunk_size = request.media_download.chunk_size
                 download_file = request.media_download.file_path
-                async with aiofiles.open(download_file, "wb+") as f:
+                pipe_to = request.media_download.pipe_to
+
+                if pipe_to:
                     async for line in response.content.iter_chunked(chunk_size):
-                        await f.write(line)
+                            pipe_to.write(line)
+
+                if download_file:
+                    async with aiofiles.open(download_file, "wb+") as f:
+                        async for line in response.content.iter_chunked(chunk_size):
+                            await f.write(line)
             else:
                 if response.status != 204:  # If no (no content)
                     try:
@@ -79,6 +87,7 @@ async def resolve_response(request, response):
                 reason=response.reason if getattr(response, "reason") else None,
                 req=request,
                 download_file=download_file,
+                pipe_to=pipe_to,
                 upload_file=upload_file,
                 session_factory=session_factory,
             )