Add messaging sendAll function

Author: ZachOrr

firebase_admin/messaging.py

@@ -14,6 +14,7 @@
 
 """Firebase Cloud Messaging module."""
 
+import json
 import requests
 import six
 
@@ -45,6 +46,7 @@
     'WebpushNotificationAction',
 
     'send',
+    'sendAll',
     'subscribe_to_topic',
     'unsubscribe_from_topic',
 ]
@@ -88,6 +90,26 @@ def send(message, dry_run=False, app=None):
     """
     return _get_messaging_service(app).send(message, dry_run)
 
+def sendAll(messages, dry_run=False, app=None):
+    """Batch sends the given messages via Firebase Cloud Messaging (FCM).
+
+    If the ``dry_run`` mode is enabled, the message will not be actually delivered to the
+    recipients. Instead FCM performs all the usual validations, and emulates the send operation.
+
+    Args:
+        messages: A list of ``messaging.Message`` instances.
+        dry_run: A boolean indicating whether to run the operation in dry run mode (optional).
+        app: An App instance (optional).
+
+    Returns:
+        list: A list of message ID string that uniquely identifies the sent the messages.
+
+    Raises:
+        ApiCallError: If an error occurs while sending the message to FCM service.
+        ValueError: If the input arguments are invalid.
+    """
+    return _get_messaging_service(app).sendAll(messages, dry_run)
+
 def subscribe_to_topic(tokens, topic, app=None):
     """Subscribes a list of registration tokens to an FCM topic.
 
@@ -196,10 +218,13 @@ class _MessagingService(object):
     """Service class that implements Firebase Cloud Messaging (FCM) functionality."""
 
     FCM_URL = 'https://fcm.googleapis.com/v1/projects/{0}/messages:send'
+    FCM_BATCH_URL = 'https://fcm.googleapis.com/batch'
     IID_URL = 'https://iid.googleapis.com'
     IID_HEADERS = {'access_token_auth': 'true'}
     JSON_ENCODER = _messaging_utils.MessageEncoder()
 
+    BATCH_PART_BOUNDARY = '__END_OF_PART__'
+
     INTERNAL_ERROR = 'internal-error'
     UNKNOWN_ERROR = 'unknown-error'
     FCM_ERROR_CODES = {
@@ -235,26 +260,25 @@ def __init__(self, app):
                 'GOOGLE_CLOUD_PROJECT environment variable.')
         self._fcm_url = _MessagingService.FCM_URL.format(project_id)
         self._client = _http_client.JsonHttpClient(credential=app.credential.get_credential())
+        self._fcm_headers = {
+            'X-GOOG-API-FORMAT-VERSION': '2',
+            'X-FIREBASE-CLIENT': 'fire-admin-python/{0}'.format(firebase_admin.__version__)
+        }
         self._timeout = app.options.get('httpTimeout')
-        self._client_version = 'fire-admin-python/{0}'.format(firebase_admin.__version__)
 
     @classmethod
     def encode_message(cls, message):
         if not isinstance(message, Message):
             raise ValueError('Message must be an instance of messaging.Message class.')
         return cls.JSON_ENCODER.default(message)
 
+    # TODO: sendMulticast
+
     def send(self, message, dry_run=False):
-        data = {'message': _MessagingService.encode_message(message)}
-        if dry_run:
-            data['validate_only'] = True
+        data = self._message_data(message, dry_run)
         try:
-            headers = {
-                'X-GOOG-API-FORMAT-VERSION': '2',
-                'X-FIREBASE-CLIENT': self._client_version,
-            }
             resp = self._client.body(
-                'post', url=self._fcm_url, headers=headers, json=data, timeout=self._timeout)
+                'post', url=self._fcm_url, headers=self._fcm_headers, json=data, timeout=self._timeout)
         except requests.exceptions.RequestException as error:
             if error.response is not None:
                 self._handle_fcm_error(error)
@@ -264,6 +288,35 @@ def send(self, message, dry_run=False):
         else:
             return resp['name']
 
+    def sendAll(self, messages, dry_run=False):
+        data = self._batch_message_data(messages, dry_run)
+        headers = {
+            'Content-Type': 'multipart/mixed; boundary={}'.format(_MessagingService.BATCH_PART_BOUNDARY)
+        }
+        try:
+            response = self._client.request(
+                'post', url=_MessagingService.FCM_BATCH_URL, headers=headers, data=data, timeout=self._timeout
+            )
+        except requests.exceptions.RequestException as error:
+            if error.response is not None:
+                self._handle_fcm_error(error)
+            else:
+                msg = 'Failed to call messaging API: {0}'.format(error)
+                raise ApiCallError(self.INTERNAL_ERROR, msg, error)
+        else:
+            import re
+            # Split our batches responses in to groups
+            response_regex = re.compile('--batch.*?$\n(.*?)(?=--batch)', re.DOTALL | re.MULTILINE)
+            responses = response_regex.findall(response.content)
+
+            # First section is headers related to the batch request
+            # Second section is headers related to the batched request
+            # Third section is data related to the batched request
+            response_sections = [r.split('\r\n\r\n') for r in responses]
+            response_json = [json.loads(r[2]) for r in response_sections]
+
+            return response_json
+
     def make_topic_management_request(self, tokens, topic, operation):
         """Invokes the IID service for topic management functionality."""
         if isinstance(tokens, six.string_types):
@@ -299,6 +352,39 @@ def make_topic_management_request(self, tokens, topic, operation):
         else:
             return TopicManagementResponse(resp)
 
+    def _message_data(self, message, dry_run):
+        data = {'message': _MessagingService.encode_message(message)}
+        if dry_run:
+            data['validate_only'] = True
+        return data
+
+    def _batch_message_data(self, messages, dry_run):
+        parts = [self._batch_request_part(request, dry_run, _MessagingService.BATCH_PART_BOUNDARY, index)\
+            for (index, request) in enumerate(messages)]
+        return '{}--{}--\r\n'.format(''.join(parts), _MessagingService.BATCH_PART_BOUNDARY)
+
+    def _batch_request_part(self, request, dry_run, boundary, index):
+        data = self._batch_request_part_data(request, dry_run)
+        part = '--{}\r\n'.format(boundary)
+        part += 'Content-Length: {}\r\n'.format(len(data))
+        part += 'Content-Type: application/http\r\n'
+        part += 'content-id: {}\r\n'.format(index + 1)
+        part += 'content-transfer-encoding: binary\r\n'
+        part += '\r\n'
+        part += '{}\r\n'.format(data)
+        return part
+
+    def _batch_request_part_data(self, request, dry_run):
+        body = json.dumps(self._message_data(request, dry_run))
+        data = 'POST {} HTTP/1.1\r\n'.format(self._fcm_url)
+        data += 'Content-Length: {}\r\n'.format(len(body))
+        data += 'Content-Type: application/json; charset=UTF-8\r\n'
+        data += '\r\n'.join(['{}: {}'.format(k, v) for k, v in self._fcm_headers.items()])
+        data += '\r\n\r\n'
+        data += body
+        return data
+
+
     def _handle_fcm_error(self, error):
         """Handles errors received from the FCM API."""
         data = {}