fix local devlopment with the Pub/Sub emulator

Author: matthewrobertson

src/functions_framework/__init__.py

@@ -126,7 +126,7 @@ def view_func(path):
             function(data, context)
         else:
             # This is a regular CloudEvent
-            event_data = request.get_json()
+            event_data = event_conversion.marshal_background_event_data(request)
             if not event_data:
                 flask.abort(400)
             event_object = BackgroundEvent(**event_data)

src/functions_framework/event_conversion.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 import re
 
-from typing import Tuple
+from datetime import datetime
+from typing import Optional, Tuple
 
 from cloudevents.http import CloudEvent
 
@@ -55,6 +56,12 @@
 _PUBSUB_CE_SERVICE = "pubsub.googleapis.com"
 _STORAGE_CE_SERVICE = "storage.googleapis.com"
 
+# Raw pubsub types
+_PUBSUB_EVENT_TYPE = 'google.pubsub.topic.publish'
+_PUBSUB_MESSAGE_TYPE = 'type.googleapis.com/google.pubsub.v1.PubsubMessage'
+
+_PUBSUB_TOPIC_REQUEST_PATH = re.compile(r"projects\/[^/?]+\/topics\/[^/?]+")
+
 # Maps background event services to their equivalent CloudEvent services.
 _SERVICE_BACKGROUND_TO_CE = {
     "providers/cloud.firestore/": _FIRESTORE_CE_SERVICE,
@@ -90,7 +97,7 @@
 
 def background_event_to_cloudevent(request) -> CloudEvent:
     """Converts a background event represented by the given HTTP request into a CloudEvent. """
-    event_data = request.get_json()
+    event_data = marshal_background_event_data(request)
     if not event_data:
         raise EventConversionException("Failed to parse JSON")
 
@@ -109,6 +116,10 @@ def background_event_to_cloudevent(request) -> CloudEvent:
     # Handle Pub/Sub events.
     if service == _PUBSUB_CE_SERVICE:
         data = {"message": data}
+        # It is possible to configure a Pub/Sub subscription to push directly to this function
+        # without passing the topic name in the URL path.
+        if resource is None:
+            resource = ""
 
     # Handle Firebase Auth events.
     if service == _FIREBASE_AUTH_CE_SERVICE:
@@ -168,3 +179,50 @@ def _split_resource(context: Context) -> Tuple[str, str, str]:
         raise EventConversionException("Resource regex did not match")
 
     return service, match.group(1), match.group(2)
+
+
+def marshal_background_event_data(request):
+    """Marshal the request body of a raw Pub/Sub HTTP request into the schema that is expected of
+    a background event"""
+    request_data = request.get_json()
+    if not _is_raw_pubsub_payload(request_data):
+        # If this in not a raw Pub/Sub request, return the unaltered request data.
+        return request_data
+    
+    return {
+        "context": {
+            "eventId": request_data["message"]["messageId"],
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+            "eventType": _PUBSUB_EVENT_TYPE,
+            "resource": {
+                "service": _PUBSUB_CE_SERVICE,
+                "type": _PUBSUB_MESSAGE_TYPE,
+                "name": _parse_pubsub_topic(request.path),
+            },
+        },
+        "data": {
+            '@type': _PUBSUB_MESSAGE_TYPE,
+            "data": request_data["message"]["data"],
+            "attributes": request_data["message"]["attributes"],
+        }
+    }
+
+
+def _is_raw_pubsub_payload(request_data) -> bool:
+    """Does the given request body match the schema of a unmarshalled Pub/Sub request"""
+    return (
+        request_data is not None and
+        "context" not in request_data and
+        "subscription" in request_data and
+        "message" in request_data and
+        "data" in request_data["message"] and
+        "messageId" in request_data["message"]
+    )
+
+
+def _parse_pubsub_topic(request_path) -> Optional[str]:
+    match = _PUBSUB_TOPIC_REQUEST_PATH.search(request_path)
+    if match:
+        return match.group(0)
+    else:
+        return None

tests/test_convert.py

@@ -65,23 +65,67 @@
 
 BACKGROUND_RESOURCE_STRING = "projects/_/buckets/some-bucket/objects/folder/Test.cs"
 
+PUBSUB_CLOUD_EVENT = {
+    "specversion": "1.0",
+    "id": "1215011316659232",
+    "source": "//pubsub.googleapis.com/projects/sample-project/topics/gcf-test",
+    "time": "2020-05-18T12:13:19Z",
+    "type": "google.cloud.pubsub.topic.v1.messagePublished",
+    "datacontenttype": "application/json",
+    "data": {
+        "message": {
+            "data": "10",
+        },
+    },
+}
+
 
 @pytest.fixture
 def pubsub_cloudevent_output():
-    event = {
-        "specversion": "1.0",
-        "id": "1215011316659232",
-        "source": "//pubsub.googleapis.com/projects/sample-project/topics/gcf-test",
-        "time": "2020-05-18T12:13:19Z",
-        "type": "google.cloud.pubsub.topic.v1.messagePublished",
-        "datacontenttype": "application/json",
+    return from_json(json.dumps(PUBSUB_CLOUD_EVENT))
+
+@pytest.fixture
+def raw_pubsub_request():
+    return {
+        "subscription": "projects/sample-project/subscriptions/gcf-test-sub",
+        "message": {
+            "data": "eyJmb28iOiJiYXIifQ==",
+            "messageId": "1215011316659232",
+            "attributes": {
+                "test": "123"
+            }
+        }
+    }
+
+@pytest.fixture
+def marshalled_pubsub_request():
+    return {
         "data": {
-            "message": {
-                "data": "10",
-            },
+            "@type": "type.googleapis.com/google.pubsub.v1.PubsubMessage",
+            "data": "eyJmb28iOiJiYXIifQ==",
+            "attributes": {
+                "test": "123"
+            }
         },
+        "context": {
+            "eventId": "1215011316659232",
+            "eventType": "google.pubsub.topic.publish",
+            "resource": {
+                "name": "projects/sample-project/topics/gcf-test",
+                "service": "pubsub.googleapis.com",
+                "type": "type.googleapis.com/google.pubsub.v1.PubsubMessage"
+            },
+            "timestamp": "2021-04-17T07:21:18.249Z",
+        }
     }
 
+@pytest.fixture
+def raw_pubsub_cloudevent_output(marshalled_pubsub_request):
+    event = PUBSUB_CLOUD_EVENT.copy()
+    # the data payload is more complex for the raw pubsub request
+    event["data"] = {
+        "message": marshalled_pubsub_request["data"]
+    }
     return from_json(json.dumps(event))
 
 
@@ -212,3 +256,64 @@ def test_split_resource_no_resource_regex_match():
     with pytest.raises(EventConversionException) as exc_info:
         event_conversion._split_resource(context)
     assert "Resource regex did not match" in exc_info.value.args[0]
+
+
+def test_marshal_background_event_data_without_topic_in_path(
+    raw_pubsub_request, marshalled_pubsub_request
+):
+    req = flask.Request.from_values(json=raw_pubsub_request, path="/myfunc/")
+    payload = event_conversion.marshal_background_event_data(req)
+
+    # Remove timestamps as they get generates on the fly
+    del marshalled_pubsub_request["context"]["timestamp"]
+    del payload["context"]["timestamp"]
+
+    # Resource name is set to None when it cannot be parsed from the request path
+    marshalled_pubsub_request["context"]["resource"]["name"] = None
+
+    assert payload == marshalled_pubsub_request
+
+def test_marshal_background_event_data_with_topic_path(
+    raw_pubsub_request, marshalled_pubsub_request
+):
+    req = flask.Request.from_values(
+        json=raw_pubsub_request, path="x/projects/sample-project/topics/gcf-test?pubsub_trigger=true"
+    )
+    payload = event_conversion.marshal_background_event_data(req)
+
+    # Remove timestamps as they are generated on the fly.
+    del marshalled_pubsub_request["context"]["timestamp"]
+    del payload["context"]["timestamp"]
+
+    assert payload == marshalled_pubsub_request
+
+def test_pubsub_emulator_request_to_cloudevent(raw_pubsub_request, raw_pubsub_cloudevent_output):
+    req = flask.Request.from_values(
+        json=raw_pubsub_request,
+        path="x/projects/sample-project/topics/gcf-test?pubsub_trigger=true"
+    )
+    cloudevent = event_conversion.background_event_to_cloudevent(req)
+
+    # Remove timestamps as they are generated on the fly.
+    del raw_pubsub_cloudevent_output['time']
+    del cloudevent['time']
+
+    assert cloudevent == raw_pubsub_cloudevent_output
+
+
+def test_pubsub_emulator_request_to_cloudevent_without_topic_path(
+    raw_pubsub_request, raw_pubsub_cloudevent_output
+):
+    req = flask.Request.from_values(
+        json=raw_pubsub_request, path="/"
+    )
+    cloudevent = event_conversion.background_event_to_cloudevent(req)
+
+    # Remove timestamps as they are generated on the fly.
+    del raw_pubsub_cloudevent_output['time']
+    del cloudevent['time']
+
+    # Default to the service name, when the topic is not configured subscription's pushEndpoint.
+    raw_pubsub_cloudevent_output['source'] = "//pubsub.googleapis.com/"
+
+    assert cloudevent == raw_pubsub_cloudevent_output

tests/test_functions.py

@@ -48,6 +48,19 @@ def background_json(tmpdir):
         "data": {"filename": str(tmpdir / "filename.txt"), "value": "some-value"},
     }
 
+@pytest.fixture
+def pubsub_emulator_request_payload(tmpdir):
+    return {
+        "subscription": "projects/FOO/subscriptions/BAR_SUB",
+        "message": {
+            "data": 'eyJmb28iOiJiYXIifQ==',
+            "messageId": "1",
+            "attributes": {
+                "filename": str(tmpdir / "filename.txt"),
+                "value": "some-value"
+            }
+        }
+    }
 
 def test_http_function_executes_success():
     source = TEST_FUNCTIONS_DIR / "http_trigger" / "main.py"
@@ -242,6 +255,21 @@ def test_pubsub_payload(background_json):
             background_json["data"]["value"]
         )
 
+def test_pubsub_emulator_payload(pubsub_emulator_request_payload):
+    source = TEST_FUNCTIONS_DIR / "background_trigger" / "main.py"
+    target = "function"
+
+    client = create_app(target, source, "event").test_client()
+
+    resp = client.post("/", json=pubsub_emulator_request_payload)
+
+    assert resp.status_code == 200
+    assert resp.data == b"OK"
+
+    with open(pubsub_emulator_request_payload["message"]["attributes"]["filename"]) as f:
+        assert f.read() == '{{"entryPoint": "function", "value": "{}"}}'.format(
+            pubsub_emulator_request_payload["message"]["attributes"]["value"]
+        )
 
 def test_background_function_no_data(background_json):
     source = TEST_FUNCTIONS_DIR / "background_trigger" / "main.py"

tests/test_functions/background_trigger/main.py

@@ -29,8 +29,9 @@ def function(
         data dictionary.
       context (google.cloud.functions.Context): The Cloud Functions event context.
     """
-    filename = event["filename"]
-    value = event["value"]
+    attributes = event.get("attributes", {})
+    filename = event.get("filename", attributes.get("filename"))
+    value = event.get("value", attributes.get("value"))
     f = open(filename, "w")
     f.write('{{"entryPoint": "function", "value": "{}"}}'.format(value))
     f.close()