package host
import (
"context"
"encoding/json"
extism "github.com/extism/go-sdk"
)
type KVStoreSetRequest struct {
Key string `json:"key"`
Value []byte `json:"value"`
}
type KVStoreSetResponse struct {
Error string `json:"error,omitempty"`
}
type KVStoreSetWithTTLRequest struct {
Key string `json:"key"`
Value []byte `json:"value"`
TtlSeconds int64 `json:"ttlSeconds"`
}
type KVStoreSetWithTTLResponse struct {
Error string `json:"error,omitempty"`
}
type KVStoreGetRequest struct {
Key string `json:"key"`
}
type KVStoreGetResponse struct {
Value []byte `json:"value,omitempty"`
Exists bool `json:"exists,omitempty"`
Error string `json:"error,omitempty"`
}
type KVStoreGetManyRequest struct {
Keys []string `json:"keys"`
}
type KVStoreGetManyResponse struct {
Values map[string][]byte `json:"values,omitempty"`
Error string `json:"error,omitempty"`
}
type KVStoreHasRequest struct {
Key string `json:"key"`
}
type KVStoreHasResponse struct {
Exists bool `json:"exists,omitempty"`
Error string `json:"error,omitempty"`
}
type KVStoreListRequest struct {
Prefix string `json:"prefix"`
}
type KVStoreListResponse struct {
Keys []string `json:"keys,omitempty"`
Error string `json:"error,omitempty"`
}
type KVStoreDeleteRequest struct {
Key string `json:"key"`
}
type KVStoreDeleteResponse struct {
Error string `json:"error,omitempty"`
}
type KVStoreDeleteByPrefixRequest struct {
Prefix string `json:"prefix"`
}
type KVStoreDeleteByPrefixResponse struct {
DeletedCount int64 `json:"deletedCount,omitempty"`
Error string `json:"error,omitempty"`
}
type KVStoreGetStorageUsedResponse struct {
Bytes int64 `json:"bytes,omitempty"`
Error string `json:"error,omitempty"`
}
func RegisterKVStoreHostFunctions(service KVStoreService) []extism.HostFunction {
return []extism.HostFunction{
newKVStoreSetHostFunction(service),
newKVStoreSetWithTTLHostFunction(service),
newKVStoreGetHostFunction(service),
newKVStoreGetManyHostFunction(service),
newKVStoreHasHostFunction(service),
newKVStoreListHostFunction(service),
newKVStoreDeleteHostFunction(service),
newKVStoreDeleteByPrefixHostFunction(service),
newKVStoreGetStorageUsedHostFunction(service),
}
}
func newKVStoreSetHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_set",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreSetRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
if svcErr := service.Set(ctx, req.Key, req.Value); svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreSetResponse{}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreSetWithTTLHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_setwithttl",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreSetWithTTLRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
if svcErr := service.SetWithTTL(ctx, req.Key, req.Value, req.TtlSeconds); svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreSetWithTTLResponse{}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreGetHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_get",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreGetRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
value, exists, svcErr := service.Get(ctx, req.Key)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreGetResponse{
Value: value,
Exists: exists,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreGetManyHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_getmany",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreGetManyRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
values, svcErr := service.GetMany(ctx, req.Keys)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreGetManyResponse{
Values: values,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreHasHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_has",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreHasRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
exists, svcErr := service.Has(ctx, req.Key)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreHasResponse{
Exists: exists,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreListHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_list",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreListRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
keys, svcErr := service.List(ctx, req.Prefix)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreListResponse{
Keys: keys,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreDeleteHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_delete",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreDeleteRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
if svcErr := service.Delete(ctx, req.Key); svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreDeleteResponse{}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreDeleteByPrefixHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_deletebyprefix",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
reqBytes, err := p.ReadBytes(stack[0])
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
var req KVStoreDeleteByPrefixRequest
if err := json.Unmarshal(reqBytes, &req); err != nil {
kvstoreWriteError(p, stack, err)
return
}
deletedcount, svcErr := service.DeleteByPrefix(ctx, req.Prefix)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreDeleteByPrefixResponse{
DeletedCount: deletedcount,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func newKVStoreGetStorageUsedHostFunction(service KVStoreService) extism.HostFunction {
return extism.NewHostFunctionWithStack(
"kvstore_getstorageused",
func(ctx context.Context, p *extism.CurrentPlugin, stack []uint64) {
bytes, svcErr := service.GetStorageUsed(ctx)
if svcErr != nil {
kvstoreWriteError(p, stack, svcErr)
return
}
resp := KVStoreGetStorageUsedResponse{
Bytes: bytes,
}
kvstoreWriteResponse(p, stack, resp)
},
[]extism.ValueType{extism.ValueTypePTR},
[]extism.ValueType{extism.ValueTypePTR},
)
}
func kvstoreWriteResponse(p *extism.CurrentPlugin, stack []uint64, resp any) {
respBytes, err := json.Marshal(resp)
if err != nil {
kvstoreWriteError(p, stack, err)
return
}
respPtr, err := p.WriteBytes(respBytes)
if err != nil {
stack[0] = 0
return
}
stack[0] = respPtr
}
func kvstoreWriteError(p *extism.CurrentPlugin, stack []uint64, err error) {
errResp := struct {
Error string `json:"error"`
}{Error: err.Error()}
respBytes, _ := json.Marshal(errResp)
respPtr, _ := p.WriteBytes(respBytes)
stack[0] = respPtr
}