在GCP Dataflow中集成Google Retail API的实践指南

# setup.py 或 requirements.txt 中
google-cloud-retail
google-cloud-retail-v2 # 推荐使用v2版本
apache-beam[gcp]

import apache_beam as beam
from apache_beam import DoFn
from google.cloud import retail_v2
from google.protobuf import timestamp_pb2
import datetime

class WriteRetailUserEventFn(DoFn):
    def __init__(self, project_id: str, location: str = "global", catalog_id: str = "default_catalog"):
        """
        初始化DoFn，传入项目ID、位置和目录ID。
        这些参数在DoFn实例化时传递，而非在setup中。
        """
        self.project_id = project_id
        self.location = location
        self.catalog_id = catalog_id
        self.user_event_client = None
        self.parent_path = None

    def setup(self):
        """
        在每个工作器实例启动时初始化Retail API客户端。
        Dataflow的Service Account将隐式处理认证。
        """
        self.user_event_client = retail_v2.UserEventServiceClient()
        self.parent_path = f"projects/{self.project_id}/locations/{self.location}/catalogs/{self.catalog_id}"

    def process(self, element: dict):
        """
        处理PCollection中的每个元素，将其转换为Retail用户事件并写入API。
        'element' 预期是一个字典，包含用户事件数据。
        """
        try:
            # 构造UserEvent对象
            user_event = retail_v2.UserEvent(
                event_type=element.get("event_type"),
                visitor_id=element.get("visitor_id"),
                event_time=self._to_timestamp_proto(element.get("event_time")), # 转换时间戳格式
                product_details=[
                    retail_v2.ProductDetail(product=f"projects/{self.project_id}/locations/{self.location}/catalogs/{self.catalog_id}/products/{pid}")
                    for pid in element.get("product_ids", [])
                ],
                uri=element.get("uri"),
                referrer_uri=element.get("referrer_uri"),
                page_view_id=element.get("page_view_id"),
                # 根据您的数据模式和Retail API要求添加其他相关字段
            )

            # 调用Retail API写入用户事件
            response = self.user_event_client.write_user_event(parent=self.parent_path, user_event=user_event)

            # Yield响应或确认消息，供下游处理/日志记录
            yield f"Successfully wrote user event for visitor_id: {user_event.visitor_id}, event_type: {user_event.event_type}"

        except Exception as e:
            # 记录错误，并可能将失败的元素发送到死信队列
            beam.metrics.Metrics.counter('retail_api_errors', 'write_event_failed').inc()
            print(f"Error writing Retail user event for element {element}: {e}")
            # 考虑yield一个错误对象或使用侧输出（Side Output）进行错误处理
            # 示例: yield beam.pvalue.TaggedOutput('errors', {'element': element, 'error': str(e)})

    def _to_timestamp_proto(self, dt_obj):
        """
        辅助方法：将datetime对象或ISO格式字符串转换为protobuf Timestamp。
        """
        if dt_obj is None:
            return None
        if isinstance(dt_obj, datetime.datetime):
            timestamp = timestamp_pb2.Timestamp()
            timestamp.FromDatetime(dt_obj)
            return timestamp
        elif isinstance(dt_obj, str):
            try:
                # 假设是ISO格式字符串，如 "2023-10-27T10:00:00Z"
                dt_obj = datetime.datetime.fromisoformat(dt_obj.replace('Z', '+00:00'))
                timestamp = timestamp_pb2.Timestamp()
                timestamp.FromDatetime(dt_obj)
                return timestamp
            except ValueError:
                # 如果无法解析，可以返回None或抛出错误
                return None
        return None

# 假设您已经定义了WriteRetailUserEventFn类

# with beam.Pipeline() as pipeline:
#     user_events_data = [
#         {"event_type": "page-view", "visitor_id": "user1", "event_time": datetime.datetime.now(), "uri": "/product/A"},
#         {"event_type": "add-to-cart", "visitor_id": "user2", "event_time": "2023-10-27T10:30:00Z", "product_ids": ["P123"]},
#     ]

#     results = (
#         pipeline
#         | 'CreateUserEvents' >> beam.Create(user_events_data)
#         | 'WriteToRetailAPI' >> beam.ParDo(WriteRetailUserEventFn(project_id="your-gcp-project-id"))
#         | 'LogResults' >> beam.Map(print)
#     )
#     pipeline.run().wait_until_finish()