expose ack id in pubsub message (#326)

buildflow/io/gcp/strategies/pubsub_strategies.py

@@ -1,7 +1,7 @@
 import dataclasses
 import datetime
 import logging
-from typing import Any, Callable, Dict, Iterable, Optional, Type, Union
+from typing import Any, Callable, Iterable, Optional, Type, Union
 
 from google.cloud.monitoring_v3 import query
 from google.cloud.pubsub_v1.types import PubsubMessage as GCPPubSubMessage
@@ -21,19 +21,14 @@
 from buildflow.io.strategies.source import AckInfo, PullResponse, SourceStrategy
 from buildflow.io.utils.clients import gcp_clients
 from buildflow.io.utils.schemas import converters
+from buildflow.types.gcp import PubsubMessage
 
 
 @dataclasses.dataclass(frozen=True)
 class _PubsubAckInfo(AckInfo):
     ack_ids: Iterable[str]
 
 
-@dataclasses.dataclass
-class PubsubMessage:
-    data: bytes
-    attributes: Dict[str, Any]
-
-
 def _timestamp_to_datetime(timestamp: Union[datetime.datetime, Timestamp]):
     if isinstance(timestamp, Timestamp):
         return timestamp.ToDatetime()
@@ -97,7 +92,10 @@ async def pull(self) -> PullResponse:
                 attributes = received_message.message.attributes
                 for key, value in attributes.items():
                     att_dict[key] = value
-                payload = PubsubMessage(received_message.message.data, att_dict)
+
+                payload = PubsubMessage(
+                    received_message.message.data, att_dict, received_message.ack_id
+                )
             elif received_message.message.data:
                 payload = received_message.message.data
             else:

buildflow/types/gcp.py

@@ -1,5 +1,6 @@
 import dataclasses
 import enum
+from typing import Any, Dict
 
 from google.cloud import storage
 from pulumi_gcp.sql import (  # noqa
@@ -48,3 +49,10 @@ class CloudSQLDatabaseVersion:
     POSTGRES_13 = "POSTGRES_13"
     POSTGRES_14 = "POSTGRES_14"
     POSTGRES_15 = "POSTGRES_15"
+
+
+@dataclasses.dataclass
+class PubsubMessage:
+    data: bytes
+    attributes: Dict[str, Any]
+    ack_id: str