upstream-sqlc/internal/ext/wasm/wasm.go

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)
}