feat: Add Auth Context enabled Firestore triggers (#184)

Author: exaby73

src/firebase_functions/firestore_fn.py

@@ -38,6 +38,11 @@
 _event_type_updated = "google.cloud.firestore.document.v1.updated"
 _event_type_deleted = "google.cloud.firestore.document.v1.deleted"
 
+_event_type_written_with_auth_context = "google.cloud.firestore.document.v1.written.withAuthContext"
+_event_type_created_with_auth_context = "google.cloud.firestore.document.v1.created.withAuthContext"
+_event_type_updated_with_auth_context = "google.cloud.firestore.document.v1.updated.withAuthContext"
+_event_type_deleted_with_auth_context = "google.cloud.firestore.document.v1.deleted.withAuthContext"
+
 
 @_dataclass.dataclass(frozen=True)
 class Event(_core.CloudEvent[_core.T]):
@@ -82,9 +87,26 @@ class Event(_core.CloudEvent[_core.T]):
 _C1 = _typing.Callable[[_E1], None]
 _C2 = _typing.Callable[[_E2], None]
 
+AuthType = _typing.Literal["service_account", "api_key", "system",
+                           "unauthenticated", "unknown"]
+
+
+@_dataclass.dataclass(frozen=True)
+class AuthEvent(Event[_core.T]):
+    auth_type: AuthType
+    """The type of principal that triggered the event"""
+    auth_id: str | None
+    """The unique identifier for the principal"""
+
+
+_E3 = AuthEvent[Change[DocumentSnapshot | None]]
+_E4 = AuthEvent[DocumentSnapshot | None]
+_C3 = _typing.Callable[[_E3], None]
+_C4 = _typing.Callable[[_E4], None]
+
 
 def _firestore_endpoint_handler(
-    func: _C1 | _C2,
+    func: _C1 | _C2 | _C3 | _C4,
     event_type: str,
     document_pattern: _path_pattern.PathPattern,
     raw: _ce.CloudEvent,
@@ -94,12 +116,14 @@ def _firestore_endpoint_handler(
     firestore_event_data: _firestore.DocumentEventData
     content_type: str = event_attributes["datacontenttype"]
     if "application/json" in content_type or isinstance(event_data, dict):
-        firestore_event_data = _firestore.DocumentEventData.from_json(
-            event_data)
+        firestore_event_data = _typing.cast(
+            _firestore.DocumentEventData,
+            _firestore.DocumentEventData.from_json(event_data))
     elif "application/protobuf" in content_type or isinstance(
             event_data, bytes):
-        firestore_event_data = _firestore.DocumentEventData.deserialize(
-            event_data)
+        firestore_event_data = _typing.cast(
+            _firestore.DocumentEventData,
+            _firestore.DocumentEventData.deserialize(event_data))
     else:
         actual_type = type(event_data)
         raise TypeError(f"Firestore: Cannot parse event payload of data type "
@@ -110,6 +134,8 @@ def _firestore_endpoint_handler(
     event_namespace = event_attributes["namespace"]
     event_document = event_attributes["document"]
     event_database = event_attributes["database"]
+    event_auth_type = event_attributes["authtype"]
+    event_auth_id = event_attributes["authid"]
 
     time = event_attributes["time"]
     event_time = _util.timestamp_conversion(time)
@@ -146,18 +172,23 @@ def _firestore_endpoint_handler(
             firestore_event_data.old_value.update_time,
         )
     if event_type == _event_type_deleted:
-        firestore_event_data = old_value_snapshot
+        firestore_event_data = _typing.cast(_firestore.DocumentEventData,
+                                            old_value_snapshot)
     if event_type == _event_type_created:
-        firestore_event_data = value_snapshot
+        firestore_event_data = _typing.cast(_firestore.DocumentEventData,
+                                            value_snapshot)
     if event_type in (_event_type_written, _event_type_updated):
-        firestore_event_data = Change(
-            before=old_value_snapshot,
-            after=value_snapshot,
-        )
+        firestore_event_data = _typing.cast(
+            _firestore.DocumentEventData,
+            Change(
+                before=old_value_snapshot,
+                after=value_snapshot,
+            ))
 
     params: dict[str, str] = {
         **document_pattern.extract_matches(event_document),
     }
+
     database_event = Event(
         project=event_project,
         namespace=event_namespace,
@@ -173,7 +204,15 @@ def _firestore_endpoint_handler(
         subject=event_attributes["subject"],
         params=params,
     )
-    func(database_event)
+
+    if event_type.endswith(".withAuthContext"):
+        database_event_with_auth_context = AuthEvent(**vars(database_event),
+                                                     auth_type=event_auth_type,
+                                                     auth_id=event_auth_id)
+        func(database_event_with_auth_context)
+    else:
+        # mypy cannot infer that the event type is correct, hence the cast
+        _typing.cast(_C1 | _C2, func)(database_event)
 
 
 @_util.copy_func_kwargs(FirestoreOptions)
@@ -224,6 +263,57 @@ def on_document_written_wrapped(raw: _ce.CloudEvent):
     return on_document_written_inner_decorator
 
 
+@_util.copy_func_kwargs(FirestoreOptions)
+def on_document_written_with_auth_context(**kwargs
+                                         ) -> _typing.Callable[[_C1], _C1]:
+    """
+    Event handler that triggers when a document is created, updated, or deleted in Firestore.
+    This trigger will also provide the authentication context of the principal who triggered 
+    the event.
+
+    Example:
+
+    .. code-block:: python
+
+      @on_document_written_with_auth_context(document="*")
+      def example(event: AuthEvent[Change[DocumentSnapshot]]) -> None:
+          pass
+
+    :param \\*\\*kwargs: Firestore options.
+    :type \\*\\*kwargs: as :exc:`firebase_functions.options.FirestoreOptions`
+    :rtype: :exc:`typing.Callable`
+            \\[ \\[ :exc:`firebase_functions.firestore_fn.AuthEvent` \\[
+            :exc:`firebase_functions.db.Change` \\] \\], `None` \\]
+            A function that takes a Firestore event and returns ``None``.
+    """
+    options = FirestoreOptions(**kwargs)
+
+    def on_document_written_with_auth_context_inner_decorator(func: _C1):
+        document_pattern = _path_pattern.PathPattern(
+            _util.normalize_path(options.document))
+
+        @_functools.wraps(func)
+        def on_document_written_with_auth_context_wrapped(raw: _ce.CloudEvent):
+            return _firestore_endpoint_handler(
+                func,
+                _event_type_written_with_auth_context,
+                document_pattern,
+                raw,
+            )
+
+        _util.set_func_endpoint_attr(
+            on_document_written_with_auth_context_wrapped,
+            options._endpoint(
+                event_type=_event_type_written,
+                func_name=func.__name__,
+                document_pattern=document_pattern,
+            ),
+        )
+        return on_document_written_with_auth_context_wrapped
+
+    return on_document_written_with_auth_context_inner_decorator
+
+
 @_util.copy_func_kwargs(FirestoreOptions)
 def on_document_updated(**kwargs) -> _typing.Callable[[_C1], _C1]:
     """
@@ -272,6 +362,57 @@ def on_document_updated_wrapped(raw: _ce.CloudEvent):
     return on_document_updated_inner_decorator
 
 
+@_util.copy_func_kwargs(FirestoreOptions)
+def on_document_updated_with_auth_context(**kwargs
+                                         ) -> _typing.Callable[[_C1], _C1]:
+    """
+    Event handler that triggers when a document is updated in Firestore.
+    This trigger will also provide the authentication context of the principal who triggered 
+    the event.
+
+    Example:
+
+    .. code-block:: python
+
+      @on_document_updated_with_auth_context(document="*")
+      def example(event: AuthEvent[Change[DocumentSnapshot]]) -> None:
+          pass
+
+    :param \\*\\*kwargs: Firestore options.
+    :type \\*\\*kwargs: as :exc:`firebase_functions.options.FirestoreOptions`
+    :rtype: :exc:`typing.Callable`
+            \\[ \\[ :exc:`firebase_functions.firestore_fn.AuthEvent` \\[
+            :exc:`firebase_functions.db.Change` \\] \\], `None` \\]
+            A function that takes a Firestore event and returns ``None``.
+    """
+    options = FirestoreOptions(**kwargs)
+
+    def on_document_updated_with_auth_context_inner_decorator(func: _C1):
+        document_pattern = _path_pattern.PathPattern(
+            _util.normalize_path(options.document))
+
+        @_functools.wraps(func)
+        def on_document_updated_with_auth_context_wrapped(raw: _ce.CloudEvent):
+            return _firestore_endpoint_handler(
+                func,
+                _event_type_updated_with_auth_context,
+                document_pattern,
+                raw,
+            )
+
+        _util.set_func_endpoint_attr(
+            on_document_updated_with_auth_context_wrapped,
+            options._endpoint(
+                event_type=_event_type_updated_with_auth_context,
+                func_name=func.__name__,
+                document_pattern=document_pattern,
+            ),
+        )
+        return on_document_updated_with_auth_context_wrapped
+
+    return on_document_updated_with_auth_context_inner_decorator
+
+
 @_util.copy_func_kwargs(FirestoreOptions)
 def on_document_created(**kwargs) -> _typing.Callable[[_C2], _C2]:
     """
@@ -320,6 +461,57 @@ def on_document_created_wrapped(raw: _ce.CloudEvent):
     return on_document_created_inner_decorator
 
 
+@_util.copy_func_kwargs(FirestoreOptions)
+def on_document_created_with_auth_context(**kwargs
+                                         ) -> _typing.Callable[[_C2], _C2]:
+    """
+    Event handler that triggers when a document is created in Firestore.
+    This trigger will also provide the authentication context of the principal who triggered 
+    the event.
+
+    Example:
+
+    .. code-block:: python
+
+        @on_document_created_with_auth_context(document="*")
+        def example(event: AuthEvent[DocumentSnapshot]):
+          pass
+
+    :param \\*\\*kwargs: Firestore options.
+    :type \\*\\*kwargs: as :exc:`firebase_functions.options.FirestoreOptions`
+    :rtype: :exc:`typing.Callable`
+            \\[ \\[ :exc:`firebase_functions.firestore_fn.AuthEvent` \\[
+            :exc:`object` \\] \\], `None` \\]
+            A function that takes a Firestore event and returns ``None``.
+    """
+    options = FirestoreOptions(**kwargs)
+
+    def on_document_created_with_auth_context_inner_decorator(func: _C2):
+        document_pattern = _path_pattern.PathPattern(
+            _util.normalize_path(options.document))
+
+        @_functools.wraps(func)
+        def on_document_created_with_auth_context_wrapped(raw: _ce.CloudEvent):
+            return _firestore_endpoint_handler(
+                func,
+                _event_type_created_with_auth_context,
+                document_pattern,
+                raw,
+            )
+
+        _util.set_func_endpoint_attr(
+            on_document_created_with_auth_context_wrapped,
+            options._endpoint(
+                event_type=_event_type_created_with_auth_context,
+                func_name=func.__name__,
+                document_pattern=document_pattern,
+            ),
+        )
+        return on_document_created_with_auth_context_wrapped
+
+    return on_document_created_with_auth_context_inner_decorator
+
+
 @_util.copy_func_kwargs(FirestoreOptions)
 def on_document_deleted(**kwargs) -> _typing.Callable[[_C2], _C2]:
     """
@@ -366,3 +558,54 @@ def on_document_deleted_wrapped(raw: _ce.CloudEvent):
         return on_document_deleted_wrapped
 
     return on_document_deleted_inner_decorator
+
+
+@_util.copy_func_kwargs(FirestoreOptions)
+def on_document_deleted_with_auth_context(**kwargs
+                                         ) -> _typing.Callable[[_C2], _C2]:
+    """
+    Event handler that triggers when a document is deleted in Firestore.
+    This trigger will also provide the authentication context of the principal who triggered 
+    the event.
+
+    Example:
+
+    .. code-block:: python
+
+      @on_document_deleted_with_auth_context(document="*")
+      def example(event: AuthEvent[DocumentSnapshot]) -> None:
+          pass
+
+    :param \\*\\*kwargs: Firestore options.
+    :type \\*\\*kwargs: as :exc:`firebase_functions.options.FirestoreOptions`
+    :rtype: :exc:`typing.Callable`
+            \\[ \\[ :exc:`firebase_functions.firestore_fn.AuthEvent` \\[
+            :exc:`object` \\] \\], `None` \\]
+            A function that takes a Firestore event and returns ``None``.
+    """
+    options = FirestoreOptions(**kwargs)
+
+    def on_document_deleted_with_auth_context_inner_decorator(func: _C2):
+        document_pattern = _path_pattern.PathPattern(
+            _util.normalize_path(options.document))
+
+        @_functools.wraps(func)
+        def on_document_deleted_with_auth_context_wrapped(raw: _ce.CloudEvent):
+            return _firestore_endpoint_handler(
+                func,
+                _event_type_deleted_with_auth_context,
+                document_pattern,
+                raw,
+            )
+
+        _util.set_func_endpoint_attr(
+            on_document_deleted_with_auth_context_wrapped,
+            options._endpoint(
+                event_type=_event_type_deleted_with_auth_context,
+                func_name=func.__name__,
+                document_pattern=document_pattern,
+            ),
+        )
+        return on_document_deleted_with_auth_context_wrapped
+
+    return on_document_deleted_with_auth_context_inner_decorator