snapshot of https://github.com/sqlc-dev/sqlc.git (main, 2026-09-16; blobs > 131072B omitted)

internal/ext/handler.go

@@ -0,0 +1,52 @@
+package ext
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/sqlc-dev/sqlc/internal/plugin"
+)
+
+type Handler interface {
+	Generate(context.Context, *plugin.GenerateRequest) (*plugin.GenerateResponse, error)
+
+	Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error
+	NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error)
+}
+
+type wrapper struct {
+	fn func(context.Context, *plugin.GenerateRequest) (*plugin.GenerateResponse, error)
+}
+
+func (w *wrapper) Generate(ctx context.Context, req *plugin.GenerateRequest) (*plugin.GenerateResponse, error) {
+	return w.fn(ctx, req)
+}
+
+func (w *wrapper) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+	req, ok := args.(*plugin.GenerateRequest)
+	if !ok {
+		return fmt.Errorf("args isn't a GenerateRequest")
+	}
+	resp, ok := reply.(*plugin.GenerateResponse)
+	if !ok {
+		return fmt.Errorf("reply isn't a GenerateResponse")
+	}
+	res, err := w.Generate(ctx, req)
+	if err != nil {
+		return err
+	}
+	resp.Files = res.Files
+	return nil
+}
+
+func (w *wrapper) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	return nil, status.Error(codes.Unimplemented, "")
+}
+
+func HandleFunc(fn func(context.Context, *plugin.GenerateRequest) (*plugin.GenerateResponse, error)) Handler {
+	return &wrapper{fn}
+}

internal/ext/process/gen.go

@@ -0,0 +1,105 @@
+package process
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+
+	"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
+	Format string
+	Env    []string
+}
+
+func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+	req, ok := args.(protoreflect.ProtoMessage)
+	if !ok {
+		return fmt.Errorf("args isn't a protoreflect.ProtoMessage")
+	}
+
+	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
+	path, err := exec.LookPath(r.Cmd)
+	if err != nil {
+		return fmt.Errorf("process: %s not found", r.Cmd)
+	}
+
+	cmd := exec.CommandContext(ctx, path, method)
+	cmd.Stdin = bytes.NewReader(stdin)
+	cmd.Env = []string{
+		fmt.Sprintf("SQLC_VERSION=%s", info.Version),
+	}
+	for _, key := range r.Env {
+		if key == "SQLC_AUTH_TOKEN" {
+			continue
+		}
+		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, os.Getenv(key)))
+	}
+
+	out, err := cmd.Output()
+	if err != nil {
+		stderr := err.Error()
+		var exit *exec.ExitError
+		if errors.As(err, &exit) {
+			stderr = string(exit.Stderr)
+		}
+		return fmt.Errorf("process: error running command %s", stderr)
+	}
+
+	resp, ok := reply.(protoreflect.ProtoMessage)
+	if !ok {
+		return fmt.Errorf("reply isn't a protoreflect.ProtoMessage")
+	}
+
+	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
+}
+
+func (r *Runner) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	return nil, status.Error(codes.Unimplemented, "")
+}

internal/ext/wasm/runner.go

@@ -0,0 +1,7 @@
+package wasm
+
+type Runner struct {
+	URL    string
+	SHA256 string
+	Env    []string
+}

internal/ext/wasm/wasm.go

@@ -0,0 +1,291 @@
+package wasm
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"os"
+	"runtime"
+	"strings"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
+	"github.com/tetratelabs/wazero/sys"
+	"golang.org/x/sync/singleflight"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
+
+	"github.com/sqlc-dev/sqlc/internal/cache"
+	"github.com/sqlc-dev/sqlc/internal/info"
+	"github.com/sqlc-dev/sqlc/internal/plugin"
+)
+
+var flight singleflight.Group
+
+type runtimeAndCode struct {
+	rt   wazero.Runtime
+	code wazero.CompiledModule
+}
+
+// Verify the provided sha256 is valid.
+func (r *Runner) getChecksum(ctx context.Context) (string, error) {
+	if r.SHA256 != "" {
+		return r.SHA256, nil
+	}
+	// TODO: Add a log line here about something
+	_, sum, err := r.fetch(ctx, r.URL)
+	if err != nil {
+		return "", err
+	}
+	slog.Warn("fetching WASM binary to calculate sha256. Set this value in sqlc.yaml to prevent unneeded work", "sha256", sum)
+	return sum, nil
+}
+
+func (r *Runner) loadAndCompile(ctx context.Context) (*runtimeAndCode, error) {
+	expected, err := r.getChecksum(ctx)
+	if err != nil {
+		return nil, err
+	}
+	store, err := cache.Open()
+	if err != nil {
+		return nil, err
+	}
+	defer store.Close()
+	value, err, _ := flight.Do(expected, func() (any, error) {
+		return r.loadAndCompileWASM(ctx, store, expected)
+	})
+	if err != nil {
+		return nil, err
+	}
+	data, ok := value.(*runtimeAndCode)
+	if !ok {
+		return nil, fmt.Errorf("returned value was not a compiled module")
+	}
+	return data, nil
+}
+
+func (r *Runner) fetch(ctx context.Context, uri string) ([]byte, string, error) {
+	var body io.ReadCloser
+
+	switch {
+
+	case strings.HasPrefix(uri, "file://"):
+		file, err := os.Open(strings.TrimPrefix(uri, "file://"))
+		if err != nil {
+			return nil, "", fmt.Errorf("os.Open: %s %w", uri, err)
+		}
+		body = file
+
+	case strings.HasPrefix(uri, "https://"):
+		req, err := http.NewRequestWithContext(ctx, "GET", uri, nil)
+		if err != nil {
+			return nil, "", fmt.Errorf("http.Get: %s %w", uri, err)
+		}
+		req.Header.Set("User-Agent", fmt.Sprintf("sqlc/%s Go/%s (%s %s)", info.Version, runtime.Version(), runtime.GOOS, runtime.GOARCH))
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			return nil, "", fmt.Errorf("http.Get: %s %w", r.URL, err)
+		}
+		body = resp.Body
+
+	default:
+		return nil, "", fmt.Errorf("unknown scheme: %s", r.URL)
+	}
+
+	defer body.Close()
+
+	wmod, err := io.ReadAll(body)
+	if err != nil {
+		return nil, "", fmt.Errorf("readall: %w", err)
+	}
+
+	sum := sha256.Sum256(wmod)
+	actual := fmt.Sprintf("%x", sum)
+
+	return wmod, actual, nil
+}
+
+func (r *Runner) loadAndCompileWASM(ctx context.Context, store *cache.Cache, expected string) (*runtimeAndCode, error) {
+	// The sha256 declared in sqlc's configuration is a content address, so
+	// the plugin binary is looked up in the CAS directly by that checksum —
+	// no action cache entry is needed, and Get re-verifies the checksum on
+	// every hit.
+	wmod, err := store.CAS.Get(cache.SHA256Digest(expected))
+	if errors.Is(err, cache.ErrNotFound) {
+		var actual string
+		wmod, actual, err = r.fetch(ctx, r.URL)
+		if err != nil {
+			return nil, err
+		}
+		if expected != actual {
+			return nil, fmt.Errorf("invalid checksum: expected %s, got %s", expected, actual)
+		}
+		if _, err := store.CAS.Put(wmod); err != nil {
+			return nil, fmt.Errorf("cache wasm: %w", err)
+		}
+	} else if err != nil {
+		return nil, err
+	}
+
+	// Compiling the module to machine code is itself a cacheable action.
+	// Its only declared input is the module's checksum: the embedded wazero
+	// version and the target platform are determined by the sqlc binary,
+	// which is an implicit input of every action. Compiled artifacts are
+	// materialized into a private exec directory for wazero's compilation
+	// cache to find; the authoritative copies live in the CAS. Once wazero
+	// has loaded the module into memory the directory is no longer needed.
+	compileAction := store.NewAction("CompileModule").
+		AddInput("wasm", []byte(expected)).
+		Digest()
+
+	execDir, err := store.ExecDir(compileAction)
+	if err != nil {
+		return nil, err
+	}
+	defer os.RemoveAll(execDir)
+	compiled := true
+	if err := store.Actions.GetTree(compileAction, execDir); errors.Is(err, cache.ErrNotFound) {
+		compiled = false
+	} else if err != nil {
+		return nil, err
+	}
+
+	wazeroCache, err := wazero.NewCompilationCacheWithDir(execDir)
+	if err != nil {
+		return nil, fmt.Errorf("wazero.NewCompilationCacheWithDir: %w", err)
+	}
+
+	config := wazero.NewRuntimeConfig().WithCompilationCache(wazeroCache)
+	rt := wazero.NewRuntimeWithConfig(ctx, config)
+
+	if _, err := wasi_snapshot_preview1.Instantiate(ctx, rt); err != nil {
+		return nil, fmt.Errorf("wasi_snapshot_preview1 instantiate: %w", err)
+	}
+
+	// Compile the Wasm binary once so that we can skip the entire compilation
+	// time during instantiation. On an action cache hit this loads the
+	// materialized machine code instead of compiling.
+	code, err := rt.CompileModule(ctx, wmod)
+	if err != nil {
+		return nil, fmt.Errorf("compile module: %w", err)
+	}
+
+	if !compiled {
+		if err := store.Actions.PutTree(compileAction, execDir); err != nil {
+			slog.Warn("caching compiled module failed", "err", err)
+		}
+	}
+
+	return &runtimeAndCode{rt: rt, code: code}, nil
+}
+
+// removePGCatalog removes the pg_catalog schema from the request. There is a
+// mysterious (reason unknown) bug with wasm plugins when a large amount of
+// tables (like there are in the catalog) are sent.
+// @see https://github.com/sqlc-dev/sqlc/pull/1748
+func removePGCatalog(req *plugin.GenerateRequest) {
+	if req.Catalog == nil || req.Catalog.Schemas == nil {
+		return
+	}
+
+	filtered := make([]*plugin.Schema, 0, len(req.Catalog.Schemas))
+	for _, schema := range req.Catalog.Schemas {
+		if schema.Name == "pg_catalog" || schema.Name == "information_schema" {
+			continue
+		}
+
+		filtered = append(filtered, schema)
+	}
+
+	req.Catalog.Schemas = filtered
+}
+
+func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+	req, ok := args.(protoreflect.ProtoMessage)
+	if !ok {
+		return status.Error(codes.InvalidArgument, "args isn't a protoreflect.ProtoMessage")
+	}
+
+	// Remove the pg_catalog schema. Its sheer size causes unknown issues with wasm plugins
+	genReq, ok := req.(*plugin.GenerateRequest)
+	if ok {
+		removePGCatalog(genReq)
+		req = genReq
+	}
+
+	stdinBlob, err := proto.Marshal(req)
+	if err != nil {
+		return fmt.Errorf("failed to encode codegen request: %w", err)
+	}
+
+	runtimeAndCode, err := r.loadAndCompile(ctx)
+	if err != nil {
+		return fmt.Errorf("loadBytes: %w", err)
+	}
+
+	var stderr, stdout bytes.Buffer
+
+	conf := wazero.NewModuleConfig().
+		WithName("").
+		WithArgs("plugin.wasm", method).
+		WithStdin(bytes.NewReader(stdinBlob)).
+		WithStdout(&stdout).
+		WithStderr(&stderr).
+		WithEnv("SQLC_VERSION", info.Version)
+	for _, key := range r.Env {
+		conf = conf.WithEnv(key, os.Getenv(key))
+	}
+
+	result, err := runtimeAndCode.rt.InstantiateModule(ctx, runtimeAndCode.code, conf)
+	if err == nil {
+		defer result.Close(ctx)
+	}
+	if cerr := checkError(err, stderr); cerr != nil {
+		return cerr
+	}
+
+	// Print WASM stdout
+	stdoutBlob := stdout.Bytes()
+
+	resp, ok := reply.(protoreflect.ProtoMessage)
+	if !ok {
+		return fmt.Errorf("reply isn't a GenerateResponse")
+	}
+
+	if err := proto.Unmarshal(stdoutBlob, resp); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (r *Runner) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	return nil, status.Error(codes.Unimplemented, "")
+}
+
+func checkError(err error, stderr bytes.Buffer) error {
+	if err == nil {
+		return err
+	}
+
+	if exitErr, ok := err.(*sys.ExitError); ok {
+		if exitErr.ExitCode() == 0 {
+			return nil
+		}
+	}
+
+	// Print WASM stdout
+	stderrBlob := stderr.String()
+	if len(stderrBlob) > 0 {
+		return errors.New(stderrBlob)
+	}
+	return fmt.Errorf("call: %w", err)
+}