Add support for json format from process plugins (#3827)

* Add support for json format from process plugins

* Update gen.go

* Fix test

---------

Co-authored-by: Kyle Gray <[email protected]>

.github/workflows/ci.yml

@@ -38,6 +38,9 @@ jobs:
     - name: install sqlc-gen-test
       run: go install github.com/sqlc-dev/[email protected]
 
+    - name: install test-json-process-plugin
+      run: go install ./scripts/test-json-process-plugin/
+
     - name: install ./...
       run: go install ./...
       env:

Makefile

@@ -32,6 +32,9 @@ sqlc-pg-gen:
 sqlc-gen-json:
 	go build -o ~/bin/sqlc-gen-json ./cmd/sqlc-gen-json
 
+test-json-process-plugin:
+	go build -o ~/bin/test-json-process-plugin ./scripts/test-json-process-plugin/
+
 start:
 	docker compose up -d
 

docs/guides/plugins.md

@@ -72,6 +72,8 @@ For a complete working example see the following files:
   - A process-based plugin that serializes the CodeGenRequest to JSON
 - [process_plugin_sqlc_gen_json](https://github.com/sqlc-dev/sqlc/tree/main/internal/endtoend/testdata/process_plugin_sqlc_gen_json)
   - An example project showing how to use a process-based plugin
+- [process_plugin_sqlc_gen_json](https://github.com/sqlc-dev/sqlc/tree/main/internal/endtoend/testdata/process_plugin_format_json/)
+  - An example project showing how to use a process-based plugin using json
 
 ## Environment variables
 
@@ -99,4 +101,4 @@ plugins:
 ```
 
 A variable named `SQLC_VERSION` is always included in the plugin's
-environment, set to the version of the `sqlc` executable invoking it.
+environment, set to the version of the `sqlc` executable invoking it.

docs/reference/config.md

@@ -273,6 +273,8 @@ Each mapping in the `plugins` collection has the following keys:
 - `process`: A mapping with a single `cmd` key
   - `cmd`:
     - The executable to call when using this plugin
+  - `format`:
+    - The format expected. Supports `json` and `protobuf` formats. Defaults to `protobuf`.
 - `wasm`: A mapping with a two keys `url` and `sha256`
   - `url`:
     - The URL to fetch the WASM file. Supports the `https://` or `file://` schemes.

internal/cmd/generate.go

@@ -349,8 +349,9 @@ func codegen(ctx context.Context, combo config.CombinedSettings, sql OutputPair,
 		switch {
 		case plug.Process != nil:
 			handler = &process.Runner{
-				Cmd: plug.Process.Cmd,
-				Env: plug.Env,
+				Cmd:    plug.Process.Cmd,
+				Env:    plug.Env,
+				Format: plug.Process.Format,
 			}
 		case plug.WASM != nil:
 			handler = &wasm.Runner{

internal/config/config.go

@@ -89,7 +89,8 @@ type Plugin struct {
 	Name    string   `json:"name" yaml:"name"`
 	Env     []string `json:"env" yaml:"env"`
 	Process *struct {
-		Cmd string `json:"cmd" yaml:"cmd"`
+		Cmd    string `json:"cmd" yaml:"cmd"`
+		Format string `json:"format" yaml:"format"`
 	} `json:"process" yaml:"process"`
 	WASM *struct {
 		URL    string `json:"url" yaml:"url"`

internal/endtoend/testdata/process_plugin_format_json/exec.json

@@ -0,0 +1,4 @@
+{
+  "process": "test-json-process-plugin",
+  "os": [ "darwin", "linux" ]
+}

internal/endtoend/testdata/process_plugin_format_json/gen/hello.txt

@@ -0,0 +1,12 @@
+SELECT id, name, bio FROM authors
+WHERE id = $1 LIMIT 1
+SELECT id, name, bio FROM authors
+ORDER BY name
+INSERT INTO authors (
+          name, bio
+) VALUES (
+  $1, $2
+)
+RETURNING id, name, bio
+DELETE FROM authors
+WHERE id = $1

internal/endtoend/testdata/process_plugin_format_json/query.sql

@@ -0,0 +1,19 @@
+-- name: GetAuthor :one
+SELECT * FROM authors
+WHERE id = $1 LIMIT 1;
+
+-- name: ListAuthors :many
+SELECT * FROM authors
+ORDER BY name;
+
+-- name: CreateAuthor :one
+INSERT INTO authors (
+          name, bio
+) VALUES (
+  $1, $2
+)
+RETURNING *;
+
+-- name: DeleteAuthor :exec
+DELETE FROM authors
+WHERE id = $1;

internal/endtoend/testdata/process_plugin_format_json/schema.sql

@@ -0,0 +1,5 @@
+CREATE TABLE authors (
+          id   BIGSERIAL PRIMARY KEY,
+          name text      NOT NULL,
+          bio  text
+);

internal/endtoend/testdata/process_plugin_format_json/sqlc.json

@@ -0,0 +1,25 @@
+{
+  "version": "2",
+  "sql": [
+    {
+      "schema": "schema.sql",
+      "queries": "query.sql",
+      "engine": "postgresql",
+      "codegen": [
+        {
+          "out": "gen",
+          "plugin": "jsonb"
+        }
+      ]
+    }
+  ],
+  "plugins": [
+    {
+      "name": "jsonb",
+      "process": {
+        "cmd": "test-json-process-plugin",
+        "format": "json"
+      }
+    }
+  ]
+}

internal/ext/process/gen.go

@@ -11,15 +11,17 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/sqlc-dev/sqlc/internal/info"
 )
 
 type Runner struct {
-	Cmd string
-	Env []string
+	Cmd    string
+	Format string
+	Env    []string
 }
 
 func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
@@ -28,9 +30,27 @@ func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any,
 		return fmt.Errorf("args isn't a protoreflect.ProtoMessage")
 	}
 
-	stdin, err := proto.Marshal(req)
-	if err != nil {
-		return fmt.Errorf("failed to encode codegen request: %w", err)
+	var stdin []byte
+	var err error
+	switch r.Format {
+	case "json":
+		m := &protojson.MarshalOptions{
+			EmitUnpopulated: true,
+			Indent:          "",
+			UseProtoNames:   true,
+		}
+		stdin, err = m.Marshal(req)
+
+		if err != nil {
+			return fmt.Errorf("failed to encode codegen request: %w", err)
+		}
+	case "", "protobuf":
+		stdin, err = proto.Marshal(req)
+		if err != nil {
+			return fmt.Errorf("failed to encode codegen request: %w", err)
+		}
+	default:
+		return fmt.Errorf("unknown plugin format: %s", r.Format)
 	}
 
 	// Check if the output plugin exists
@@ -66,8 +86,15 @@ func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any,
 		return fmt.Errorf("reply isn't a protoreflect.ProtoMessage")
 	}
 
-	if err := proto.Unmarshal(out, resp); err != nil {
-		return fmt.Errorf("process: failed to read codegen resp: %w", err)
+	switch r.Format {
+	case "json":
+		if err := protojson.Unmarshal(out, resp); err != nil {
+			return fmt.Errorf("process: failed to read codegen resp: %w", err)
+		}
+	default:
+		if err := proto.Unmarshal(out, resp); err != nil {
+			return fmt.Errorf("process: failed to read codegen resp: %w", err)
+		}
 	}
 
 	return nil

scripts/test-json-process-plugin/main.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"os"
+)
+
+type Out struct {
+	Files []File `json:"files"`
+}
+
+type File struct {
+	Name     string `json:"name"`
+	Contents []byte `json:"contents"`
+}
+
+func main() {
+	in := make(map[string]interface{})
+	decoder := json.NewDecoder(os.Stdin)
+	err := decoder.Decode(&in)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "error generating JSON: %s", err)
+		os.Exit(2)
+	}
+
+	buf := bytes.NewBuffer(nil)
+	queries := in["queries"].([]interface{})
+	for _, q := range queries {
+		text := q.(map[string]interface{})["text"].(string)
+		buf.WriteString(text)
+		buf.WriteString("\n")
+	}
+
+	e := json.NewEncoder(os.Stdout)
+	e.SetIndent("", "  ")
+	e.Encode(&Out{Files: []File{{Name: "hello.txt", Contents: buf.Bytes()}}})
+}