protoc-gen-go-aip-dataloader

protoc-gen-go-aip-dataloader

Generate GraphQL Go dataloaders for AIP BatchGet methods.

Heavily inspired by vektah/dataloaden.

How to

Step 1: Declare a Batch method in your API

  // Batch get sites.
  // See: https://google.aip.dev/231 (Batch methods: Get).
  rpc BatchGetSites(BatchGetSitesRequest) returns (BatchGetSitesResponse) {
    option (google.api.http) = {
      get: "/v1/{parent=shippers/*}/sites:batchGet"
    };
  }

Step 2: Install the plugin

$ GOBIN=$PWD/build go install go.einride.tech/protoc-gen-go-aip-dataloader

Step 3: Configure code generation

The following example uses a buf generate template to configure the CLI generator.

buf.gen.example.yaml:

version: v1

managed:
  enabled: true
  go_package_prefix:
    default: go.einride.tech/protoc-gen-go-aip-dataloader/example/internal/proto/gen
    except:
      - buf.build/googleapis/googleapis

plugins:
  # Dataloaders require the stubs generated by protoc-gen-go.
  - name: go
    out: example/internal/proto/gen
    opt: module=go.einride.tech/protoc-gen-go-aip-dataloader/example/internal/proto/gen
    path: ./build/protoc-gen-go

  # Dataloaders require the stubs generated by protoc-gen-go-grpc.
  - name: go-grpc
    out: example/internal/proto/gen
    opt: module=go.einride.tech/protoc-gen-go-aip-dataloader/example/internal/proto/gen
    path: ./build/protoc-gen-go-grpc

  # Generate dataloaders for Batch methods.
  - name: go-aip-dataloader
    out: example/internal/proto/gen
    opt: module=go.einride.tech/protoc-gen-go-aip-dataloader/example/internal/proto/gen
    path: ./build/protoc-gen-go-aip-dataloader

Step 4: Generate the code

$ buf generate buf.build/einride/aip \
  --template buf.gen.example.yaml \
  --path einride/example/freight

Step 5: Use the dataloader

package main

import (
	"context"
	"time"

	freightv1 "go.einride.tech/protoc-gen-go-aip-dataloader/example/internal/proto/gen/einride/example/freight/v1"
	"google.golang.org/grpc"
)

func main() {
	ctx := context.Background()
	// Connect to service.
	conn, err := grpc.Dial("ADDRESS")
	if err != nil {
		panic(err) // TODO: Handle error.
	}
	// Create a client.
	client := freightv1.NewFreightServiceClient(conn)
	// Wrap the client in a dataloader.
	sitesDataloader := freightv1.NewSitesDataloader(
		ctx,
		client,
		&freightv1.BatchGetSitesRequest{},
		100*time.Millisecond, // wait
		1_000,                // max batch
	)
	// The dataloader can batch together Get calls into BatchGet calls and return thunks.
	// This is useful in e.g. GraphQL resolver implementations.
	thunk1 := sitesDataloader.LoadThunk("shippers/folkfood/sites/sthlm")
	thunk2 := sitesDataloader.LoadThunk("shippers/folkfood/sites/gbg")
	// Evaluate the thunks to wait for the BatchGet call to go through.
	sthlm, err := thunk1()
	if err != nil {
		panic(err)
	}
	gbg, err := thunk2()
	if err != nil {
		panic(err)
	}
	_, _ = sthlm, gbg
}