feat: Codegen plugins, powered by WASM

cmd/sqlc-gen-json/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -26,7 +27,7 @@ func run() error {
 	if err := req.UnmarshalVT(reqBlob); err != nil {
 		return err
 	}
-	resp, err := json.Generate(&req)
+	resp, err := json.Generate(context.Background(), &req)
 	if err != nil {
 		return err
 	}

go.mod

@@ -20,6 +20,7 @@ require (
 )
 
 require (
+	github.com/bytecodealliance/wasmtime-go v0.37.0 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/jackc/chunkreader/v2 v2.0.1 // indirect

go.sum

@@ -5,6 +5,8 @@ github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20220209173558-ad29539cd2e9 h1:z
 github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20220209173558-ad29539cd2e9/go.mod h1:F7bn7fEU90QkQ3tnmaTx3LTKLEDqnwWODIYppRQ5hnY=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
+github.com/bytecodealliance/wasmtime-go v0.37.0 h1:eNP2Snp5UFMuGuunRPxwVETJ/WpC8LhWonZAklXJfjk=
+github.com/bytecodealliance/wasmtime-go v0.37.0/go.mod h1:q320gUxqyI8yB+ZqRuaJOEnGkAnHh6WtJjMaT2CW4wI=
 github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I=
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=

internal/cmd/generate.go

@@ -20,8 +20,10 @@ import (
 	"github.com/kyleconroy/sqlc/internal/debug"
 	"github.com/kyleconroy/sqlc/internal/ext"
 	"github.com/kyleconroy/sqlc/internal/ext/process"
+	"github.com/kyleconroy/sqlc/internal/ext/wasm"
 	"github.com/kyleconroy/sqlc/internal/multierr"
 	"github.com/kyleconroy/sqlc/internal/opts"
+	"github.com/kyleconroy/sqlc/internal/plugin"
 )
 
 const errMessageNoVersion = `The configuration file must have a version number.
@@ -51,6 +53,15 @@ type outPair struct {
 	config.SQL
 }
 
+func findPlugin(conf config.Config, name string) (*config.Plugin, error) {
+	for _, plug := range conf.Plugins {
+		if plug.Name == name {
+			return &plug, nil
+		}
+	}
+	return nil, fmt.Errorf("plugin not found")
+}
+
 func readConfig(stderr io.Writer, dir, filename string) (string, *config.Config, error) {
 	configPath := ""
 	if filename != "" {
@@ -216,43 +227,7 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
 			break
 		}
 
-		var region *trace.Region
-		if debug.Traced {
-			region = trace.StartRegion(ctx, "codegen")
-		}
-		var handler ext.Handler
-		var out string
-		switch {
-		case sql.Gen.Go != nil:
-			out = combo.Go.Out
-			handler = ext.HandleFunc(golang.Generate)
-
-		case sql.Gen.Kotlin != nil:
-			out = combo.Kotlin.Out
-			handler = ext.HandleFunc(kotlin.Generate)
-
-		case sql.Gen.Python != nil:
-			out = combo.Python.Out
-			handler = ext.HandleFunc(python.Generate)
-
-		case sql.Gen.JSON != nil:
-			out = combo.JSON.Out
-			handler = ext.HandleFunc(json.Generate)
-
-		case sql.Plugin != nil:
-			out = sql.Plugin.Out
-			handler = &process.Runner{
-				Config: combo.Global,
-				Plugin: sql.Plugin.Plugin,
-			}
-
-		default:
-			panic("missing language backend")
-		}
-		resp, err := handler.Generate(codeGenRequest(result, combo))
-		if region != nil {
-			region.End()
-		}
+		out, resp, err := codegen(ctx, combo, sql, result)
 		if err != nil {
 			fmt.Fprintf(stderr, "# package %s\n", name)
 			fmt.Fprintf(stderr, "error generating code: %s\n", err)
@@ -262,6 +237,7 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
 			}
 			continue
 		}
+
 		files := map[string]string{}
 		for _, file := range resp.Files {
 			files[file.Name] = string(file.Contents)
@@ -313,3 +289,58 @@ func parse(ctx context.Context, e Env, name, dir string, sql config.SQL, combo c
 	}
 	return c.Result(), false
 }
+
+func codegen(ctx context.Context, combo config.CombinedSettings, sql outPair, result *compiler.Result) (string, *plugin.CodeGenResponse, error) {
+	var region *trace.Region
+	if debug.Traced {
+		region = trace.StartRegion(ctx, "codegen")
+	}
+	var handler ext.Handler
+	var out string
+	switch {
+	case sql.Gen.Go != nil:
+		out = combo.Go.Out
+		handler = ext.HandleFunc(golang.Generate)
+
+	case sql.Gen.Kotlin != nil:
+		out = combo.Kotlin.Out
+		handler = ext.HandleFunc(kotlin.Generate)
+
+	case sql.Gen.Python != nil:
+		out = combo.Python.Out
+		handler = ext.HandleFunc(python.Generate)
+
+	case sql.Gen.JSON != nil:
+		out = combo.JSON.Out
+		handler = ext.HandleFunc(json.Generate)
+
+	case sql.Plugin != nil:
+		out = sql.Plugin.Out
+		plug, err := findPlugin(combo.Global, sql.Plugin.Plugin)
+		if err != nil {
+			return "", nil, fmt.Errorf("plugin not found: %s", err)
+		}
+
+		switch {
+		case plug.Process != nil:
+			handler = &process.Runner{
+				Cmd: plug.Process.Cmd,
+			}
+		case plug.WASM != nil:
+			handler = &wasm.Runner{
+				URL:      plug.WASM.URL,
+				Checksum: plug.WASM.Checksum,
+			}
+		default:
+			return "", nil, fmt.Errorf("unsupported plugin type")
+		}
+
+	default:
+		return "", nil, fmt.Errorf("missing language backend")
+	}
+	resp, err := handler.Generate(ctx, codeGenRequest(result, combo))
+	if region != nil {
+		region.End()
+	}
+	return out, resp, err
+}

internal/codegen/golang/gen.go

@@ -3,6 +3,7 @@ package golang
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"go/format"
@@ -42,7 +43,7 @@ func (t *tmplCtx) OutputQuery(sourceName string) bool {
 	return t.SourceName == sourceName
 }
 
-func Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+func Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
 	enums := buildEnums(req)
 	structs := buildStructs(req)
 	queries, err := buildQueries(req, structs)

internal/codegen/json/gen.go

@@ -2,6 +2,7 @@ package json
 
 import (
 	"bytes"
+	"context"
 	ejson "encoding/json"
 	"fmt"
 
@@ -31,7 +32,7 @@ func parseOptions(req *plugin.CodeGenRequest) (plugin.JSONCode, error) {
 	return plugin.JSONCode{}, nil
 }
 
-func Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+func Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
 	options, err := parseOptions(req)
 	if err != nil {
 		return nil, err

internal/codegen/kotlin/gen.go

@@ -3,6 +3,7 @@ package kotlin
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"regexp"
@@ -759,7 +760,7 @@ func ktFormat(s string) string {
 	return o
 }
 
-func Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+func Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
 	enums := buildEnums(req)
 	structs := buildDataClasses(req)
 	queries, err := buildQueries(req, structs)

internal/codegen/python/gen.go

@@ -1,6 +1,7 @@
 package python
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log"
@@ -1080,7 +1081,7 @@ func HashComment(s string) string {
 	return "# " + strings.ReplaceAll(s, "\n", "\n# ")
 }
 
-func Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+func Generate(_ context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
 	enums := buildEnums(req)
 	models := buildModels(req)
 	queries, err := buildQueries(req, models)

internal/config/config.go

@@ -90,7 +90,8 @@ type Plugin struct {
 		Cmd string `json:"cmd" yaml:"cmd"`
 	} `json:"process" yaml:"process"`
 	WASM *struct {
-		URL string `json:"url" yaml:"url"`
+		URL      string `json:"url" yaml:"url"`
+		Checksum string `json:"checksum" yaml:"checksum"`
 	} `json:"wasm" yaml:"wasm"`
 }
 

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

@@ -0,0 +1 @@
+Hello World

internal/endtoend/testdata/wasm_plugin_sqlc_gen_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/wasm_plugin_sqlc_gen_json/schema.sql

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

internal/endtoend/testdata/wasm_plugin_sqlc_gen_json/sqlc.json

@@ -0,0 +1,29 @@
+{
+  "version": "2",
+  "sql": [
+    {
+      "schema": "schema.sql",
+      "queries": "query.sql",
+      "engine": "postgresql",
+      "codegen": [
+        {
+          "out": "gen",
+          "plugin": "jsonb",
+          "options": {
+            "indent": "  ",
+            "filename": "codegen.json"
+          }
+        }
+      ]
+    }
+  ],
+  "plugins": [
+    {
+      "name": "jsonb",
+      "wasm": {
+        "url": "file:///Users/kyle/Downloads/sqlc-gen-json-wasm.wasm",
+        "checksum": "sha256/afc486dac2068d741d7a4110146559d12a013fd0286f42a2fc7dcd802424ad07"
+      }
+    }
+  ]
+}

internal/ext/handler.go

@@ -1,21 +1,23 @@
 package ext
 
 import (
+	"context"
+
 	"github.com/kyleconroy/sqlc/internal/plugin"
 )
 
 type Handler interface {
-	Generate(*plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)
+	Generate(context.Context, *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)
 }
 
 type wrapper struct {
-	fn func(*plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)
+	fn func(context.Context, *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)
 }
 
-func (w *wrapper) Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
-	return w.fn(req)
+func (w *wrapper) Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+	return w.fn(ctx, req)
 }
 
-func HandleFunc(fn func(*plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)) Handler {
+func HandleFunc(fn func(context.Context, *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error)) Handler {
 	return &wrapper{fn}
 }

internal/ext/process/gen.go

@@ -9,47 +9,26 @@ import (
 
 	"google.golang.org/protobuf/proto"
 
-	"github.com/kyleconroy/sqlc/internal/config"
 	"github.com/kyleconroy/sqlc/internal/plugin"
 )
 
 type Runner struct {
-	Config config.Config
-	Plugin string
-}
-
-func (r Runner) pluginCmd() (string, error) {
-	for _, plug := range r.Config.Plugins {
-		if plug.Name != r.Plugin {
-			continue
-		}
-		if plug.Process == nil {
-			continue
-		}
-		return plug.Process.Cmd, nil
-	}
-	return "", fmt.Errorf("plugin not found")
+	Cmd string
 }
 
 // TODO: Update the gen func signature to take a ctx
-func (r Runner) Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
+func (r Runner) Generate(ctx context.Context, req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
 	stdin, err := proto.Marshal(req)
 	if err != nil {
 		return nil, fmt.Errorf("failed to encode codegen request: %s", err)
 	}
 
-	name, err := r.pluginCmd()
-	if err != nil {
-		return nil, fmt.Errorf("process: unknown plugin %s", r.Plugin)
-	}
-
 	// Check if the output plugin exists
-	path, err := exec.LookPath(name)
+	path, err := exec.LookPath(r.Cmd)
 	if err != nil {
-		return nil, fmt.Errorf("process: %s not found", name)
+		return nil, fmt.Errorf("process: %s not found", r.Cmd)
 	}
 
-	ctx := context.Background()
 	cmd := exec.CommandContext(ctx, path)
 	cmd.Stdin = bytes.NewReader(stdin)
 	cmd.Env = []string{