Changeset 0.29.2 (#197)
Co-authored-by: Jeffrey Smith <jasafpro@gmail.com> Co-committed-by: Jeffrey Smith <jasafpro@gmail.com>
This commit is contained in:
254
server/sandbox/db_module_test.go
Normal file
254
server/sandbox/db_module_test.go
Normal file
@@ -0,0 +1,254 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"go.starlark.net/starlark"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// newTestDB creates an in-memory SQLite database with a test extension table.
|
||||
func newTestDB(t *testing.T) (*sql.DB, DBModuleConfig) {
|
||||
t.Helper()
|
||||
db, err := sql.Open("sqlite", ":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { db.Close() })
|
||||
|
||||
// Create the extension table the tests will use.
|
||||
// physicalTable("logs") for package "test-ext" → ext_test_ext_logs
|
||||
_, err = db.Exec(`CREATE TABLE ext_test_ext_logs (
|
||||
id TEXT PRIMARY KEY,
|
||||
message TEXT,
|
||||
user_id TEXT,
|
||||
count INTEGER,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
)`)
|
||||
if err != nil {
|
||||
t.Fatalf("create test table: %v", err)
|
||||
}
|
||||
|
||||
cfg := DBModuleConfig{
|
||||
PackageID: "test-ext",
|
||||
CanWrite: true,
|
||||
DB: db,
|
||||
IsPostgres: false,
|
||||
}
|
||||
return db, cfg
|
||||
}
|
||||
|
||||
// execScript runs a Starlark snippet with the db module injected.
|
||||
func execScript(t *testing.T, cfg DBModuleConfig, script string) (*Result, error) {
|
||||
t.Helper()
|
||||
sb := New(DefaultConfig())
|
||||
ctx := context.Background()
|
||||
modules := map[string]starlark.Value{
|
||||
"db": BuildDBModule(ctx, cfg),
|
||||
}
|
||||
return sb.Exec(ctx, "test.star", script, modules)
|
||||
}
|
||||
|
||||
// ── physicalTable ───────────────────────────
|
||||
|
||||
func TestPhysicalTable_ValidName(t *testing.T) {
|
||||
cfg := DBModuleConfig{PackageID: "my-ext"}
|
||||
got, err := cfg.physicalTable("logs")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "ext_my_ext_logs" {
|
||||
t.Errorf("got %q, want ext_my_ext_logs", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPhysicalTable_EmptyName(t *testing.T) {
|
||||
cfg := DBModuleConfig{PackageID: "my-ext"}
|
||||
_, err := cfg.physicalTable("")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty table name")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPhysicalTable_InvalidChars(t *testing.T) {
|
||||
cfg := DBModuleConfig{PackageID: "my-ext"}
|
||||
_, err := cfg.physicalTable("bad name")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for table name with space")
|
||||
}
|
||||
_, err = cfg.physicalTable("drop;table")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for table name with semicolon")
|
||||
}
|
||||
}
|
||||
|
||||
// ── db.insert / db.query ────────────────────
|
||||
|
||||
func TestDBInsertAndQuery(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
|
||||
result, err := execScript(t, cfg, `
|
||||
row = db.insert("logs", {"message": "hello", "user_id": "u1"})
|
||||
rows = db.query("logs")
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
|
||||
rowsVal, ok := result.Globals["rows"]
|
||||
if !ok {
|
||||
t.Fatal("rows not in globals")
|
||||
}
|
||||
if _, ok := rowsVal.(*starlark.List); !ok {
|
||||
t.Fatalf("rows is %T, want *starlark.List", rowsVal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDBQueryFilters(t *testing.T) {
|
||||
db, cfg := newTestDB(t)
|
||||
|
||||
// Pre-populate with two rows
|
||||
db.Exec(`INSERT INTO ext_test_ext_logs (id, message, user_id) VALUES ('a', 'one', 'u1')`)
|
||||
db.Exec(`INSERT INTO ext_test_ext_logs (id, message, user_id) VALUES ('b', 'two', 'u2')`)
|
||||
|
||||
result, err := execScript(t, cfg, `
|
||||
rows = db.query("logs", filters={"user_id": "u1"})
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
|
||||
rowsStr := result.Globals["rows"].String()
|
||||
if !strings.Contains(rowsStr, "one") {
|
||||
t.Errorf("expected row with message=one, got: %s", rowsStr)
|
||||
}
|
||||
if strings.Contains(rowsStr, "two") {
|
||||
t.Errorf("unexpected row with message=two: %s", rowsStr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDBQueryLimit(t *testing.T) {
|
||||
db, cfg := newTestDB(t)
|
||||
for i := 0; i < 5; i++ {
|
||||
db.Exec(`INSERT INTO ext_test_ext_logs (id, message) VALUES (?, ?)`,
|
||||
generateID(), "msg")
|
||||
}
|
||||
|
||||
result, err := execScript(t, cfg, `rows = db.query("logs", limit=2)`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
// String representation of a 2-element list has exactly 2 entries
|
||||
rowsStr := result.Globals["rows"].String()
|
||||
// Count occurrences of "id" key — each row has one
|
||||
count := strings.Count(rowsStr, `"id"`)
|
||||
if count != 2 {
|
||||
t.Errorf("expected 2 rows with limit=2, got repr: %s", rowsStr)
|
||||
}
|
||||
}
|
||||
|
||||
// ── db.update ───────────────────────────────
|
||||
|
||||
func TestDBUpdate(t *testing.T) {
|
||||
db, cfg := newTestDB(t)
|
||||
db.Exec(`INSERT INTO ext_test_ext_logs (id, message) VALUES ('r1', 'original')`)
|
||||
|
||||
_, err := execScript(t, cfg, `db.update("logs", "r1", {"message": "updated"})`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
|
||||
var msg string
|
||||
db.QueryRow(`SELECT message FROM ext_test_ext_logs WHERE id = 'r1'`).Scan(&msg)
|
||||
if msg != "updated" {
|
||||
t.Errorf("expected message=updated, got %q", msg)
|
||||
}
|
||||
}
|
||||
|
||||
// ── db.delete ───────────────────────────────
|
||||
|
||||
func TestDBDelete(t *testing.T) {
|
||||
db, cfg := newTestDB(t)
|
||||
db.Exec(`INSERT INTO ext_test_ext_logs (id, message) VALUES ('r1', 'to-delete')`)
|
||||
|
||||
_, err := execScript(t, cfg, `db.delete("logs", "r1")`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
|
||||
var count int
|
||||
db.QueryRow(`SELECT COUNT(*) FROM ext_test_ext_logs WHERE id = 'r1'`).Scan(&count)
|
||||
if count != 0 {
|
||||
t.Errorf("expected row deleted, count = %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
// ── write guard ─────────────────────────────
|
||||
|
||||
func TestDBReadOnlyModuleNoWriteFunctions(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
cfg.CanWrite = false
|
||||
|
||||
mod := BuildDBModule(context.Background(), cfg)
|
||||
modStr := mod.String()
|
||||
|
||||
// insert/update/delete should not be present in a read-only module
|
||||
for _, fn := range []string{"insert", "update", "delete"} {
|
||||
if strings.Contains(modStr, fn) {
|
||||
t.Errorf("read-only module exposes write function %q", fn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDBWriteFunctionBlockedByReadOnly(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
cfg.CanWrite = false
|
||||
|
||||
_, err := execScript(t, cfg, `db.insert("logs", {"message": "x"})`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error calling insert on read-only module")
|
||||
}
|
||||
}
|
||||
|
||||
// ── view allowlist ───────────────────────────
|
||||
|
||||
func TestDBViewRejectsUnknownView(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
_, err := execScript(t, cfg, `db.view("secrets")`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown view")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unknown view") {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── namespace isolation ──────────────────────
|
||||
|
||||
func TestDBQueryRejectsOtherExtensionTable(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
// Attempt to query a table that belongs to a different extension
|
||||
_, err := execScript(t, cfg, `db.query("../../etc/passwd")`)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for path traversal attempt")
|
||||
}
|
||||
}
|
||||
|
||||
// ── list_tables ─────────────────────────────
|
||||
|
||||
func TestDBListTables(t *testing.T) {
|
||||
_, cfg := newTestDB(t)
|
||||
result, err := execScript(t, cfg, `names = db.list_tables()`)
|
||||
if err != nil {
|
||||
t.Fatalf("script error: %v", err)
|
||||
}
|
||||
namesStr := result.Globals["names"].String()
|
||||
// The test table ext_test_ext_logs should appear as "logs"
|
||||
if !strings.Contains(namesStr, "logs") {
|
||||
t.Errorf("expected 'logs' in list_tables output, got: %s", namesStr)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user