扩展 gRPC 服务

在开发 gRPC 服务器时，经常需要包含一些无法通过 Ent 模式自动生成的方法。为了实现此目的，可以在 entpb 目录下的额外 .proto 文件中定义这些方法，并将其置于额外的服务中。

信息

本节描述的更改可在 此拉取请求 中查看。

例如，假设需要添加名为 TopUser 的方法来返回 ID 号最大的用户。为此，在 entpb 目录中创建新的 .proto 文件，并定义新服务：

ent/proto/entpb/ext.proto

syntax = "proto3";

package entpb;

option go_package = "github.com/rotemtam/ent-grpc-example/ent/proto/entpb";

import "entpb/entpb.proto";

import "google/protobuf/empty.proto";


service ExtService {
  rpc TopUser ( google.protobuf.Empty ) returns ( User );
}

接着更新 entpb/generate.go，将新文件加入 protoc 命令的输入中：

ent/proto/entpb/generate.go

- //go:generate protoc -I=.. --go_out=.. --go-grpc_out=.. --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative --entgrpc_out=.. --entgrpc_opt=paths=source_relative,schema_path=../../schema entpb/entpb.proto
+ //go:generate protoc -I=.. --go_out=.. --go-grpc_out=.. --go_opt=paths=source_relative --go-grpc_opt=paths=source_relative --entgrpc_out=.. --entgrpc_opt=paths=source_relative,schema_path=../../schema entpb/entpb.proto entpb/ext.proto

随后重新运行代码生成：

go generate ./...

可观察到 ent/proto/entpb 目录中生成了新文件：

tree
.
|-- entpb.pb.go
|-- entpb.proto
|-- entpb_grpc.pb.go
|-- entpb_user_service.go
|-- ext.pb.go
|-- ext.proto
|-- ext_grpc.pb.go
`-- generate.go

0 directories, 9 files

现在，可以在 ent/proto/entpb/ext.go 中实现 TopUser 方法：

ent/proto/entpb/ext.go

package entpb

import (
    "context"

    "github.com/rotemtam/ent-grpc-example/ent"
    "github.com/rotemtam/ent-grpc-example/ent/user"
    "google.golang.org/protobuf/types/known/emptypb"
)

// ExtService 实现 ExtServiceServer。
type ExtService struct {
    client *ent.Client
    UnimplementedExtServiceServer
}

// TopUser 返回 ID 最大的用户。
func (s *ExtService) TopUser(ctx context.Context, _ *emptypb.Empty) (*User, error) {
    id := s.client.User.Query().Aggregate(ent.Max(user.FieldID)).IntX(ctx)
    user := s.client.User.GetX(ctx, id)
    return toProtoUser(user)
}

// NewExtService 返回新的 ExtService。
func NewExtService(client *ent.Client) *ExtService {
    return &ExtService{
        client: client,
    }
}

将新服务添加到 gRPC 服务器

最后更新 cmd/server.go 以包含新服务：

cmd/server.go

package main

import (
    "context"
    "log"
    "net"

    _ "github.com/mattn/go-sqlite3"
    "github.com/rotemtam/ent-grpc-example/ent"
    "github.com/rotemtam/ent-grpc-example/ent/proto/entpb"
    "google.golang.org/grpc"
)

func main() {
    // 初始化 ent 客户端。
    client, err := ent.Open("sqlite3", "file:ent?mode=memory&cache=shared&_fk=1")
    if err != nil {
        log.Fatalf("failed opening connection to sqlite: %v", err)
    }
    defer client.Close()

    // 运行迁移工具（创建表等）。
    if err := client.Schema.Create(context.Background()); err != nil {
        log.Fatalf("failed creating schema resources: %v", err)
    }

    // 初始化生成的 User 服务。
    svc := entpb.NewUserService(client)

    // 创建新的 gRPC 服务器（可将多个服务连接至单个服务器）。
    server := grpc.NewServer()

    // 向服务器注册 User 服务。
    entpb.RegisterUserServiceServer(server, svc)

    // 向服务器注册外部 ExtService 服务。
    entpb.RegisterExtServiceServer(server, entpb.NewExtService(client))

    // 开放 5000 端口监听流量。
    lis, err := net.Listen("tcp", ":5000")
    if err != nil {
        log.Fatalf("failed listening: %s", err)
    }

    // 持续监听流量。
    if err := server.Serve(lis); err != nil {
        log.Fatalf("server ended: %s", err)
    }
}