Add messaging send_all and send_multicast functions

Author: ZachOrr

firebase_admin/_messaging_utils.py

@@ -54,6 +54,26 @@ def __init__(self, data=None, notification=None, android=None, webpush=None, apn
         self.condition = condition
 
 
+class MulticastMessage(Message):
+    """A message that can be sent to multiple tokens via Firebase Cloud Messaging.
+
+    Contains payload information as well as recipient information. In particular, the message must
+    contain exactly one of token, topic or condition fields.
+
+    Args:
+        tokens: A list of registration token of the device to which the message should be sent (optional).
+        data: A dictionary of data fields (optional). All keys and values in the dictionary must be
+            strings.
+        notification: An instance of ``messaging.Notification`` (optional).
+        android: An instance of ``messaging.AndroidConfig`` (optional).
+        webpush: An instance of ``messaging.WebpushConfig`` (optional).
+        apns: An instance of ``messaging.ApnsConfig`` (optional).
+    """
+    def __init__(self, tokens=[], data=None, notification=None, android=None, webpush=None, apns=None):
+        super(MulticastMessage, self).__init__(data=data, notification=notification, android=android, webpush=webpush, apns=apns)
+        self.tokens = tokens
+
+
 class Notification(object):
     """A notification that can be included in a message.
 

firebase_admin/messaging.py

@@ -16,6 +16,10 @@
 
 import requests
 import six
+import threading
+
+import googleapiclient
+from googleapiclient.discovery import build
 
 import firebase_admin
 from firebase_admin import _http_client
@@ -34,17 +38,22 @@
     'ApiCallError',
     'Aps',
     'ApsAlert',
+    'BatchResponse',
     'CriticalSound',
     'ErrorInfo',
     'Message',
+    'MulticastMessage',
     'Notification',
+    'SendResponse',
     'TopicManagementResponse',
     'WebpushConfig',
     'WebpushFcmOptions',
     'WebpushNotification',
     'WebpushNotificationAction',
 
     'send',
+    'send_all',
+    'send_multicast',
     'subscribe_to_topic',
     'unsubscribe_from_topic',
 ]
@@ -58,6 +67,7 @@
 ApsAlert = _messaging_utils.ApsAlert
 CriticalSound = _messaging_utils.CriticalSound
 Message = _messaging_utils.Message
+MulticastMessage = _messaging_utils.MulticastMessage
 Notification = _messaging_utils.Notification
 WebpushConfig = _messaging_utils.WebpushConfig
 WebpushFcmOptions = _messaging_utils.WebpushFcmOptions
@@ -88,6 +98,54 @@ def send(message, dry_run=False, app=None):
     """
     return _get_messaging_service(app).send(message, dry_run)
 
+def send_all(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:
+        BatchResponse: A ``messaging.BatchResponse`` instance.
+
+    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).send_all(messages, dry_run)
+
+def send_multicast(multicast_message, dry_run=False, app=None):
+    """Sends the given mutlicast message to the mutlicast message tokens 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:
+        message: An instance of ``messaging.MulticastMessage``.
+        dry_run: A boolean indicating whether to run the operation in dry run mode (optional).
+        app: An App instance (optional).
+
+    Returns:
+        BatchResponse: A ``messaging.BatchResponse`` instance.
+
+    Raises:
+        ApiCallError: If an error occurs while sending the message to FCM service.
+        ValueError: If the input arguments are invalid.
+    """
+    messages = map(lambda token: Message(
+        data=multicast_message.data,
+        notification=multicast_message.notification,
+        android=multicast_message.android,
+        webpush=multicast_message.webpush,
+        apns=multicast_message.apns,
+        token=token
+    ), multicast_message.tokens)
+    return _get_messaging_service(app).send_all(messages, dry_run)
+
 def subscribe_to_topic(tokens, topic, app=None):
     """Subscribes a list of registration tokens to an FCM topic.
 
@@ -192,10 +250,72 @@ def __init__(self, code, message, detail=None):
         self.detail = detail
 
 
+class BatchResponse(object):
+
+    def __init__(self, responses):
+        if not isinstance(responses, list):
+            raise ValueError('Unexpected responses: {0}.'.format(responses))
+        self._responses = responses
+        self._success_count = 0
+        self._failure_count = 0
+        for response in responses:
+            if response.success:
+                self._success_count += 1
+            else:
+                self._failure_count += 1
+
+    @property
+    def responses(self):
+        """A list of ``messaging.SendResponse`` objects (possibly empty)."""
+        return self._responses
+
+    @property
+    def success_count(self):
+        return self._success_count
+
+    @property
+    def failure_count(self):
+        return self._failure_count
+
+
+class SendResponse(object):
+
+    def __init__(self, resp, exception):
+        if resp and not isinstance(resp, dict):
+            raise ValueError('Unexpected response: {0}.'.format(resp))
+
+        self._message_id = None
+        self._exception = None
+
+        if resp:
+            self._message_id = resp.get('name', None)
+
+        if exception:
+            if exception.content is not None:
+                self._exception = _MessagingService._parse_fcm_error(exception)
+            else:
+                msg = 'Failed to call messaging API: {0}'.format(exception)
+                self._exception = ApiCallError(_MessagingService.INTERNAL_ERROR, msg, exception)
+
+    @property
+    def message_id(self):
+        """A message ID string that uniquely identifies the sent the message."""
+        return self._message_id
+
+    @property
+    def success(self):
+        """A boolean indicating if the request was successful."""
+        return self._message_id is not None and not self._exception
+
+    @property
+    def exception(self):
+        """A ApiCallError if an error occurs while sending the message to FCM service."""
+        return self._exception
+
+
 class _MessagingService(object):
     """Service class that implements Firebase Cloud Messaging (FCM) functionality."""
 
-    FCM_URL = 'https://fcm.googleapis.com/v1/projects/{0}/messages:send'
     IID_URL = 'https://iid.googleapis.com'
     IID_HEADERS = {'access_token_auth': 'true'}
     JSON_ENCODER = _messaging_utils.MessageEncoder()
@@ -233,10 +353,14 @@ def __init__(self, app):
                 'Project ID is required to access Cloud Messaging service. Either set the '
                 'projectId option, or use service account credentials. Alternatively, set the '
                 'GOOGLE_CLOUD_PROJECT environment variable.')
-        self._fcm_url = _MessagingService.FCM_URL.format(project_id)
+        self._fcm_service = build('fcm', 'v1', credentials=app.credential.get_credential())
+        self._fcm_parent = 'projects/{}'.format(project_id)
+        self._fcm_headers = {
+            'X-GOOG-API-FORMAT-VERSION': '2',
+            'X-FIREBASE-CLIENT': 'fire-admin-python/{0}'.format(firebase_admin.__version__)
+        }
         self._client = _http_client.JsonHttpClient(credential=app.credential.get_credential())
         self._timeout = app.options.get('httpTimeout')
-        self._client_version = 'fire-admin-python/{0}'.format(firebase_admin.__version__)
 
     @classmethod
     def encode_message(cls, message):
@@ -245,25 +369,37 @@ def encode_message(cls, message):
         return cls.JSON_ENCODER.default(message)
 
     def send(self, message, dry_run=False):
-        data = {'message': _MessagingService.encode_message(message)}
-        if dry_run:
-            data['validate_only'] = True
+        request = self._message_request(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)
-        except requests.exceptions.RequestException as error:
-            if error.response is not None:
-                self._handle_fcm_error(error)
+            resp = request.execute()
+        except googleapiclient.errors.HttpError as error:
+            if error.content is not None:
+                raise _MessagingService._parse_fcm_error(error)
             else:
                 msg = 'Failed to call messaging API: {0}'.format(error)
                 raise ApiCallError(self.INTERNAL_ERROR, msg, error)
         else:
             return resp['name']
 
+    def send_all(self, messages, dry_run=False):
+        message_count = len(messages)
+        send_all_complete = threading.Event()
+        responses = []
+
+        def send_all_callback(request_id, response, exception):
+            send_response = SendResponse(response, exception)
+            responses.append(send_response)
+            if len(responses) == message_count:
+                send_all_complete.set()
+
+        batch = self._fcm_service.new_batch_http_request(callback=send_all_callback)
+        for message in messages:
+            batch.add(self._message_request(message, dry_run))
+        batch.execute()
+
+        send_all_complete.wait()
+        return BatchResponse(responses)
+
     def make_topic_management_request(self, tokens, topic, operation):
         """Invokes the IID service for topic management functionality."""
         if isinstance(tokens, six.string_types):
@@ -299,11 +435,21 @@ def make_topic_management_request(self, tokens, topic, operation):
         else:
             return TopicManagementResponse(resp)
 
-    def _handle_fcm_error(self, error):
+    def _message_request(self, message, dry_run):
+        data = {'message': _MessagingService.encode_message(message)}
+        if dry_run:
+            data['validate_only'] = True
+        request = self._fcm_service.projects().messages().send(parent=self._fcm_parent, body=data)
+        request.headers.update(self._fcm_headers)
+        return request
+
+    @classmethod
+    def _parse_fcm_error(cls, error):
         """Handles errors received from the FCM API."""
         data = {}
         try:
-            parsed_body = error.response.json()
+            import json
+            parsed_body = json.loads(error.content)
             if isinstance(parsed_body, dict):
                 data = parsed_body
         except ValueError:
@@ -322,8 +468,8 @@ def _handle_fcm_error(self, error):
         msg = error_dict.get('message')
         if not msg:
             msg = 'Unexpected HTTP response with status: {0}; body: {1}'.format(
-                error.response.status_code, error.response.content.decode())
-        raise ApiCallError(code, msg, error)
+                error.resp.status, error.content)
+        return ApiCallError(code, msg, error)
 
     def _handle_iid_error(self, error):
         """Handles errors received from the Instance ID API."""

requirements.txt

@@ -6,6 +6,7 @@ tox >= 3.6.0
 
 cachecontrol >= 0.12.4
 google-api-core[grpc] >= 1.7.0, < 2.0.0dev; platform.python_implementation != 'PyPy'
+google-api-python-client >= 1.7.8
 google-cloud-firestore >= 0.31.0; platform.python_implementation != 'PyPy'
 google-cloud-storage >= 1.13.0
 six >= 1.6.1