Skip to content

Commit 609b313

Browse files
kyleconroykyleconroy
committed
WIP: Python code generation via WASM
1 parent 64c2e4d commit 609b313

File tree

13 files changed

+4706
-5
lines changed

13 files changed

+4706
-5
lines changed

‎Makefile‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,5 +34,18 @@ mysqlsh:
3434
# $ protoc --version
3535
# libprotoc 3.17.3
3636
# go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
37-
proto:
38-
protoc -I ./protos --go_out=. --go_opt=module=github.com/kyleconroy/sqlc ./protos/**/*.proto
37+
proto: internal/plugin/codegen.pb.go internal/python/ast/ast.pb.go
38+
39+
internal/plugin/codegen.pb.go: protos/plugin/codegen.proto
40+
protoc -I ./protos \
41+
--go_out=. \
42+
--go_opt=module=github.com/kyleconroy/sqlc \
43+
--go-vtproto_out=. \
44+
--go-vtproto_opt=module=github.com/kyleconroy/sqlc,features=marshal+unmarshal+size \
45+
./protos/plugin/codegen.proto
46+
47+
internal/python/ast/ast.pb.go: protos/python/ast.proto
48+
protoc -I ./protos \
49+
--go_out=. \
50+
--go_opt=module=github.com/kyleconroy/sqlc \
51+
./protos/python/ast.proto

‎cmd/sqlc-codegen-python/main.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package main
2+
3+
import (
4+
"bufio"
5+
"io"
6+
"os"
7+
8+
"github.com/kyleconroy/sqlc/internal/codegen/python"
9+
"github.com/kyleconroy/sqlc/internal/plugin"
10+
)
11+
12+
func main() {
13+
var req plugin.CodeGenRequest
14+
reqBlob, err := io.ReadAll(os.Stdin)
15+
if err != nil {
16+
os.Exit(10)
17+
}
18+
if err := req.UnmarshalVT(reqBlob); err != nil {
19+
os.Exit(11)
20+
}
21+
resp, err := python.Generate(&req)
22+
if err != nil {
23+
os.Exit(12)
24+
}
25+
respBlob, err := resp.MarshalVT()
26+
if err != nil {
27+
os.Exit(13)
28+
}
29+
w := bufio.NewWriter(os.Stdout)
30+
if _, err := w.Write(respBlob); err != nil {
31+
os.Exit(14)
32+
}
33+
if err := w.Flush(); err != nil {
34+
os.Exit(15)
35+
}
36+
}

‎cmd/sqlc-codegen-python/sqlc-codegen-python‎

5.69 MB
Binary file not shown.

‎cmd/sqlc-codegen-python/sqlc-codegen-python.wasm‎

1.23 MB
Binary file not shown.

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ require (
1919
)
2020

2121
require (
22+
github.com/bytecodealliance/wasmtime-go v0.33.0 // indirect
2223
github.com/golang/protobuf v1.5.2 // indirect
2324
github.com/inconshreveable/mousetrap v1.0.0 // indirect
2425
github.com/jackc/chunkreader/v2 v2.0.1 // indirect

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
7070
github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
7171
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
7272
github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
73+
github.com/bytecodealliance/wasmtime-go v0.33.0 h1:lRhyaHpBKx+5swZKHc5uTxCPWtDRAAYope5sQPfi7Tw=
74+
github.com/bytecodealliance/wasmtime-go v0.33.0/go.mod h1:q320gUxqyI8yB+ZqRuaJOEnGkAnHh6WtJjMaT2CW4wI=
7375
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
7476
github.com/census-instrumentation/opencensus-proto v0.3.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
7577
github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=

‎internal/cmd/generate.go‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"errors"
77
"fmt"
88
"io"
9+
"log"
910
"os"
1011
"path/filepath"
1112
"runtime/trace"
@@ -199,7 +200,20 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
199200
files, err = kotlin.Generate(result, combo)
200201
case sql.Gen.Python != nil:
201202
out = combo.Python.Out
202-
resp, err = python.Generate(codeGenRequest(result, combo))
203+
req := codeGenRequest(result, combo)
204+
if os.Getenv("WASM") != "" {
205+
resp, err = pythonGenerate(req)
206+
} else {
207+
blob, err := req.MarshalVT()
208+
if err != nil {
209+
log.Fatal(err)
210+
}
211+
var codeReq plugin.CodeGenRequest
212+
if err := codeReq.UnmarshalVT(blob); err != nil {
213+
log.Fatal(err)
214+
}
215+
resp, err = python.Generate(&codeReq)
216+
}
203217
default:
204218
panic("missing language backend")
205219
}

‎internal/cmd/sqlc-codegen-python.wasm‎

1.23 MB
Binary file not shown.

‎internal/cmd/wasm.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
package cmd
2+
3+
import (
4+
"context"
5+
_ "embed"
6+
"fmt"
7+
"io/ioutil"
8+
"os"
9+
"path/filepath"
10+
"runtime/trace"
11+
12+
wasmtime "github.com/bytecodealliance/wasmtime-go"
13+
14+
"github.com/kyleconroy/sqlc/internal/plugin"
15+
)
16+
17+
//go:embed sqlc-codegen-python.wasm
18+
var pythonCodeGen []byte
19+
20+
func pythonGenerate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
21+
stdinBlob, err := req.MarshalVT()
22+
if err != nil {
23+
return nil, err
24+
}
25+
26+
cctx := context.Background()
27+
28+
ctx, task := trace.NewTask(cctx, "pythonGenerate")
29+
defer task.End()
30+
31+
dir, err := ioutil.TempDir("", "out")
32+
if err != nil {
33+
return nil, fmt.Errorf("temp dir: %w", err)
34+
}
35+
36+
defer os.RemoveAll(dir)
37+
stdinPath := filepath.Join(dir, "stdin")
38+
stderrPath := filepath.Join(dir, "stderr")
39+
stdoutPath := filepath.Join(dir, "stdout")
40+
41+
if err := os.WriteFile(stdinPath, stdinBlob, 0755); err != nil {
42+
return nil, fmt.Errorf("write file: %w", err)
43+
}
44+
45+
engine := wasmtime.NewEngine()
46+
linker := wasmtime.NewLinker(engine)
47+
48+
// Link WASI
49+
if err := linker.DefineWasi(); err != nil {
50+
return nil, fmt.Errorf("define wasi: %w", err)
51+
}
52+
53+
// Configure WASI imports to write stdout into a file.
54+
wasiConfig := wasmtime.NewWasiConfig()
55+
wasiConfig.SetStdinFile(stdinPath)
56+
wasiConfig.SetStdoutFile(stdoutPath)
57+
wasiConfig.SetStderrFile(stderrPath)
58+
59+
store := wasmtime.NewStore(engine)
60+
store.SetWasi(wasiConfig)
61+
62+
// Set the version to the same as in the WAT.
63+
// wasi, err := wasmtime.NewWasiInstance(store, wasiConfig, "wasi_snapshot_preview1")
64+
// if err != nil {
65+
// return fmt.Errorf("new wasi instances: %w", err)
66+
// }
67+
68+
// Create our module
69+
//
70+
// Compiling modules requires WebAssembly binary input, but the wasmtime
71+
// package also supports converting the WebAssembly text format to the
72+
// binary format.
73+
// wasm, err := os.ReadFile("plugin.wasm")
74+
// if err != nil {
75+
// return fmt.Errorf("read file: %w", err)
76+
// }
77+
78+
moduRegion := trace.StartRegion(ctx, "wasmtime.NewModule")
79+
module, err := wasmtime.NewModule(store.Engine, pythonCodeGen)
80+
moduRegion.End()
81+
if err != nil {
82+
return nil, fmt.Errorf("define wasi: %w", err)
83+
}
84+
85+
linkRegion := trace.StartRegion(ctx, "linker.Instantiate")
86+
instance, err := linker.Instantiate(store, module)
87+
linkRegion.End()
88+
if err != nil {
89+
return nil, fmt.Errorf("define wasi: %w", err)
90+
}
91+
92+
// Run the function
93+
94+
callRegion := trace.StartRegion(ctx, "call _start")
95+
nom := instance.GetExport(store, "_start").Func()
96+
_, err = nom.Call(store)
97+
callRegion.End()
98+
if err != nil {
99+
return nil, fmt.Errorf("call: %w", err)
100+
}
101+
102+
// Print WASM stdout
103+
stdoutBlob, err := os.ReadFile(stdoutPath)
104+
if err != nil {
105+
return nil, fmt.Errorf("read file: %w", err)
106+
}
107+
var resp plugin.CodeGenResponse
108+
return &resp, resp.UnmarshalVT(stdoutBlob)
109+
}

‎internal/codegen/python/gen.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1090,6 +1090,15 @@ func HashComment(s string) string {
10901090
}
10911091

10921092
func Generate(req *plugin.CodeGenRequest) (*plugin.CodeGenResponse, error) {
1093+
// empty := plugin.CodeGenResponse{
1094+
// Files: []*plugin.File{
1095+
// {
1096+
// Name: "marker-one.py",
1097+
// Contents: []byte(fmt.Sprintf("%t", req.Catalog == nil)),
1098+
// },
1099+
// },
1100+
// }
1101+
// return &empty, nil
10931102
enums := buildEnums(req)
10941103
models := buildModels(req)
10951104
queries, err := buildQueries(req, models)

0 commit comments

Comments
 (0)