1package policy
  2
  3import (
  4	"context"
  5	"database/sql"
  6	"errors"
  7	"fmt"
  8
  9	"postern/internal/db"
 10
 11	"go.starlark.net/starlark"
 12)
 13
 14type KVModule struct {
 15	db *db.DB
 16}
 17
 18func NewKVModule(db *db.DB) *KVModule {
 19	return &KVModule{
 20		db: db,
 21	}
 22}
 23
 24func (k *KVModule) set(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 25	ctx, ok := thread.Local("context").(context.Context)
 26	if !ok || ctx == nil {
 27		ctx = context.Background()
 28	}
 29
 30	if err := ctx.Err(); err != nil {
 31		return nil, err
 32	}
 33
 34	var keyStr string
 35	var val starlark.Value
 36
 37	if err := starlark.UnpackArgs("kv.set", args, kwargs, "key", &keyStr, "value", &val); err != nil {
 38		return nil, err
 39	}
 40
 41	// Convert Starlark value to Go/SQLite compatible interface
 42	var goVal any
 43	switch v := val.(type) {
 44	case starlark.String:
 45		goVal = string(v)
 46	case starlark.Int:
 47		i, _ := v.Int64()
 48		goVal = i
 49	case starlark.Bool:
 50		goVal = bool(v)
 51	case starlark.Bytes:
 52		goVal = []byte(v)
 53	default:
 54		return nil, fmt.Errorf("kv.set: unsupported value type %s", val.Type())
 55	}
 56
 57	err := k.db.KVSet(ctx, keyStr, goVal)
 58	if err != nil {
 59		return nil, fmt.Errorf("kv.set: %w", err)
 60	}
 61
 62	return starlark.None, nil
 63}
 64
 65func (k *KVModule) get(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 66	ctx, ok := thread.Local("context").(context.Context)
 67	if !ok || ctx == nil {
 68		ctx = context.Background()
 69	}
 70
 71	if err := ctx.Err(); err != nil {
 72		return nil, err
 73	}
 74
 75	var keyStr string
 76	var defaultVal starlark.Value = starlark.None
 77
 78	if err := starlark.UnpackArgs("kv.get", args, kwargs, "key", &keyStr, "default?", &defaultVal); err != nil {
 79		return nil, err
 80	}
 81
 82	rawVal, err := k.db.KVGet(ctx, keyStr)
 83	if errors.Is(err, sql.ErrNoRows) {
 84		return defaultVal, nil
 85	} else if err != nil {
 86		return nil, err
 87	}
 88
 89	return rawToStarlark(rawVal)
 90}
 91
 92func (k *KVModule) inc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 93	ctx, ok := thread.Local("context").(context.Context)
 94	if !ok || ctx == nil {
 95		ctx = context.Background()
 96	}
 97
 98	if err := ctx.Err(); err != nil {
 99		return nil, err
100	}
101
102	var keyStr string
103	var delta int64 = 1
104
105	if err := starlark.UnpackArgs("kv.inc", args, kwargs, "key", &keyStr, "delta?", &delta); err != nil {
106		return nil, err
107	}
108
109	updatedVal, err := k.db.KVInc(ctx, keyStr, delta)
110	if err != nil {
111		return nil, err
112	}
113
114	return starlark.MakeInt64(updatedVal), nil
115}
116
117func (k *KVModule) pop(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
118	ctx, ok := thread.Local("context").(context.Context)
119	if !ok || ctx == nil {
120		ctx = context.Background()
121	}
122
123	if err := ctx.Err(); err != nil {
124		return nil, err
125	}
126
127	var keyStr string
128	if err := starlark.UnpackArgs("kv.pop", args, kwargs, "key", &keyStr); err != nil {
129		return nil, err
130	}
131
132	rawVal, err := k.db.KVPop(ctx, keyStr)
133	if errors.Is(err, sql.ErrNoRows) {
134		return starlark.None, nil
135	} else if err != nil {
136		return nil, err
137	}
138
139	return rawToStarlark(rawVal)
140}
141
142func (k *KVModule) String() string       { return "<module kv>" }
143func (k *KVModule) Type() string         { return "module" }
144func (k *KVModule) Freeze()              {}
145func (k *KVModule) Truth() starlark.Bool { return true }
146func (k *KVModule) Hash() (uint32, error) {
147	return 0, errors.New("unhashable")
148}
149
150func (k *KVModule) Attr(name string) (starlark.Value, error) {
151	switch name {
152	case "set":
153		return starlark.NewBuiltin("set", k.set), nil
154	case "get":
155		return starlark.NewBuiltin("get", k.get), nil
156	case "inc":
157		return starlark.NewBuiltin("inc", k.inc), nil
158	case "pop":
159		return starlark.NewBuiltin("pop", k.pop), nil
160	}
161	return nil, nil
162}
163
164func (k *KVModule) AttrNames() []string {
165	return []string{"get", "set", "inc", "pop"}
166}
167
168func rawToStarlark(rawVal any) (starlark.Value, error) {
169	switch v := rawVal.(type) {
170	case int64:
171		return starlark.MakeInt64(v), nil
172	case string:
173		return starlark.String(v), nil
174	case []byte:
175		return starlark.Bytes(v), nil
176	case bool:
177		return starlark.Bool(v), nil
178	default:
179		return starlark.None, errors.New("kv.get: unsupported value from database")
180	}
181}