Generate Python code via a WASM plugin

Makefile

@@ -21,6 +21,9 @@ sqlc-pg-gen:
 start:
 	docker-compose up -d
 
+fmt:
+	go fmt ./...
+
 psql:
 	PGPASSWORD=mysecretpassword psql --host=127.0.0.1 --port=5432 --username=postgres dinotest
 
@@ -31,5 +34,18 @@ mysqlsh:
 # $ protoc --version
 # libprotoc 3.17.3
 # go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
-proto:
-	protoc -I ./protos --go_out=. --go_opt=module=github.com/kyleconroy/sqlc ./protos/**/*.proto
+proto: internal/plugin/codegen.pb.go internal/python/ast/ast.pb.go
+
+internal/plugin/codegen.pb.go: protos/plugin/codegen.proto
+	protoc -I ./protos \
+		--go_out=. \
+		--go_opt=module=github.com/kyleconroy/sqlc \
+		--go-vtproto_out=. \
+		--go-vtproto_opt=module=github.com/kyleconroy/sqlc,features=marshal+unmarshal+size \
+		./protos/plugin/codegen.proto
+
+internal/python/ast/ast.pb.go: protos/python/ast.proto
+	protoc -I ./protos \
+		--go_out=. \
+		--go_opt=module=github.com/kyleconroy/sqlc \
+		./protos/python/ast.proto

cmd/sqlc-codegen-python/main.go

@@ -0,0 +1,36 @@
+package main
+
+import (
+	"bufio"
+	"io"
+	"os"
+
+	"github.com/kyleconroy/sqlc/internal/codegen/python"
+	"github.com/kyleconroy/sqlc/internal/plugin"
+)
+
+func main() {
+	var req plugin.CodeGenRequest
+	reqBlob, err := io.ReadAll(os.Stdin)
+	if err != nil {
+		os.Exit(10)
+	}
+	if err := req.UnmarshalVT(reqBlob); err != nil {
+		os.Exit(11)
+	}
+	resp, err := python.Generate(&req)
+	if err != nil {
+		os.Exit(12)
+	}
+	respBlob, err := resp.MarshalVT()
+	if err != nil {
+		os.Exit(13)
+	}
+	w := bufio.NewWriter(os.Stdout)
+	if _, err := w.Write(respBlob); err != nil {
+		os.Exit(14)
+	}
+	if err := w.Flush(); err != nil {
+		os.Exit(15)
+	}
+}

examples/authors/mysql/db_test.go

@@ -1,3 +1,4 @@
+//go:build examples
 // +build examples
 
 package authors

examples/authors/postgresql/db_test.go

@@ -1,3 +1,4 @@
+//go:build examples
 // +build examples
 
 package authors

examples/booktest/postgresql/db_test.go

@@ -1,3 +1,4 @@
+//go:build examples
 // +build examples
 
 package booktest

examples/ondeck/mysql/db_test.go

@@ -1,3 +1,4 @@
+//go:build examples
 // +build examples
 
 package ondeck

examples/ondeck/postgresql/db_test.go

@@ -1,3 +1,4 @@
+//go:build examples
 // +build examples
 
 package ondeck

go.mod

@@ -19,6 +19,7 @@ require (
 )
 
 require (
+	github.com/bytecodealliance/wasmtime-go v0.33.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

@@ -70,6 +70,8 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
+github.com/bytecodealliance/wasmtime-go v0.33.0 h1:lRhyaHpBKx+5swZKHc5uTxCPWtDRAAYope5sQPfi7Tw=
+github.com/bytecodealliance/wasmtime-go v0.33.0/go.mod h1:q320gUxqyI8yB+ZqRuaJOEnGkAnHh6WtJjMaT2CW4wI=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/census-instrumentation/opencensus-proto v0.3.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=

internal/cmd/generate.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"os"
 	"path/filepath"
 	"runtime/trace"
@@ -19,6 +20,7 @@ import (
 	"github.com/kyleconroy/sqlc/internal/debug"
 	"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.
@@ -187,6 +189,7 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
 			region = trace.StartRegion(ctx, "codegen")
 		}
 		var files map[string]string
+		var resp *plugin.CodeGenResponse
 		var out string
 		switch {
 		case sql.Gen.Go != nil:
@@ -197,14 +200,34 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
 			files, err = kotlin.Generate(result, combo)
 		case sql.Gen.Python != nil:
 			out = combo.Python.Out
-			files, err = python.Generate(result, combo)
+			req := codeGenRequest(result, combo)
+			if os.Getenv("WASM") != "" {
+				resp, err = pythonGenerate(ctx, req)
+			} else {
+				blob, err := req.MarshalVT()
+				if err != nil {
+					log.Fatal(err)
+				}
+				var codeReq plugin.CodeGenRequest
+				if err := codeReq.UnmarshalVT(blob); err != nil {
+					log.Fatal(err)
+				}
+				resp, err = python.Generate(&codeReq)
+			}
 		default:
 			panic("missing language backend")
 		}
 		if region != nil {
 			region.End()
 		}
 
+		if resp != nil {
+			files = map[string]string{}
+			for _, file := range resp.Files {
+				files[file.Name] = string(file.Contents)
+			}
+		}
+
 		if err != nil {
 			fmt.Fprintf(stderr, "# package %s\n", name)
 			fmt.Fprintf(stderr, "error generating code: %s\n", err)

internal/cmd/shim.go

@@ -0,0 +1,217 @@
+package cmd
+
+import (
+	"strings"
+
+	"github.com/kyleconroy/sqlc/internal/compiler"
+	"github.com/kyleconroy/sqlc/internal/config"
+	"github.com/kyleconroy/sqlc/internal/plugin"
+	"github.com/kyleconroy/sqlc/internal/sql/catalog"
+)
+
+func pluginOverride(o config.Override) *plugin.Override {
+	var column string
+	var table plugin.Identifier
+
+	if o.Column != "" {
+		colParts := strings.Split(o.Column, ".")
+		switch len(colParts) {
+		case 2:
+			table.Schema = "public"
+			table.Name = colParts[0]
+			column = colParts[1]
+		case 3:
+			table.Schema = colParts[0]
+			table.Name = colParts[1]
+			column = colParts[2]
+		case 4:
+			table.Catalog = colParts[0]
+			table.Schema = colParts[1]
+			table.Name = colParts[2]
+			column = colParts[3]
+		}
+	}
+	return &plugin.Override{
+		CodeType:   "", // FIXME
+		DbType:     o.DBType,
+		Nullable:   o.Nullable,
+		Column:     o.Column,
+		ColumnName: column,
+		Table:      &table,
+		PythonType: pluginPythonType(o.PythonType),
+	}
+}
+
+func pluginSettings(cs config.CombinedSettings) *plugin.Settings {
+	var over []*plugin.Override
+	for _, o := range cs.Overrides {
+		over = append(over, pluginOverride(o))
+	}
+	return &plugin.Settings{
+		Version:   cs.Global.Version,
+		Engine:    string(cs.Package.Engine),
+		Schema:    []string(cs.Package.Schema),
+		Queries:   []string(cs.Package.Queries),
+		Overrides: over,
+		Rename:    cs.Rename,
+		Python:    pluginPythonCode(cs.Python),
+	}
+}
+
+func pluginPythonCode(s config.SQLPython) *plugin.PythonCode {
+	return &plugin.PythonCode{
+		Out:                 s.Out,
+		Package:             s.Package,
+		EmitExactTableNames: s.EmitExactTableNames,
+		EmitSyncQuerier:     s.EmitSyncQuerier,
+		EmitAsyncQuerier:    s.EmitAsyncQuerier,
+	}
+}
+
+func pluginPythonType(pt config.PythonType) *plugin.PythonType {
+	return &plugin.PythonType{
+		Module: pt.Module,
+		Name:   pt.Name,
+	}
+}
+
+func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
+	var schemas []*plugin.Schema
+	for _, s := range c.Schemas {
+		var enums []*plugin.Enum
+		for _, typ := range s.Types {
+			enum, ok := typ.(*catalog.Enum)
+			if !ok {
+				continue
+			}
+			enums = append(enums, &plugin.Enum{
+				Name:    enum.Name,
+				Comment: enum.Comment,
+				Vals:    enum.Vals,
+			})
+		}
+		var tables []*plugin.Table
+		for _, t := range s.Tables {
+			var columns []*plugin.Column
+			for _, c := range t.Columns {
+				l := -1
+				if c.Length != nil {
+					l = *c.Length
+				}
+				columns = append(columns, &plugin.Column{
+					Name: c.Name,
+					Type: &plugin.Identifier{
+						Catalog: c.Type.Catalog,
+						Schema:  c.Type.Schema,
+						Name:    c.Type.Name,
+					},
+					Comment: c.Comment,
+					NotNull: c.IsNotNull,
+					IsArray: c.IsArray,
+					Length:  int32(l),
+					Table: &plugin.Identifier{
+						Catalog: t.Rel.Catalog,
+						Schema:  t.Rel.Schema,
+						Name:    t.Rel.Name,
+					},
+				})
+			}
+			tables = append(tables, &plugin.Table{
+				Rel: &plugin.Identifier{
+					Catalog: t.Rel.Catalog,
+					Schema:  t.Rel.Schema,
+					Name:    t.Rel.Name,
+				},
+				Columns: columns,
+				Comment: t.Comment,
+			})
+		}
+		schemas = append(schemas, &plugin.Schema{
+			Comment: s.Comment,
+			Name:    s.Name,
+			Tables:  tables,
+			Enums:   enums,
+		})
+	}
+	return &plugin.Catalog{
+		Name:          c.Name,
+		DefaultSchema: c.DefaultSchema,
+		Comment:       c.Comment,
+		Schemas:       schemas,
+	}
+}
+
+func pluginQueries(r *compiler.Result) []*plugin.Query {
+	var out []*plugin.Query
+	for _, q := range r.Queries {
+		var params []*plugin.Parameter
+		var columns []*plugin.Column
+		for _, c := range q.Columns {
+			columns = append(columns, pluginQueryColumn(c))
+		}
+		for _, p := range q.Params {
+			params = append(params, pluginQueryParam(p))
+		}
+		out = append(out, &plugin.Query{
+			Name:     q.Name,
+			Cmd:      q.Cmd,
+			Text:     q.SQL,
+			Comments: q.Comments,
+			Columns:  columns,
+			Params:   params,
+			Filename: q.Filename,
+		})
+	}
+	return out
+}
+
+func pluginQueryColumn(c *compiler.Column) *plugin.Column {
+	l := -1
+	if c.Length != nil {
+		l = *c.Length
+	}
+	out := &plugin.Column{
+		Name:    c.Name,
+		Comment: c.Comment,
+		NotNull: c.NotNull,
+		IsArray: c.IsArray,
+		Length:  int32(l),
+	}
+
+	if c.Type != nil {
+		out.Type = &plugin.Identifier{
+			Catalog: c.Type.Catalog,
+			Schema:  c.Type.Schema,
+			Name:    c.Type.Name,
+		}
+	} else {
+		out.Type = &plugin.Identifier{
+			Name: c.DataType,
+		}
+	}
+
+	if c.Table != nil {
+		out.Table = &plugin.Identifier{
+			Catalog: c.Table.Catalog,
+			Schema:  c.Table.Schema,
+			Name:    c.Table.Name,
+		}
+	}
+
+	return out
+}
+
+func pluginQueryParam(p compiler.Parameter) *plugin.Parameter {
+	return &plugin.Parameter{
+		Number: int32(p.Number),
+		Column: pluginQueryColumn(p.Column),
+	}
+}
+
+func codeGenRequest(r *compiler.Result, settings config.CombinedSettings) *plugin.CodeGenRequest {
+	return &plugin.CodeGenRequest{
+		Settings: pluginSettings(settings),
+		Catalog:  pluginCatalog(r.Catalog),
+		Queries:  pluginQueries(r),
+	}
+}