Lua query fixes

common/space_lua/eval.ts

@@ -583,7 +583,12 @@ export async function evalStatement(
         .map((lval) => evalLValue(lval, env, sf)));
 
       for (let i = 0; i < lvalues.length; i++) {
-        luaSet(lvalues[i].env, lvalues[i].key, values[i], sf.withCtx(s.ctx));
+        await luaSet(
+          lvalues[i].env,
+          lvalues[i].key,
+          values[i],
+          sf.withCtx(s.ctx),
+        );
       }
 
       break;

common/space_lua/language_core_test.lua

@@ -142,7 +142,9 @@ assert_equal(t.foo, "Key not found: foo")
 t = setmetatable(
     {}, {
         __newindex = function(table, key, value)
+            print("Raw set", key, value)
             rawset(table, key, "Value: " .. value)
+            print("Raw set done")
         end
     }
 )
@@ -150,8 +152,8 @@ t = setmetatable(
 t.name = "John"
 -- rawset ignores the metamethod
 rawset(t, "age", 100)
-assert(t.name == "Value: John")
-assert(t.age == 100)
+assert_equal(t.name, "Value: John")
+assert_equal(t.age, 100)
 
 -- Test some of the operator metamethods
 t = setmetatable(

common/space_lua/query_collection.test.ts

@@ -1,5 +1,8 @@
 import { parseExpressionString } from "$common/space_lua/parse.ts";
-import { ArrayQueryCollection } from "./query_collection.ts";
+import {
+  ArrayQueryCollection,
+  findAllQueryVariables,
+} from "./query_collection.ts";
 import {
   LuaEnv,
   LuaNativeJSFunction,
@@ -137,3 +140,13 @@ Deno.test("ArrayQueryCollection", async () => {
   assertEquals(result9[2], "Jane Doe");
   assertEquals(result9[3], "Bob Johnson");
 });
+
+Deno.test("findAllQueryVariables", () => {
+  const query = {
+    where: parseExpressionString("p.x >= 2 and b.x >= 2"),
+    select: parseExpressionString("p.x + b.x"),
+    orderBy: [{ expr: parseExpressionString("q.x"), desc: false }],
+  };
+  const variables = findAllQueryVariables(query);
+  assertEquals(variables, ["p", "b", "q"]);
+});

common/space_lua/query_collection.ts

@@ -51,6 +51,73 @@ export type LuaCollectionQuery = {
   offset?: number;
 };
 
+export function findAllQueryVariables(query: LuaCollectionQuery): string[] {
+  const variables = new Set<string>();
+
+  // Helper to traverse an expression and collect variables
+  function findVariables(expr: LuaExpression) {
+    if (!expr) return;
+
+    switch (expr.type) {
+      case "Variable":
+        variables.add(expr.name);
+        break;
+      case "Binary":
+        findVariables(expr.left);
+        findVariables(expr.right);
+        break;
+      case "Unary":
+        findVariables(expr.argument);
+        break;
+      case "TableAccess":
+        findVariables(expr.object);
+        findVariables(expr.key);
+        break;
+      case "FunctionCall":
+        findVariables(expr.prefix);
+        expr.args.forEach(findVariables);
+        break;
+      case "TableConstructor":
+        expr.fields.forEach((field) => {
+          switch (field.type) {
+            case "DynamicField":
+              findVariables(field.key);
+              findVariables(field.value);
+              break;
+            case "PropField":
+              findVariables(field.value);
+              break;
+            case "ExpressionField":
+              findVariables(field.value);
+              break;
+          }
+        });
+        break;
+      case "PropertyAccess":
+        findVariables(expr.object);
+        break;
+      case "Parenthesized":
+        findVariables(expr.expression);
+        break;
+    }
+  }
+
+  // Check all parts of the query that can contain expressions
+  if (query.where) {
+    findVariables(query.where);
+  }
+
+  if (query.orderBy) {
+    query.orderBy.forEach((ob) => findVariables(ob.expr));
+  }
+
+  if (query.select) {
+    findVariables(query.select);
+  }
+
+  return Array.from(variables);
+}
+
 export interface LuaQueryCollection {
   query(
     query: LuaCollectionQuery,

common/space_lua/runtime.test.ts

@@ -48,4 +48,6 @@ Deno.test("Test Lua Rutime", async () => {
   assertEquals(await luaToString(new Promise((resolve) => resolve(1))), "1");
   assertEquals(await luaToString({ a: 1 }), "{a = 1}");
   assertEquals(await luaToString([{ a: 1 }]), "{{a = 1}}");
+  // Ensure simple cases are not returning promises
+  assertEquals(luaToString(10), "10");
 });

common/space_lua/runtime.ts

@@ -328,7 +328,13 @@ export class LuaTable implements ILuaSettable, ILuaGettable {
     return false;
   }
 
-  rawSet(key: LuaValue, value: LuaValue) {
+  rawSet(key: LuaValue, value: LuaValue): void | Promise<void> {
+    if (key instanceof Promise) {
+      return key.then((key) => this.rawSet(key, value));
+    }
+    if (value instanceof Promise) {
+      return value.then(() => this.rawSet(key, value));
+    }
     if (typeof key === "string") {
       this.stringKeys[key] = value;
     } else if (Number.isInteger(key) && key >= 1) {
@@ -360,7 +366,7 @@ export class LuaTable implements ILuaSettable, ILuaGettable {
     }
 
     // Just set the value
-    this.rawSet(key, value);
+    return this.rawSet(key, value);
   }
 
   rawGet(key: LuaValue): LuaValue | null {
@@ -480,7 +486,12 @@ export class LuaTable implements ILuaSettable, ILuaGettable {
 
 export type LuaLValueContainer = { env: ILuaSettable; key: LuaValue };
 
-export function luaSet(obj: any, key: any, value: any, sf: LuaStackFrame) {
+export async function luaSet(
+  obj: any,
+  key: any,
+  value: any,
+  sf: LuaStackFrame,
+): Promise<void> {
   if (!obj) {
     throw new LuaRuntimeError(
       `Not a settable object: nil`,
@@ -489,7 +500,7 @@ export function luaSet(obj: any, key: any, value: any, sf: LuaStackFrame) {
   }
 
   if (obj instanceof LuaTable || obj instanceof LuaEnv) {
-    obj.set(key, value, sf);
+    await obj.set(key, value, sf);
   } else {
     obj[key] = value;
   }
@@ -678,54 +689,56 @@ export function luaTruthy(value: any): boolean {
   return true;
 }
 
-export async function luaToString(value: any): Promise<string> {
+export function luaToString(value: any): string | Promise<string> {
   if (value === null || value === undefined) {
     return "nil";
   }
   if (value instanceof Promise) {
-    return luaToString(await value);
+    return value.then(luaToString);
   }
   if (value.toStringAsync) {
     return value.toStringAsync();
   }
   // Handle plain JavaScript objects in a Lua-like format
   if (typeof value === "object") {
-    let result = "{";
-    let first = true;
+    return (async () => {
+      let result = "{";
+      let first = true;
 
-    // Handle arrays
-    if (Array.isArray(value)) {
-      for (const val of value) {
+      // Handle arrays
+      if (Array.isArray(value)) {
+        for (const val of value) {
+          if (first) {
+            first = false;
+          } else {
+            result += ", ";
+          }
+          // Recursively stringify the value
+          const strVal = await luaToString(val);
+          result += strVal;
+        }
+        return result + "}";
+      }
+
+      // Handle objects
+      for (const [key, val] of Object.entries(value)) {
         if (first) {
           first = false;
         } else {
           result += ", ";
         }
+        if (/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(key)) {
+          result += `${key} = `;
+        } else {
+          result += `["${key}"] = `;
+        }
         // Recursively stringify the value
         const strVal = await luaToString(val);
         result += strVal;
       }
-      return result + "}";
-    }
-
-    // Handle objects
-    for (const [key, val] of Object.entries(value)) {
-      if (first) {
-        first = false;
-      } else {
-        result += ", ";
-      }
-      if (/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(key)) {
-        result += `${key} = `;
-      } else {
-        result += `["${key}"] = `;
-      }
-      // Recursively stringify the value
-      const strVal = await luaToString(val);
-      result += strVal;
-    }
-    result += "}";
-    return result;
+      result += "}";
+      return result;
+    })();
   }
   return String(value);
 }

common/space_lua/stdlib.ts

@@ -11,6 +11,7 @@ import {
   luaToString,
   luaTypeOf,
   type LuaValue,
+  luaValueToJS,
 } from "$common/space_lua/runtime.ts";
 import { stringApi } from "$common/space_lua/stdlib/string.ts";
 import { tableApi } from "$common/space_lua/stdlib/table.ts";
@@ -20,11 +21,13 @@ import {
   interpolateLuaString,
   spaceLuaApi,
 } from "$common/space_lua/stdlib/space_lua.ts";
-import type {
-  LuaCollectionQuery,
-  LuaQueryCollection,
+import {
+  findAllQueryVariables,
+  type LuaCollectionQuery,
+  type LuaQueryCollection,
 } from "$common/space_lua/query_collection.ts";
 import { templateApi } from "$common/space_lua/stdlib/template.ts";
+import { json } from "@codemirror/legacy-modes/X-ZUBjb2RlbWlycm9yL2xhbmd1YWdl/mode/javascript.d.ts";
 
 const printFunction = new LuaBuiltinFunction(async (_sf, ...args) => {
   console.log("[Lua]", ...(await Promise.all(args.map(luaToString))));
@@ -135,8 +138,7 @@ const setmetatableFunction = new LuaBuiltinFunction(
 
 const rawsetFunction = new LuaBuiltinFunction(
   (_sf, table: LuaTable, key: LuaValue, value: LuaValue) => {
-    table.rawSet(key, value);
-    return table;
+    return table.rawSet(key, value);
   },
 );
 
@@ -152,7 +154,28 @@ const tagFunction = new LuaBuiltinFunction(
       throw new LuaRuntimeError("Global not found", sf);
     }
     return {
-      query: async (query: LuaCollectionQuery): Promise<any[]> => {
+      query: async (query: LuaCollectionQuery, env: LuaEnv): Promise<any[]> => {
+        const localVars = findAllQueryVariables(query).filter((v) =>
+          !global.has(v) && v !== "_"
+        );
+        const scopedVariables: Record<string, any> = {};
+        for (const v of localVars) {
+          try {
+            const jsonValue = await luaValueToJS(env.get(v));
+            // Ensure this is JSON serializable
+            JSON.stringify(jsonValue);
+            scopedVariables[v] = jsonValue;
+          } catch (e: any) {
+            console.error(
+              "Failed to JSON serialize variable",
+              v,
+            );
+            throw new LuaRuntimeError(
+              `Failed to JSON serialize variable ${v} in query`,
+              sf,
+            );
+          }
+        }
         return (await global.get("datastore").get("query_lua").call(
           sf,
           [
@@ -160,6 +183,7 @@ const tagFunction = new LuaBuiltinFunction(
             tagName,
           ],
           query,
+          scopedVariables,
         )).toJSArray();
       },
     };

common/space_lua/stdlib/string.ts

@@ -6,6 +6,43 @@ import {
   luaToString,
 } from "$common/space_lua/runtime.ts";
 
+function createLuaMatcher(pattern: string, global = false) {
+  const jsPattern = pattern
+    .replace(/%(.)/g, (_, char) => {
+      switch (char) {
+        case ".":
+          return "[.]";
+        case "%":
+          return "%";
+        case "d":
+          return "\\d";
+        case "D":
+          return "\\D";
+        case "s":
+          return "\\s";
+        case "S":
+          return "\\S";
+        case "w":
+          return "\\w";
+        case "a":
+          return "[A-Za-z]";
+        case "l":
+          return "[a-z]";
+        case "u":
+          return "[A-Z]";
+        case "p":
+          return "[\\p{P}]";
+        default:
+          return char;
+      }
+    });
+
+  const regex = new RegExp(jsPattern, global ? "g" : undefined);
+  return (s: string) => {
+    return regex.exec(s);
+  };
+}
+
 export const stringApi = new LuaTable({
   byte: new LuaBuiltinFunction((_sf, s: string, i?: number, j?: number) => {
     i = i ?? 1;
@@ -34,27 +71,9 @@ export const stringApi = new LuaTable({
     },
   ),
   gmatch: new LuaBuiltinFunction((_sf, s: string, pattern: string) => {
-    const jsPattern = pattern
-      .replace(/%(.)/g, (_, char) => {
-        switch (char) {
-          case ".":
-            return "[.]";
-          case "%":
-            return "%";
-          case "d":
-            return "\\d";
-          case "s":
-            return "\\s";
-          case "w":
-            return "\\w";
-          default:
-            return char;
-        }
-      });
-
-    const regex = new RegExp(jsPattern, "g");
+    const matcher = createLuaMatcher(pattern, true);
     return () => {
-      const result = regex.exec(s);
+      const result = matcher(s);
       if (!result) {
         return;
       }
@@ -155,11 +174,12 @@ export const stringApi = new LuaTable({
   match: new LuaBuiltinFunction(
     (_sf, s: string, pattern: string, init?: number) => {
       init = init ?? 1;
-      const result = s.slice(init - 1).match(pattern);
+      const result = createLuaMatcher(pattern)(s.slice(init - 1));
       if (!result) {
         return new LuaMultiRes([]);
       }
-      return new LuaMultiRes(result.slice(1));
+      const captures = result.slice(1);
+      return new LuaMultiRes(captures.length > 0 ? captures : [result[0]]);
     },
   ),
   rep: new LuaBuiltinFunction((_sf, s: string, n: number, sep?: string) => {

common/space_lua/stdlib/string_test.lua

@@ -106,4 +106,38 @@ assert_equal(string.startswith("hello world", "world"), false)
 assert_equal(string.endswith("hello world", "world"), true)
 assert_equal(string.endswith("hello world", "hello"), false)
 
- 
+-- Extended string.match tests
+-- Basic pattern matching
+assert_equal(string.match("hello", "h"), "h")
+assert_equal(string.match("hello", "hello"), "hello")
+
+-- Test with no matches
+assert_equal(string.match("hello", "x"), nil)
+
+-- Test with captures
+local m1, m2 = string.match("hello", "(h)(ello)")
+assert_equal(m1, "h")
+assert_equal(m2, "ello")
+
+-- Test with init position
+local init_match = string.match("hello world", "(world)", 7)
+assert_equal(init_match, "world")
+
+-- Test init position with no match
+assert_equal(string.match("hello world", "hello", 7), nil)
+
+-- Test pattern characters
+assert_equal(string.match("123", "%d+"), "123")
+assert_equal(string.match("abc123", "%a+"), "abc")
+assert_equal(string.match("   abc", "%s+"), "   ")
+
+-- Test multiple captures
+local day, month, year = string.match("2024-03-14", "(%d+)-(%d+)-(%d+)")
+assert_equal(day, "2024")
+assert_equal(month, "03")
+assert_equal(year, "14")
+
+-- Test optional captures
+local word = string.match("The quick brown fox", "%s*(%w+)%s*")
+assert_equal(word, "The")
+

common/space_lua/stdlib/table.ts

@@ -81,6 +81,9 @@ export const tableApi = new LuaTable({
    */
   includes: new LuaBuiltinFunction(
     (sf, tbl: LuaTable | Record<string, any>, value: LuaValue) => {
+      if (!tbl) {
+        return false;
+      }
       if (tbl instanceof LuaTable) {
         // Iterate over the table
         for (const key of tbl.keys()) {

lib/plugos/syscalls/datastore.ts

@@ -6,7 +6,12 @@ import type { CommonSystem } from "$common/common_system.ts";
 import type { KV, KvKey, KvQuery } from "../../../plug-api/types.ts";
 import type { DataStore } from "../../data/datastore.ts";
 import type { SysCallMapping } from "../system.ts";
-import { LuaStackFrame, luaValueToJS } from "$common/space_lua/runtime.ts";
+import {
+  jsToLuaValue,
+  LuaEnv,
+  LuaStackFrame,
+  luaValueToJS,
+} from "$common/space_lua/runtime.ts";
 
 /**
  * Exposes the datastore API to plugs, but scoping everything to a prefix based on the plug's name
@@ -41,11 +46,16 @@ export function dataStoreReadSyscalls(
       _ctx,
       prefix: string[],
       query: LuaCollectionQuery,
+      scopeVariables: Record<string, any> = {},
     ): Promise<KV[]> => {
       const dsQueryCollection = new DataStoreQueryCollection(ds, prefix);
+      const env = new LuaEnv(commonSystem.spaceLuaEnv.env);
+      for (const [key, value] of Object.entries(scopeVariables)) {
+        env.set(key, jsToLuaValue(value));
+      }
       return (await dsQueryCollection.query(
         query,
-        commonSystem.spaceLuaEnv.env,
+        env,
         LuaStackFrame.lostFrame,
       )).map((item) => luaValueToJS(item));
     },

plug-api/syscalls/datastore.ts

@@ -74,8 +74,9 @@ export function query(
 export function queryLua(
   prefix: string[],
   query: LuaCollectionQuery,
+  scopeVariables: Record<string, any>,
 ): Promise<KV[]> {
-  return syscall("datastore.queryLua", prefix, query);
+  return syscall("datastore.queryLua", prefix, query, scopeVariables);
 }
 
 /**

server/server_system.ts

@@ -64,7 +64,7 @@ export class ServerSystem extends CommonSystem {
   }
 
   // Always needs to be invoked right after construction
-  async init(awaitIndex = false) {
+  async init(awaitIndex = false, performIndex = true) {
     this.system = new System(
       "server",
       {
@@ -219,10 +219,12 @@ export class ServerSystem extends CommonSystem {
 
     space.updatePageList().catch(console.error);
 
-    // Ensure a valid index
-    const indexPromise = ensureSpaceIndex(this.ds, this.system);
-    if (awaitIndex) {
-      await indexPromise;
+    if (performIndex) {
+      // Ensure a valid index
+      const indexPromise = ensureSpaceIndex(this.ds, this.system);
+      if (awaitIndex) {
+        await indexPromise;
+      }
     }
 
     await this.eventHook.dispatchEvent("system:ready");