@@ -29,24 +29,13 @@ import (
29
29
"github.com/sqlc-dev/sqlc/internal/plugin"
30
30
)
31
31
32
- func cacheDir () (string , error ) {
33
- cache := os .Getenv ("SQLCCACHE" )
34
- if cache != "" {
35
- return cache , nil
36
- }
37
- cacheHome := os .Getenv ("XDG_CACHE_HOME" )
38
- if cacheHome == "" {
39
- home , err := os .UserHomeDir ()
40
- if err != nil {
41
- return "" , err
42
- }
43
- cacheHome = filepath .Join (home , ".cache" )
44
- }
45
- return filepath .Join (cacheHome , "sqlc" ), nil
46
- }
47
-
48
32
var flight singleflight.Group
49
33
34
+ type runtimeAndCode struct {
35
+ rt wazero.Runtime
36
+ code wazero.CompiledModule
37
+ }
38
+
50
39
// Verify the provided sha256 is valid.
51
40
func (r * Runner ) getChecksum (ctx context.Context ) (string , error ) {
52
41
if r .SHA256 != "" {
@@ -61,7 +50,7 @@ func (r *Runner) getChecksum(ctx context.Context) (string, error) {
61
50
return sum , nil
62
51
}
63
52
64
- func (r * Runner ) loadBytes (ctx context.Context ) ([] byte , error ) {
53
+ func (r * Runner ) loadAndCompile (ctx context.Context ) (* runtimeAndCode , error ) {
65
54
expected , err := r .getChecksum (ctx )
66
55
if err != nil {
67
56
return nil , err
@@ -71,14 +60,14 @@ func (r *Runner) loadBytes(ctx context.Context) ([]byte, error) {
71
60
return nil , err
72
61
}
73
62
value , err , _ := flight .Do (expected , func () (interface {}, error ) {
74
- return r .loadWASM (ctx , cacheDir , expected )
63
+ return r .loadAndCompileWASM (ctx , cacheDir , expected )
75
64
})
76
65
if err != nil {
77
66
return nil , err
78
67
}
79
- data , ok := value .([] byte )
68
+ data , ok := value .(* runtimeAndCode )
80
69
if ! ok {
81
- return nil , fmt .Errorf ("returned value was not a byte slice " )
70
+ return nil , fmt .Errorf ("returned value was not a compiled module " )
82
71
}
83
72
return data , nil
84
73
}
@@ -124,7 +113,7 @@ func (r *Runner) fetch(ctx context.Context, uri string) ([]byte, string, error)
124
113
return wmod , actual , nil
125
114
}
126
115
127
- func (r * Runner ) loadWASM (ctx context.Context , cache string , expected string ) ([] byte , error ) {
116
+ func (r * Runner ) loadAndCompileWASM (ctx context.Context , cache string , expected string ) (* runtimeAndCode , error ) {
128
117
pluginDir := filepath .Join (cache , expected )
129
118
pluginPath := filepath .Join (pluginDir , "plugin.wasm" )
130
119
_ , staterr := os .Stat (pluginPath )
@@ -153,7 +142,22 @@ func (r *Runner) loadWASM(ctx context.Context, cache string, expected string) ([
153
142
}
154
143
}
155
144
156
- return wmod , nil
145
+ wazeroCache , err := wazero .NewCompilationCacheWithDir (filepath .Join (cache , "wazero" ))
146
+ if err != nil {
147
+ return nil , fmt .Errorf ("wazero.NewCompilationCacheWithDir: %w" , err )
148
+ }
149
+ config := wazero .NewRuntimeConfig ().WithCompilationCache (wazeroCache )
150
+ rt := wazero .NewRuntimeWithConfig (ctx , config )
151
+ // TODO: Handle error
152
+ wasi_snapshot_preview1 .MustInstantiate (ctx , rt )
153
+
154
+ // Compile the Wasm binary once so that we can skip the entire compilation time during instantiation.
155
+ code , err := rt .CompileModule (ctx , wmod )
156
+ if err != nil {
157
+ return nil , fmt .Errorf ("compile module: %w" , err )
158
+ }
159
+
160
+ return & runtimeAndCode {rt : rt , code : code }, nil
157
161
}
158
162
159
163
// removePGCatalog removes the pg_catalog schema from the request. There is a
@@ -195,47 +199,25 @@ func (r *Runner) Invoke(ctx context.Context, method string, args any, reply any,
195
199
return fmt .Errorf ("failed to encode codegen request: %w" , err )
196
200
}
197
201
198
- cacheDir , err := cache . PluginsDir ( )
202
+ runtimeAndCode , err := r . loadAndCompile ( ctx )
199
203
if err != nil {
200
- return err
201
- }
202
-
203
- cache , err := wazero .NewCompilationCacheWithDir (filepath .Join (cacheDir , "wazero" ))
204
- if err != nil {
205
- return err
206
- }
207
-
208
- wasmBytes , err := r .loadBytes (ctx )
209
- if err != nil {
210
- return fmt .Errorf ("loadModule: %w" , err )
211
- }
212
-
213
- config := wazero .NewRuntimeConfig ().WithCompilationCache (cache )
214
- rt := wazero .NewRuntimeWithConfig (ctx , config )
215
- defer rt .Close (ctx )
216
-
217
- // TODO: Handle error
218
- wasi_snapshot_preview1 .MustInstantiate (ctx , rt )
219
-
220
- // Compile the Wasm binary once so that we can skip the entire compilation time during instantiation.
221
- mod , err := rt .CompileModule (ctx , wasmBytes )
222
- if err != nil {
223
- return err
204
+ return fmt .Errorf ("loadBytes: %w" , err )
224
205
}
225
206
226
207
var stderr , stdout bytes.Buffer
227
208
228
- conf := wazero .NewModuleConfig ()
229
- conf = conf .WithArgs ("plugin.wasm" , method )
230
- conf = conf .WithEnv ("SQLC_VERSION" , info .Version )
209
+ conf := wazero .NewModuleConfig ().
210
+ WithName ("" ).
211
+ WithArgs ("plugin.wasm" , method ).
212
+ WithStdin (bytes .NewReader (stdinBlob )).
213
+ WithStdout (& stdout ).
214
+ WithStderr (& stderr ).
215
+ WithEnv ("SQLC_VERSION" , info .Version )
231
216
for _ , key := range r .Env {
232
217
conf = conf .WithEnv (key , os .Getenv (key ))
233
218
}
234
- conf = conf .WithStdin (bytes .NewReader (stdinBlob ))
235
- conf = conf .WithStdout (& stdout )
236
- conf = conf .WithStderr (& stderr )
237
219
238
- result , err := rt .InstantiateModule (ctx , mod , conf )
220
+ result , err := runtimeAndCode . rt .InstantiateModule (ctx , runtimeAndCode . code , conf )
239
221
if result != nil {
240
222
defer result .Close (ctx )
241
223
}
0 commit comments