Fix upload and cencellation

Author: shemogumbe

src/msgraph_core/models/large_file_upload_session.py

@@ -18,12 +18,52 @@ def __init__(
         is_cancelled: Optional[bool] = False,
         next_expected_ranges: Optional[List[str]] = None
     ):
-        self.upload_url = upload_url
-        self.expiration_date_time = expiration_date_time
+        self._upload_url = upload_url
+        self._expiration_date_time = expiration_date_time
         self.additional_data = additional_data if additional_data is not None else []
         self.is_cancelled = is_cancelled
         self.next_expected_ranges = next_expected_ranges if next_expected_ranges is not None else []
 
+    @property
+    def upload_url(self):
+        return self._upload_url
+
+    @upload_url.setter
+    def upload_url(self, value):
+        self._upload_url = value
+
+    @property
+    def expiration_date_time(self):
+        return self._expiration_date_time
+
+    @expiration_date_time.setter
+    def expiration_date_time(self, value):
+        self._expiration_date_time = value
+
+    @property
+    def additional_data(self):
+        return self._additional_data
+
+    @additional_data.setter
+    def additional_data(self, value):
+        self._additional_data = value if value is not None else []
+
+    @property
+    def is_cancelled(self):
+        return self._is_cancelled
+
+    @is_cancelled.setter
+    def is_cancelled(self, value):
+        self._is_cancelled = value
+
+    @property
+    def next_expected_ranges(self):
+        return self._next_expected_ranges
+
+    @next_expected_ranges.setter
+    def next_expected_ranges(self, value):
+        self._next_expected_ranges = value if value is not None else []
+
     @staticmethod
     def create_from_discriminator_value(
         parse_node: ParseNode

src/msgraph_core/tasks/large_file_upload.py

@@ -2,6 +2,7 @@
 from io import BytesIO
 from asyncio import Future
 from datetime import datetime, timedelta
+import logging
 
 from kiota_abstractions.serialization.parsable import Parsable
 from kiota_abstractions.method import Method
@@ -22,8 +23,10 @@ def __init__(
         stream: BytesIO,  # counter check this
         max_chunk_size: int = 1024  # 4 * 1024 * 1024 - use smaller chnuks for testing
     ):
-        self._upload_session = upload_session
-        self._request_adapter = request_adapter
+        if not isinstance(upload_session, LargeFileUploadSession):
+            raise TypeError("upload_session must be an instance of LargeFileUploadSession")
+        self.upload_session = upload_session
+        self.request_adapter = request_adapter
         self.stream = stream
         self.file_size = stream.getbuffer().nbytes
         self.max_chunk_size = max_chunk_size
@@ -36,15 +39,14 @@ def __init__(
 
     @property
     def upload_session(self) -> Parsable:
-        return self._upload_session
+        return self.upload_session
 
     @staticmethod
     async def create_upload_session(request_adapter: RequestAdapter, request_body, url: str):
         request_information = RequestInformation()
         base_url = request_adapter.base_url.rstrip('/')
         path = url.lstrip('/')
         new_url = f"{base_url}/{path}"
-        print(f"New url {new_url}")
         request_information.url = new_url
         request_information.http_method = Method.POST
         request_information.set_content_from_parsable(
@@ -58,19 +60,18 @@ async def create_upload_session(request_adapter: RequestAdapter, request_body, u
 
     @property
     def request_adapter(self) -> RequestAdapter:
-        return self._request_adapter
+        return self.request_adapter
 
     @property
     def chunks(self) -> int:
         return self._chunks
 
     def upload_session_expired(self, upload_session: LargeFileUploadSession = None) -> bool:
         now = datetime.now()
-        upload_session = upload_session or self._upload_session
+        upload_session = upload_session or self.upload_session
         if not hasattr(upload_session, "expiration_date_time"):
             raise ValueError("Upload session does not have an expiration date time")
         expiry = upload_session.expiration_date_time
-        print(expiry)
         if expiry is None:
             raise ValueError("Expiry is None")
         if isinstance(expiry, str):
@@ -86,66 +87,53 @@ def upload_session_expired(self, upload_session: LargeFileUploadSession = None)
             return True
         return False
 
-    async def upload(self, after_chunk_upload=None):
+    async def upload(self, after_chunk_upload: Optional[Callable] = None):
         # Rewind at this point to take care of failures.
         self.stream.seek(0)
         if self.upload_session_expired(self.upload_session):
             raise RuntimeError('The upload session is expired.')
 
-        self.on_chunk_upload_complete = after_chunk_upload if self.on_chunk_upload_complete is None else self.on_chunk_upload_complete
+        self.on_chunk_upload_complete = after_chunk_upload or self.on_chunk_upload_complete
         session = await self.next_chunk(
             self.stream, 0, max(0, min(self.max_chunk_size - 1, self.file_size - 1))
         )
         process_next = session
         # determine the range  to be uploaded
         # even when resuming existing upload sessions.
-        range_parts = self.next_range[0].split("-") if self.next_range else ['0', '0']
+        range_parts = self.next_range[0].split("-") if self.next_range else ['0']
         end = min(int(range_parts[0]) + self.max_chunk_size - 1, self.file_size)
         uploaded_range = [range_parts[0], end]
-        print(f"File size {self.file_size}")
         while self.chunks > 0:
             session = process_next
-            future = Future()
-
-            def success_callback(large_file_uploasd_session):
-                nonlocal process_next, uploaded_range
-
-                if large_file_uploasd_session is None:
-                    return large_file_uploasd_session
-
-                next_range = large_file_uploasd_session.next_expected_ranges
+            try:
+                lfu_session: Optional[LargeFileUploadSession] = await session
+                if lfu_session is None:
+                    continue
+                next_range = lfu_session.next_expected_ranges
                 old_url = self.get_validated_upload_url(self.upload_session)
-                large_file_uploasd_session.upload_url = old_url
-
+                lfu_session.upload_url = old_url
                 if self.on_chunk_upload_complete is not None:
                     self.on_chunk_upload_complete(uploaded_range)
-
                 if not next_range:
-                    return large_file_uploasd_session
-
-                range_parts = next_range[0].split("-")
+                    continue
+                range_parts = str(next_range[0]).split("-")
                 end = min(int(range_parts[0]) + self.max_chunk_size, self.file_size)
                 uploaded_range = [range_parts[0], end]
-                self.set_next_range(next_range[0] + "-")
+                self.next_range = next_range[0] + "-"
                 process_next = self.next_chunk(self.stream)
-
-                return large_file_uploasd_session
-
-        def failure_callback(error):
-            raise error
-
-        future.add_done_callback(success_callback)
-        future.set_exception_callback(failure_callback)
-
-        if future is not None:
-            future.result()  # This will block until the future is resolved
-
-        self.chunks -= 1
-
+            except Exception as error:
+                logging.error(f"Error uploading chunk {error}")
+                raise  # remove after manual testing
+            self.chunks -= 1
         return session
 
-    def set_next_range(self, next_range: Optional[str]) -> None:
-        self.next_range = next_range
+    @property
+    def next_range(self):
+        return self._next_range
+
+    @next_range.setter
+    def next_range(self, value: Optional[str]) -> None:
+        self._next_range = value
 
     async def next_chunk(self, file: BytesIO, range_start: int = 0, range_end: int = 0) -> Future:
         upload_url = self.get_validated_upload_url(self.upload_session)
@@ -156,7 +144,7 @@ async def next_chunk(self, file: BytesIO, range_start: int = 0, range_end: int =
         info.url = upload_url
         info.http_method = Method.PUT
         if not self.next_range:
-            self.set_next_range(f'{range_start}-{range_end}')
+            self.next_range = f'{range_start}-{range_end}'
         range_parts = self.next_range.split('-') if self.next_range else ['-']
         start = int(range_parts[0])
         end = int(range_parts[1]) if len(range_parts) > 1 else 0
@@ -184,12 +172,6 @@ async def next_chunk(self, file: BytesIO, range_start: int = 0, range_end: int =
 
         parsable_factory: LargeFileUploadSession[Any] = self.upload_session
 
-        print(f"Body {LargeFileUploadSession.create_from_discriminator_value}")
-        print(f"Upload session url {self.upload_session.upload_url}")
-        print(f"Expiraton date {self.upload_session.expiration_date_time}")
-        print(f"Additional data {self.upload_session.additional_data}")
-        print(f"Session is canceled {self.upload_session.is_cancelled}")
-        print(f"Next expected ranges {self.upload_session.next_expected_ranges}")
         return await self.request_adapter.send_async(info, parsable_factory, error_map)
 
     def get_file(self) -> BytesIO: