Merge pull request #683 from microsoftgraph/shem/batch_requests

Batch Requests

Author: shemogumbe

src/msgraph_core/requests/__init__.py

@@ -0,0 +1,144 @@
+from typing import TypeVar, Type, Dict, Optional, Union
+import logging
+
+from kiota_abstractions.request_adapter import RequestAdapter
+from kiota_abstractions.request_information import RequestInformation
+from kiota_abstractions.method import Method
+from kiota_abstractions.serialization import Parsable
+from kiota_abstractions.headers_collection import HeadersCollection
+from kiota_abstractions.api_error import APIError
+
+from .batch_request_content import BatchRequestContent
+from .batch_request_content_collection import BatchRequestContentCollection
+from .batch_response_content import BatchResponseContent
+from .batch_response_content_collection import BatchResponseContentCollection
+
+T = TypeVar('T', bound='Parsable')
+
+APPLICATION_JSON = "application/json"
+
+
+class BatchRequestBuilder:
+    """
+    Provides operations to call the batch method.
+    """
+
+    def __init__(
+        self,
+        request_adapter: RequestAdapter,
+        error_map: Optional[Dict[str, Type[Parsable]]] = None
+    ):
+        if request_adapter is None:
+            raise ValueError("request_adapter cannot be Null.")
+        self._request_adapter = request_adapter
+        self.url_template = f"{self._request_adapter.base_url}/$batch"
+        self.error_map = error_map or {}
+
+    async def post(
+        self,
+        batch_request_content: Union[BatchRequestContent, BatchRequestContentCollection],
+        error_map: Optional[Dict[str, Type[Parsable]]] = None,
+    ) -> Union[T, BatchResponseContentCollection]:
+        """
+        Sends a batch request and returns the batch response content.
+        
+        Args:
+            batch_request_content (Union[BatchRequestContent, 
+            BatchRequestContentCollection]): The batch request content.
+            response_type: Optional[Type[T]] : The type to deserialize the response into.
+            Optional[Dict[str, Type[Parsable]]] = None: 
+                Error mappings for response handling.
+
+        Returns:
+            Union[T, BatchResponseContentCollection]: The batch response content
+             or the specified response type.
+
+        """
+        if batch_request_content is None:
+            raise ValueError("batch_request_content cannot be Null.")
+        response_type = BatchResponseContent
+
+        if isinstance(batch_request_content, BatchRequestContent):
+            request_info = await self.to_post_request_information(batch_request_content)
+            bytes_content = request_info.content
+            json_content = bytes_content.decode("utf-8")
+            updated_str = '{"requests":' + json_content + '}'
+            updated_bytes = updated_str.encode("utf-8")
+            request_info.content = updated_bytes
+            error_map = error_map or self.error_map
+            response = None
+            try:
+                response = await self._request_adapter.send_async(
+                    request_info, response_type, error_map
+                )
+
+            except APIError as e:
+                logging.error("API Error: %s", e)
+                raise e
+            if response is None:
+                raise ValueError("Failed to get a valid response from the API.")
+            return response
+        if isinstance(batch_request_content, BatchRequestContentCollection):
+            batch_responses = await self._post_batch_collection(batch_request_content, error_map)
+            return batch_responses
+
+        raise ValueError("Invalid type for batch_request_content.")
+
+    async def _post_batch_collection(
+        self,
+        batch_request_content_collection: BatchRequestContentCollection,
+        error_map: Optional[Dict[str, Type[Parsable]]] = None,
+    ) -> BatchResponseContentCollection:
+        """
+        Sends a collection of batch requests and returns a collection of batch response contents.
+        
+        Args:
+            batch_request_content_collection (BatchRequestContentCollection): The 
+            collection of batch request contents.
+            Optional[Dict[str, Type[Parsable]]] = None: 
+                Error mappings for response handling.
+
+        Returns:
+            BatchResponseContentCollection: The collection of batch response contents.
+        """
+        if batch_request_content_collection is None:
+            raise ValueError("batch_request_content_collection cannot be Null.")
+
+        batch_responses = BatchResponseContentCollection()
+
+        for batch_request_content in batch_request_content_collection.batches:
+            request_info = await self.to_post_request_information(batch_request_content)
+            response = await self._request_adapter.send_async(
+                request_info, BatchResponseContent, error_map or self.error_map
+            )
+            batch_responses.add_response(response)
+
+        return batch_responses
+
+    async def to_post_request_information(
+        self, batch_request_content: BatchRequestContent
+    ) -> RequestInformation:
+        """
+        Creates request information for a batch POST request.
+        
+        Args:
+            batch_request_content (BatchRequestContent): The batch request content.
+
+        Returns:
+            RequestInformation: The request information.
+        """
+
+        if batch_request_content is None:
+            raise ValueError("batch_request_content cannot be Null.")
+        batch_request_items = list(batch_request_content.requests.values())
+
+        request_info = RequestInformation()
+        request_info.http_method = Method.POST
+        request_info.url_template = self.url_template
+        request_info.headers = HeadersCollection()
+        request_info.headers.try_add("Content-Type", APPLICATION_JSON)
+        request_info.set_content_from_parsable(
+            self._request_adapter, APPLICATION_JSON, batch_request_items
+        )
+
+        return request_info

src/msgraph_core/requests/batch_request_builder.py

@@ -0,0 +1,140 @@
+import uuid
+from typing import List, Dict, Union, Optional
+
+from kiota_abstractions.request_information import RequestInformation
+from kiota_abstractions.serialization import Parsable, ParseNode
+from kiota_abstractions.serialization import SerializationWriter
+
+from .batch_request_item import BatchRequestItem
+
+
+class BatchRequestContent(Parsable):
+    """
+    Provides operations to call the batch method.
+    """
+
+    MAX_REQUESTS = 20
+
+    def __init__(self, requests: Dict[str, Union['BatchRequestItem', 'RequestInformation']] = {}):
+        """
+        Initializes a new instance of the BatchRequestContent class.
+        """
+        self._requests: Dict[str, Union[BatchRequestItem, 'RequestInformation']] = requests or {}
+
+        self.is_finalized = False
+        for request_id, request in requests.items():
+            self.add_request(request_id, request)
+
+    @property
+    def requests(self) -> Dict:
+        """
+        Gets  the requests.
+        """
+        return self._requests
+
+    @requests.setter
+    def requests(self, requests: List[BatchRequestItem]) -> None:
+        """
+        Sets the requests.
+        """
+        if len(requests) >= BatchRequestContent.MAX_REQUESTS:
+            raise ValueError(f"Maximum number of requests is {BatchRequestContent.MAX_REQUESTS}")
+        for request in requests:
+            self.add_request(request.id, request)
+
+    def add_request(self, request_id: Optional[str], request: BatchRequestItem) -> None:
+        """
+        Adds a request to the batch request content.
+        """
+        if len(self.requests) >= BatchRequestContent.MAX_REQUESTS:
+            raise RuntimeError(f"Maximum number of requests is {BatchRequestContent.MAX_REQUESTS}")
+        if not request.id:
+            request.id = str(uuid.uuid4())
+        if hasattr(request, 'depends_on') and request.depends_on:
+            for dependent_id in request.depends_on:
+                if dependent_id not in [req.id for req in self.requests]:
+                    dependent_request = self._request_by_id(dependent_id)
+                    if dependent_request:
+                        self._requests[dependent_id] = dependent_request
+        self._requests[request.id] = request
+
+    def add_request_information(self, request_information: RequestInformation) -> None:
+        """ 
+        Adds a request to the batch request content.
+        Args:
+            request_information (RequestInformation): The request information to add.
+        """
+        request_id = str(uuid.uuid4())
+        self.add_request(request_id, BatchRequestItem(request_information))
+
+    def add_urllib_request(self, request) -> None:
+        """
+        Adds a request to the batch request content.
+        """
+        request_id = str(uuid.uuid4())
+        self.add_request(request_id, BatchRequestItem.create_with_urllib_request(request))
+
+    def remove(self, request_id: str) -> None:
+        """
+        Removes a request from the batch request content.
+        Also removes the request from the depends_on list of 
+            other requests.
+        """
+        request_to_remove = None
+        for request in self.requests:
+            if request.id == request_id:
+                request_to_remove = request
+            if hasattr(request, 'depends_on') and request.depends_on:
+                if request_id in request.depends_on:
+                    request.depends_on.remove(request_id)
+        if request_to_remove:
+            del self._requests[request_to_remove.id]
+        else:
+            raise ValueError(f"Request ID {request_id} not found in requests.")
+
+    def remove_batch_request_item(self, item: BatchRequestItem) -> None:
+        """
+        Removes a request from the batch request content.
+        """
+        self.remove(item.id)
+
+    def finalize(self):
+        """
+        Finalizes the batch request content.
+        """
+        self.is_finalized = True
+        return self._requests
+
+    def _request_by_id(self, request_id: str) -> Optional[BatchRequestItem]:
+        """
+        Finds a request by its ID.
+        
+        Args:
+            request_id (str): The ID of the request to find.
+
+        Returns:
+            The request with the given ID, or None if not found.
+        """
+        return self._requests.get(request_id)
+
+    @staticmethod
+    def create_from_discriminator_value(
+        parse_node: Optional[ParseNode] = None
+    ) -> 'BatchRequestContent':
+        if parse_node is None:
+            raise ValueError("parse_node cannot be None")
+        return BatchRequestContent()
+
+    def get_field_deserializers(self, ) -> Dict:
+        """
+        The deserialization information for the current model
+        """
+        return {}
+
+    def serialize(self, writer: SerializationWriter) -> None:
+        """
+        Serializes information the current object
+        Args:
+            writer: Serialization writer to use to serialize this model
+        """
+        writer.write_collection_of_object_values("requests", self.requests)

src/msgraph_core/requests/batch_request_content.py

@@ -0,0 +1,84 @@
+from typing import List, Optional
+
+from kiota_abstractions.request_information import RequestInformation
+from kiota_abstractions.serialization import SerializationWriter
+
+from .batch_request_content import BatchRequestContent
+from .batch_request_item import BatchRequestItem
+
+
+class BatchRequestContentCollection:
+    """A collection of request content objects."""
+
+    def __init__(self) -> None:
+        """
+        Initializes a new instance of the BatchRequestContentCollection class.
+         
+        
+        """
+        self.max_requests_per_batch = BatchRequestContent.MAX_REQUESTS
+        self.batches: List[BatchRequestContent] = []
+        self.current_batch: BatchRequestContent = BatchRequestContent()
+
+    def add_batch_request_item(self, request: BatchRequestItem) -> None:
+        """ 
+        Adds a request item to the collection.
+        Args:
+            request (BatchRequestItem): The request item to add.
+        """
+        if len(self.current_batch.requests) >= self.max_requests_per_batch:
+            self.batches.append(self.current_batch.finalize())
+            self.current_batch = BatchRequestContent()
+        self.current_batch.add_request(request.id, request)
+        self.batches.append(self.current_batch)
+
+    def remove_batch_request_item(self, request_id: str) -> None:
+        """ 
+        Removes a request item from the collection.
+        Args:
+            request_id (str): The ID of the request item to remove.
+        """
+        for batch in self.batches:
+            if request_id in batch.requests:
+                del batch.requests[request_id]
+                return
+        if request_id in self.current_batch.requests:
+            del self.current_batch.requests[request_id]
+
+    def new_batch_with_failed_requests(self) -> Optional[BatchRequestContent]:
+        """
+        Creates a new batch with failed requests.
+        Returns:
+            Optional[BatchRequestContent]: A new batch with failed requests.
+        """
+        # Use IDs to get response status codes, generate new batch with failed requests
+        batch_with_failed_responses: Optional[BatchRequestContent] = BatchRequestContent()
+        for batch in self.batches:
+            for request in batch.requests:
+                if request.status_code not in [200, 201, 202, 203, 204, 205, 206, 207, 208, 226]:
+                    if batch_with_failed_responses is not None:
+                        batch_with_failed_responses.add_request(request.id, request)
+                    else:
+                        raise ValueError("batch_with_failed_responses is None")
+        return batch_with_failed_responses
+
+    def get_batch_requests_for_execution(self) -> List[BatchRequestContent]:
+        """
+        Gets the batch requests for execution.
+        Returns:
+            List[BatchRequestContent]: The batch requests for execution.
+        """
+        # if not self.current_batch.is_finalized:
+        #     self.current_batch.finalize()
+        #     self.batches.append(self.current_batch)
+        return self.batches
+
+    def serialize(self, writer: SerializationWriter) -> None:
+        """
+        Serializes information the current object
+        Args:
+            writer: Serialization writer to use to serialize this model
+        """
+        pass
+        # print(f"serializing {self.batches}")
+        # writer.write_collection_of_object_values("requests", self.batches)