在Apache Beam Python管道中获取Google Datastore自动生成的实体ID

我来帮你搞定这个问题！要让Google Datastore自动生成实体ID，核心就是不要手动指定ID值，让Datastore在插入实体时自动分配唯一整数ID，之后再从插入响应中获取这个生成的ID就行。下面是具体的实现步骤和代码示例：

1. 修改实体构建逻辑：不指定ID #

你原来的代码手动传入了uuid.uuid4()作为ID，现在只需要在调用datastore_helper.add_key_path时，只传入key和实体类型（kind），省略ID参数即可：

from googledatastore import helper as datastore_helper
from googledatastore import entity_pb2

# 创建实体对象
entity = entity_pb2.Entity()
# 只指定实体类型，不传入ID参数（让Datastore自动生成）
datastore_helper.add_key_path(entity.key, entityName.get())

# 这里可以继续添加实体的属性，比如：
# datastore_helper.add_property(entity, "your_property", "property_value")

这样构建出来的实体key没有ID部分，Datastore在插入时会自动为你生成一个唯一的整数ID。

2. 在Apache Beam管道中插入实体并获取自动ID #

在Beam里，你可以通过ParDo函数结合Datastore客户端来插入实体，并直接从插入响应中拿到生成的ID。示例代码如下：

from apache_beam import DoFn
from googledatastore import client as datastore_client
from googledatastore import helper as datastore_helper
from googledatastore import entity_pb2

class WriteEntityWithAutoID(DoFn):
    def start_bundle(self):
        # 在bundle开始时初始化Datastore客户端，避免重复创建提升性能
        self.datastore_client = datastore_client.Client()

    def process(self, element):
        # 构建不带ID的实体
        entity = entity_pb2.Entity()
        datastore_helper.add_key_path(entity.key, entityName.get())
        
        # 假设element是包含实体属性的字典，遍历添加属性
        for prop_name, prop_value in element.items():
            datastore_helper.add_property(entity, prop_name, prop_value)
        
        # 执行插入操作，获取响应
        mutation_response = self.datastore_client.mutate(insert=[entity])
        
        # 从响应中提取自动生成的ID
        # insert_auto_id_keys里返回的是插入后带ID的key对象
        generated_key = mutation_response.insert_auto_id_keys[0]
        auto_id = generated_key.path[-1].id  # ID是整数类型
        
        # 可以返回包含ID的结果，供后续管道步骤使用
        yield {"auto_generated_id": auto_id, "original_data": element}

之后你就可以在Beam管道中使用这个ParDo来写入实体，并拿到自动生成的ID啦。

关键点提醒 #

• Datastore自动生成的ID是整数类型，和你之前用UUID生成的字符串ID类型不同，后续处理要注意类型适配。
• 确保你的Datastore客户端有正确的权限（比如在GCP环境下使用服务账号，本地测试时配置好环境变量）。

内容的提问来源于stack exchange，提问作者user3433050