feat: add init function support (#188)

Author: exaby73

src/firebase_functions/alerts_fn.py

@@ -24,7 +24,7 @@
 
 import firebase_functions.private.util as _util
 
-from firebase_functions.core import T, CloudEvent as _CloudEvent
+from firebase_functions.core import T, CloudEvent as _CloudEvent, _with_init
 from firebase_functions.options import FirebaseAlertOptions
 
 # Explicitly import AlertType to make it available in the public API.
@@ -95,7 +95,7 @@ def on_alert_published_inner_decorator(func: OnAlertPublishedCallable):
         @_functools.wraps(func)
         def on_alert_published_wrapped(raw: _ce.CloudEvent):
             from firebase_functions.private._alerts_fn import alerts_event_from_ce
-            func(alerts_event_from_ce(raw))
+            _with_init(func)(alerts_event_from_ce(raw))
 
         _util.set_func_endpoint_attr(
             on_alert_published_wrapped,

src/firebase_functions/core.py

@@ -18,6 +18,8 @@
 import datetime as _datetime
 import typing as _typing
 
+from . import logger as _logger
+
 T = _typing.TypeVar("T")
 
 
@@ -80,3 +82,46 @@ class Change(_typing.Generic[T]):
     """
     The state of data after the change.
     """
+
+
+_did_init = False
+_init_callback: _typing.Callable[[], _typing.Any] | None = None
+
+
+def init(callback: _typing.Callable[[], _typing.Any]) -> None:
+    """
+     Registers a function that should be run when in a production environment
+     before executing any functions code.
+     Calling this decorator more than once leads to undefined behavior.
+    """
+
+    global _did_init
+    global _init_callback
+
+    if _did_init:
+        _logger.warn(
+            "Setting init callback more than once. Only the most recent callback will be called"
+        )
+
+    _init_callback = callback
+    _did_init = False
+
+
+def _with_init(
+    fn: _typing.Callable[...,
+                         _typing.Any]) -> _typing.Callable[..., _typing.Any]:
+    """
+    A decorator that runs the init callback before running the decorated function.
+    """
+
+    def wrapper(*args, **kwargs):
+        global _did_init
+
+        if not _did_init:
+            if _init_callback is not None:
+                _init_callback()
+            _did_init = True
+
+        return fn(*args, **kwargs)
+
+    return wrapper

src/firebase_functions/db_fn.py

@@ -119,7 +119,7 @@ def _db_endpoint_handler(
         subject=event_attributes["subject"],
         params=params,
     )
-    func(database_event)
+    _core._with_init(func)(database_event)
 
 
 @_util.copy_func_kwargs(DatabaseOptions)

src/firebase_functions/eventarc_fn.py

@@ -21,7 +21,7 @@
 
 import firebase_functions.options as _options
 import firebase_functions.private.util as _util
-from firebase_functions.core import CloudEvent
+from firebase_functions.core import CloudEvent, _with_init
 
 
 @_util.copy_func_kwargs(_options.EventarcTriggerOptions)
@@ -73,7 +73,7 @@ def on_custom_event_published_wrapped(raw: _ce.CloudEvent):
                 ),
                 type=event_dict["type"],
             )
-            func(event)
+            _with_init(func)(event)
 
         _util.set_func_endpoint_attr(
             on_custom_event_published_wrapped,

src/firebase_functions/firestore_fn.py

@@ -205,6 +205,8 @@ def _firestore_endpoint_handler(
         params=params,
     )
 
+    func = _core._with_init(func)
+
     if event_type.endswith(".withAuthContext"):
         database_event_with_auth_context = AuthEvent(**vars(database_event),
                                                      auth_type=event_auth_type,

src/firebase_functions/https_fn.py

@@ -403,7 +403,7 @@ def _on_call_handler(func: _C2,
                 instance_id_token=request.headers.get(
                     "Firebase-Instance-ID-Token"),
             )
-        result = func(context)
+        result = _core._with_init(func)(context)
         return _jsonify(result=result)
     # Disable broad exceptions lint since we want to handle all exceptions here
     # and wrap as an HttpsError.
@@ -447,7 +447,7 @@ def on_request_wrapped(request: Request) -> Response:
                     methods=options.cors.cors_methods,
                     origins=options.cors.cors_origins,
                 )(func)(request)
-            return func(request)
+            return _core._with_init(func)(request)
 
         _util.set_func_endpoint_attr(
             on_request_wrapped,

src/firebase_functions/private/_identity_fn.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Cloud functions to handle Eventarc events."""
-
 # pylint: disable=protected-access
 import typing as _typing
 import datetime as _dt
 import time as _time
 import json as _json
+
+from firebase_functions.core import _with_init
 from firebase_functions.https_fn import HttpsError, FunctionsErrorCode
 
 import firebase_functions.private.util as _util
@@ -351,8 +352,8 @@ def before_operation_handler(
         jwt_token = request.json["data"]["jwt"]
         decoded_token = _token_verifier.verify_auth_blocking_token(jwt_token)
         event = _auth_blocking_event_from_token_data(decoded_token)
-        auth_response: BeforeCreateResponse | BeforeSignInResponse | None = func(
-            event)
+        auth_response: BeforeCreateResponse | BeforeSignInResponse | None = _with_init(
+            func)(event)
         if not auth_response:
             return _jsonify({})
         auth_response_dict = _validate_auth_response(event_type, auth_response)
@@ -362,7 +363,7 @@ def before_operation_handler(
     # pylint: disable=broad-except
     except Exception as exception:
         if not isinstance(exception, HttpsError):
-            _logging.error("Unhandled error", exception)
+            _logging.error("Unhandled error %s", exception)
             exception = HttpsError(FunctionsErrorCode.INTERNAL, "INTERNAL")
         status = exception._http_error_code.status
         return _make_response(_jsonify(error=exception._as_dict()), status)

src/firebase_functions/pubsub_fn.py

@@ -25,7 +25,7 @@
 
 import firebase_functions.private.util as _util
 
-from firebase_functions.core import CloudEvent, T
+from firebase_functions.core import CloudEvent, T, _with_init
 from firebase_functions.options import PubSubOptions
 
 
@@ -151,7 +151,7 @@ def _message_handler(
         type=event_dict["type"],
     )
 
-    func(event)
+    _with_init(func)(event)
 
 
 @_util.copy_func_kwargs(PubSubOptions)

src/firebase_functions/remote_config_fn.py

@@ -24,7 +24,7 @@
 
 import firebase_functions.private.util as _util
 
-from firebase_functions.core import CloudEvent
+from firebase_functions.core import CloudEvent, _with_init
 from firebase_functions.options import EventHandlerOptions
 
 
@@ -189,7 +189,7 @@ def _config_handler(func: _C1, raw: _ce.CloudEvent) -> None:
         type=event_dict["type"],
     )
 
-    func(event)
+    _with_init(func)(event)
 
 
 @_util.copy_func_kwargs(EventHandlerOptions)

src/firebase_functions/scheduler_fn.py

@@ -27,6 +27,7 @@
     make_response as _make_response,
 )
 
+from firebase_functions.core import _with_init
 # Export for user convenience.
 # pylint: disable=unused-import
 from firebase_functions.options import Timezone
@@ -108,7 +109,7 @@ def on_schedule_wrapped(request: _Request) -> _Response:
                 schedule_time=schedule_time,
             )
             try:
-                func(event)
+                _with_init(func)(event)
                 return _make_response()
             # Disable broad exceptions lint since we want to handle all exceptions.
             # pylint: disable=broad-except