core/server/sandbox/db_module_test.go

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