Lua integrated query work (#1205)

* Lua query embeddings

build_web.ts

@@ -16,7 +16,7 @@ export async function bundleAll(
   await buildCopyBundleAssets();
   let timer;
   if (watch) {
-    const watcher = Deno.watchFs(["web", "dist_plug_bundle"]);
+    const watcher = Deno.watchFs(["web", "common", "dist_plug_bundle"]);
     for await (const _event of watcher) {
       if (timer) {
         clearTimeout(timer);

common/space_lua.ts

@@ -16,7 +16,6 @@ import {
 import type { ScriptEnvironment } from "$common/space_script.ts";
 import { luaValueToJS } from "$common/space_lua/runtime.ts";
 import type { ASTCtx } from "$common/space_lua/ast.ts";
-import type { ObjectQuery } from "@silverbulletmd/silverbullet/types";
 import { buildLuaEnv } from "$common/space_lua_api.ts";
 
 export class SpaceLuaEnvironment {
@@ -26,21 +25,23 @@ export class SpaceLuaEnvironment {
    * Loads all Lua scripts from the database and evaluates them in a new environment
    * @param system
    */
-  async reload(system: System<any>, scriptEnv: ScriptEnvironment) {
+  async reload(
+    system: System<any>,
+    scriptEnv: ScriptEnvironment,
+  ) {
     const allScripts: ScriptObject[] = await system.invokeFunction(
       "index.queryObjects",
-      ["space-lua", {
-        // This is a bit silly, but at least makes the order deterministic
-        orderBy: [{ expr: ["attr", "ref"] }],
-      } as ObjectQuery],
+      ["space-lua", {}],
     );
     this.env = buildLuaEnv(system, scriptEnv);
+    const tl = new LuaEnv();
     for (const script of allScripts) {
       try {
+        console.log("Now evaluating", script.ref);
         const ast = parseLua(script.script, { ref: script.ref });
         // We create a local scope for each script
         const scriptEnv = new LuaEnv(this.env);
-        const sf = new LuaStackFrame(new LuaEnv(), ast.ctx);
+        const sf = new LuaStackFrame(tl, ast.ctx);
         await evalStatement(ast, scriptEnv, sf);
       } catch (e: any) {
         if (e instanceof LuaRuntimeError) {
@@ -66,7 +67,7 @@ export class SpaceLuaEnvironment {
           `[Lua] Registering global function '${globalName}' (source: ${value.body.ctx.ref})`,
         );
         scriptEnv.registerFunction({ name: globalName }, (...args: any[]) => {
-          const sf = new LuaStackFrame(new LuaEnv(), value.body.ctx);
+          const sf = new LuaStackFrame(tl, value.body.ctx);
           return luaValueToJS(value.call(sf, ...args.map(jsToLuaValue)));
         });
       }

common/space_lua/ast.ts

@@ -150,7 +150,8 @@ export type LuaExpression =
   | LuaBinaryExpression
   | LuaUnaryExpression
   | LuaTableConstructor
-  | LuaFunctionDefinition;
+  | LuaFunctionDefinition
+  | LuaQueryExpression;
 
 export type LuaNilLiteral = {
   type: "Nil";
@@ -254,3 +255,49 @@ export type LuaFunctionDefinition = {
   type: "FunctionDefinition";
   body: LuaFunctionBody;
 } & ASTContext;
+
+// Query stuff
+export type LuaQueryExpression = {
+  type: "Query";
+  clauses: LuaQueryClause[];
+} & ASTContext;
+
+export type LuaQueryClause =
+  | LuaFromClause
+  | LuaWhereClause
+  | LuaLimitClause
+  | LuaOrderByClause
+  | LuaSelectClause;
+
+export type LuaFromClause = {
+  type: "From";
+  name: string;
+  expression: LuaExpression;
+} & ASTContext;
+
+export type LuaWhereClause = {
+  type: "Where";
+  expression: LuaExpression;
+} & ASTContext;
+
+export type LuaLimitClause = {
+  type: "Limit";
+  limit: LuaExpression;
+  offset?: LuaExpression;
+} & ASTContext;
+
+export type LuaOrderByClause = {
+  type: "OrderBy";
+  orderBy: LuaOrderBy[];
+} & ASTContext;
+
+export type LuaOrderBy = {
+  type: "Order";
+  expression: LuaExpression;
+  direction: "asc" | "desc";
+} & ASTContext;
+
+export type LuaSelectClause = {
+  type: "Select";
+  tableConstructor: LuaTableConstructor;
+} & ASTContext;

common/space_lua/eval.ts

@@ -29,6 +29,12 @@ import {
   type LuaValue,
   singleResult,
 } from "./runtime.ts";
+import {
+  ArrayQueryCollection,
+  type LuaCollectionQuery,
+} from "$common/space_lua/query_collection.ts";
+import { luaValueToJS } from "$common/space_lua/runtime.ts";
+import { jsToLuaValue } from "$common/space_lua/runtime.ts";
 
 export function evalExpression(
   e: LuaExpression,
@@ -244,6 +250,90 @@ export function evalExpression(
       case "FunctionDefinition": {
         return new LuaFunction(e.body, env);
       }
+      case "Query": {
+        // console.log("Query", e);
+        const findFromClause = e.clauses.find((c) => c.type === "From");
+        if (!findFromClause) {
+          throw new LuaRuntimeError("No from clause found", sf.withCtx(e.ctx));
+        }
+        const objectVariable = findFromClause.name;
+        const objectExpression = findFromClause.expression;
+        return Promise.resolve(evalExpression(objectExpression, env, sf)).then(
+          async (collection: LuaValue) => {
+            if (!collection) {
+              throw new LuaRuntimeError(
+                "Collection is nil",
+                sf.withCtx(e.ctx),
+              );
+            }
+            // Check if collection is a queryable collection
+            if (!collection.query) {
+              // If not, try to convert it to JS and see if it's an array
+              collection = await luaValueToJS(collection);
+              if (!Array.isArray(collection)) {
+                throw new LuaRuntimeError(
+                  "Collection does not support query",
+                  sf.withCtx(e.ctx),
+                );
+              }
+              collection = new ArrayQueryCollection(collection);
+            }
+            // Build up query object
+            const query: LuaCollectionQuery = {
+              objectVariable,
+            };
+
+            // Map clauses to query parameters
+            for (const clause of e.clauses) {
+              switch (clause.type) {
+                case "Where": {
+                  query.where = clause.expression;
+                  break;
+                }
+                case "OrderBy": {
+                  query.orderBy = clause.orderBy.map((o) => ({
+                    expr: o.expression,
+                    desc: o.direction === "desc",
+                  }));
+                  break;
+                }
+                case "Select": {
+                  query.select = clause.tableConstructor.fields.map((f) => {
+                    if (f.type === "PropField") {
+                      return {
+                        name: f.key,
+                        expr: f.value,
+                      };
+                    } else {
+                      throw new LuaRuntimeError(
+                        "Select fields must be named",
+                        sf.withCtx(f.ctx),
+                      );
+                    }
+                  });
+
+                  break;
+                }
+                case "Limit": {
+                  const limitVal = await evalExpression(clause.limit, env, sf);
+                  query.limit = Number(limitVal);
+                  if (clause.offset) {
+                    const offsetVal = await evalExpression(
+                      clause.offset,
+                      env,
+                      sf,
+                    );
+                    query.offset = Number(offsetVal);
+                  }
+                  break;
+                }
+              }
+            }
+
+            return collection.query(query, env, sf).then(jsToLuaValue);
+          },
+        );
+      }
       default:
         throw new Error(`Unknown expression type ${e.type}`);
     }

common/space_lua/language.test.ts

@@ -1,6 +1,10 @@
 import { parse } from "$common/space_lua/parse.ts";
 import { luaBuildStandardEnv } from "$common/space_lua/stdlib.ts";
-import { LuaEnv, LuaStackFrame } from "$common/space_lua/runtime.ts";
+import {
+  LuaEnv,
+  LuaRuntimeError,
+  LuaStackFrame,
+} from "$common/space_lua/runtime.ts";
 import { evalStatement } from "$common/space_lua/eval.ts";
 import { assert } from "@std/assert/assert";
 import { fileURLToPath } from "node:url";
@@ -18,7 +22,11 @@ Deno.test("Lua language tests", async () => {
   try {
     await evalStatement(chunk, env, sf);
   } catch (e: any) {
-    console.error(`Error evaluating script:`, e.toPrettyString(luaFile));
+    if (e instanceof LuaRuntimeError) {
+      console.error(`Error evaluating script:`, e.toPrettyString(luaFile));
+    } else {
+      console.error(`Error evaluating script:`, e);
+    }
     assert(false);
   }
 });

common/space_lua/language_test.lua

@@ -700,3 +700,33 @@ assert(evalResult == 2, "Eval should return 2")
 local parsedExpr = space_lua.parse_expression("tostring(a + 1)")
 local evalResult = space_lua.eval_expression(parsedExpr, { a = 1 })
 assert(evalResult == "2", "Eval should return 2 as a string")
+
+-- Test query
+local data = { { name = "John", lastModified = 1, age = 20 }, { name = "Jane", lastModified = 2, age = 21 } }
+local r = query [[from p = data limit 1]]
+assert_equal(#r, 1)
+assert_equal(r[1].name, "John")
+assert_equal(r[1].lastModified, 1)
+
+local r = query [[from p = data order by p.lastModified desc]]
+assert_equal(#r, 2)
+assert_equal(r[1].name, "Jane")
+assert_equal(r[1].lastModified, 2)
+assert_equal(r[2].name, "John")
+assert_equal(r[2].lastModified, 1)
+
+local r = query [[from p = data order by p.lastModified]]
+assert_equal(#r, 2)
+assert_equal(r[1].name, "John")
+assert_equal(r[1].lastModified, 1)
+assert_equal(r[2].name, "Jane")
+assert_equal(r[2].lastModified, 2)
+
+local r = query [[from p = data order by p.age select {name=p.name, age=p.age}]]
+assert_equal(#r, 2)
+assert_equal(r[1].name, "John")
+assert_equal(r[1].age, 20)
+assert_equal(r[2].name, "Jane")
+assert_equal(r[2].age, 21)
+assert_equal(r[1].lastModified, nil)
+assert_equal(r[2].lastModified, nil)

common/space_lua/lua.grammar

@@ -17,6 +17,10 @@
 
 @top Chunk { Block }
 
+kw<term> { @specialize[@name={term}]<identifier, term> }
+ckw<term> { @extend[@name={term}]<identifier, term> }
+list<term> { term ("," term)* }
+
 Block { statement* ReturnStatement? }
 
 ReturnStatement { kw<"return"> ExpList? ";"?}
@@ -58,7 +62,6 @@ ForStatement {
 FuncName { Name ("." Name)* (":" Name)? }
 FuncBody { "(" ArgList ")" Block kw<"end"> }
 
-list<term> { term ("," term)* }
 
 NameList { list<Name> }
 ExpList { list<exp> }
@@ -78,29 +81,29 @@ exp {
   BinaryExpression |
   UnaryExpression |
   TableConstructor |
-  FunctionDef { kw<"function"> FuncBody }
-  // |  Query
+  FunctionDef { kw<"function"> FuncBody } |
+  Query
 }
 
 Query {
-  "[" exp QueryClause* "]" 
+  kw<"query"> "[[" QueryClause* "]]" 
 }
 
 QueryClause {
+  FromClause |
   WhereClause |
   OrderByClause |
   SelectClause |
-  RenderClause |
   LimitClause
 }
 
-WhereClause { kw<"where"> exp }
-LimitClause { kw<"limit"> exp }
-OrderByClause { kw<"order"> kw<"by"> exp kw<"desc">? }
-SelectClause { kw<"select"> list<Select> }
-RenderClause { kw<"render"> ( kw<"each"> | kw<"all"> )? simpleString }
+FromClause { ckw<"from"> Name "=" exp }
+WhereClause { ckw<"where"> exp }
+LimitClause { ckw<"limit"> exp ("," exp)? }
+OrderByClause { ckw<"order">  ckw<"by"> list<OrderBy> }
+OrderBy { exp ckw<"desc">? }
+SelectClause { ckw<"select"> TableConstructor }
 
-Select { Name | exp kw<"as"> Name }
 
 field[@isGroup=Field] {
   FieldDynamic { "[" exp "]" "=" exp } |
@@ -113,6 +116,7 @@ prefixexp {
   Parens { "(" exp ")" ~parens } |
   FunctionCall ~fcall
 }
+
 FunctionCall { prefixexp (":" Name)? !call args }
 args {
   LiteralString |
@@ -124,8 +128,6 @@ var {
   Name | Property { (prefixexp "." Name) } | MemberExpression { (prefixexp "[" exp "]") }
 }
 
-kw<term> { @specialize[@name={term}]<identifier, term> }
-
 Name { identifier }
 Label { "::" Name "::" }
 LiteralString { simpleString }
@@ -154,11 +156,7 @@ TableConstructor { "{" (field (fieldsep field)* fieldsep?)? "}" }
 @tokens {
   CompareOp { "<" | ">" | $[<>=~/!] "=" }
 
-  TagIdentifier { @asciiLetter (@asciiLetter | @digit | "-" | "_" | "/" )* }
-
-  word { (std.asciiLetter | "_") (std.digit | std.asciiLetter | "_")* }
-
-  identifier { word }
+  identifier { (std.asciiLetter | "_") (std.digit | std.asciiLetter | "_")* }
 
   stringEscape {
     "\\" ($[abfnz"'\\] | digit digit? digit?) |

common/space_lua/parse-lua.terms.js

@@ -14,15 +14,23 @@ export const
   UnaryExpression = 51,
   FuncBody = 58,
   ArgList = 59,
-  IfStatement = 63,
-  ForStatement = 68,
-  ForNumeric = 70,
-  ForGeneric = 71,
-  NameList = 72,
-  ExpList = 74,
-  FuncName = 76,
-  VarList = 80,
-  AttNameList = 82,
-  AttName = 83,
-  Attrib = 84,
-  ReturnStatement = 85
+  Query = 60,
+  QueryClause = 62,
+  FromClause = 63,
+  WhereClause = 65,
+  OrderByClause = 67,
+  OrderBy = 70,
+  SelectClause = 72,
+  LimitClause = 74,
+  IfStatement = 79,
+  ForStatement = 84,
+  ForNumeric = 86,
+  ForGeneric = 87,
+  NameList = 88,
+  ExpList = 90,
+  FuncName = 92,
+  VarList = 96,
+  AttNameList = 98,
+  AttName = 99,
+  Attrib = 100,
+  ReturnStatement = 101

common/space_lua/parse.test.ts

@@ -21,6 +21,9 @@ Deno.test("Test Lua parser", () => {
   parse(`e(a.b.c)`);
   parse(`e((1+2))`);
 
+  // Use keywordy variables
+  parse(`e(order, limit, where)`);
+
   // Table expressions
   parse(`e({})`);
   parse(`e({1, 2, 3, })`);
@@ -99,3 +102,12 @@ Deno.test("Test comment handling", () => {
         -- yo
       ]])`);
 });
+
+Deno.test("Test query parsing", () => {
+  parse(`_(query[[from p = tag("page") where p.name == "John" limit 10, 3]])`);
+  parse(`_(query[[from p = tag("page") select {name="hello", age=10}]])`);
+  parse(
+    `_(query[[from p = tag("page") order by p.lastModified desc, p.name]])`,
+  );
+  parse(`_(query[[from p = tag("page") order by p.lastModified]])`);
+});

common/space_lua/parse.ts

@@ -16,8 +16,11 @@ import type {
   LuaFunctionCallStatement,
   LuaFunctionName,
   LuaLValue,
+  LuaOrderBy,
   LuaPrefixExpression,
+  LuaQueryClause,
   LuaStatement,
+  LuaTableConstructor,
   LuaTableField,
 } from "./ast.ts";
 import { tags as t } from "@lezer/highlight";
@@ -29,7 +32,7 @@ const luaStyleTags = styleTags({
   CompareOp: t.operator,
   "true false": t.bool,
   Comment: t.lineComment,
-  "return break goto do end while repeat until function local if then else elseif in for nil or and not":
+  "return break goto do end while repeat until function local if then else elseif in for nil or and not query from where limit select order by desc":
     t.keyword,
 });
 
@@ -464,12 +467,84 @@ function parseExpression(t: ParseTree, ctx: ASTCtx): LuaExpression {
       };
     case "nil":
       return { type: "Nil", ctx: context(t, ctx) };
+    case "Query":
+      return {
+        type: "Query",
+        clauses: t.children!.slice(2, -1).map((c) => parseQueryClause(c, ctx)),
+        ctx: context(t, ctx),
+      };
     default:
       console.error(t);
       throw new Error(`Unknown expression type: ${t.type}`);
   }
 }
 
+function parseQueryClause(t: ParseTree, ctx: ASTCtx): LuaQueryClause {
+  if (t.type !== "QueryClause") {
+    throw new Error(`Expected QueryClause, got ${t.type}`);
+  }
+  t = t.children![0];
+  switch (t.type) {
+    case "FromClause": {
+      return {
+        type: "From",
+        name: t.children![1].children![0].text!,
+        expression: parseExpression(t.children![3], ctx),
+        ctx: context(t, ctx),
+      };
+    }
+    case "WhereClause":
+      return {
+        type: "Where",
+        expression: parseExpression(t.children![1], ctx),
+        ctx: context(t, ctx),
+      };
+    case "LimitClause": {
+      const limit = parseExpression(t.children![1], ctx);
+      const offset = t.children![2]
+        ? parseExpression(t.children![3], ctx)
+        : undefined;
+      return {
+        type: "Limit",
+        limit,
+        offset,
+        ctx: context(t, ctx),
+      };
+    }
+    case "OrderByClause": {
+      const orderBy: LuaOrderBy[] = [];
+      for (const child of t.children!) {
+        if (child.type === "OrderBy") {
+          orderBy.push({
+            type: "Order",
+            expression: parseExpression(child.children![0], ctx),
+            direction: child.children![1]?.type === "desc" ? "desc" : "asc",
+            ctx: context(child, ctx),
+          });
+        }
+      }
+      return {
+        type: "OrderBy",
+        orderBy,
+        ctx: context(t, ctx),
+      };
+    }
+    case "SelectClause": {
+      return {
+        type: "Select",
+        tableConstructor: parseExpression(
+          t.children![1],
+          ctx,
+        ) as LuaTableConstructor,
+        ctx: context(t, ctx),
+      };
+    }
+    default:
+      console.error(t);
+      throw new Error(`Unknown query clause type: ${t.type}`);
+  }
+}
+
 function parseFunctionArgs(ts: ParseTree[], ctx: ASTCtx): LuaExpression[] {
   return ts.filter((t) => ![",", "(", ")"].includes(t.type!)).map(
     (e) => parseExpression(e, ctx),
@@ -639,7 +714,8 @@ export function parse(s: string, ctx: ASTCtx = {}): LuaBlock {
 }
 
 export function parseToCrudeAST(t: string): ParseTree {
-  return cleanTree(lezerToParseTree(t, parser.parse(t).topNode), true);
+  const n = lezerToParseTree(t, parser.parse(t).topNode);
+  return cleanTree(n, true);
 }
 
 /**

common/space_lua/query_collection.test.ts

@@ -22,7 +22,8 @@ Deno.test("ArrayQueryCollection", async () => {
   }]);
   const result = await collection.query(
     {
-      where: parseExpressionString("x >= 2"),
+      objectVariable: "p",
+      where: parseExpressionString("p.x >= 2"),
     },
     rootEnv,
     LuaStackFrame.lostFrame,
@@ -33,6 +34,7 @@ Deno.test("ArrayQueryCollection", async () => {
   // Test limit
   const result2 = await collection.query(
     {
+      objectVariable: "p",
       limit: 1,
     },
     rootEnv,
@@ -44,6 +46,7 @@ Deno.test("ArrayQueryCollection", async () => {
   // Test offset
   const result3 = await collection.query(
     {
+      objectVariable: "p",
       offset: 1,
     },
     rootEnv,
@@ -55,7 +58,8 @@ Deno.test("ArrayQueryCollection", async () => {
   // Test order by
   const result4 = await collection.query(
     {
-      orderBy: [{ expr: parseExpressionString("x"), desc: false }],
+      objectVariable: "p",
+      orderBy: [{ expr: parseExpressionString("p.x"), desc: false }],
     },
     rootEnv,
     LuaStackFrame.lostFrame,
@@ -68,7 +72,8 @@ Deno.test("ArrayQueryCollection", async () => {
   // Test order by desc
   const result5 = await collection.query(
     {
-      orderBy: [{ expr: parseExpressionString("x"), desc: true }],
+      objectVariable: "p",
+      orderBy: [{ expr: parseExpressionString("p.x"), desc: true }],
     },
     rootEnv,
     LuaStackFrame.lostFrame,
@@ -87,9 +92,10 @@ Deno.test("ArrayQueryCollection", async () => {
   ]);
   const result6 = await collection2.query(
     {
+      objectVariable: "p",
       orderBy: [
-        { expr: parseExpressionString("lastName"), desc: false },
-        { expr: parseExpressionString("firstName"), desc: true },
+        { expr: parseExpressionString("p.lastName"), desc: false },
+        { expr: parseExpressionString("p.firstName"), desc: true },
       ],
     },
     rootEnv,
@@ -104,23 +110,13 @@ Deno.test("ArrayQueryCollection", async () => {
   assertEquals(result6[3].firstName, "Alice");
   assertEquals(result6[3].lastName, "Johnson");
 
-  // Test select
-  const result7 = await collection2.query(
-    {
-      select: [{ name: "firstName" }],
-    },
-    rootEnv,
-    LuaStackFrame.lostFrame,
-  );
-  assertEquals(result7[0].firstName, "John");
-  assertEquals(result7[0].lastName, undefined);
-
   // Test select with expression
   const result8 = await collection2.query(
     {
+      objectVariable: "p",
       select: [{
         name: "fullName",
-        expr: parseExpressionString("firstName .. ' ' .. lastName"),
+        expr: parseExpressionString("p.firstName .. ' ' .. p.lastName"),
       }],
     },
     rootEnv,
@@ -134,9 +130,10 @@ Deno.test("ArrayQueryCollection", async () => {
   // Test select with native function
   const result9 = await collection2.query(
     {
+      objectVariable: "p",
       select: [{
         name: "fullName",
-        expr: parseExpressionString("build_name(firstName, lastName)"),
+        expr: parseExpressionString("build_name(p.firstName, p.lastName)"),
       }],
     },
     rootEnv,

common/space_lua/query_collection.ts

@@ -1,18 +1,12 @@
 import type { LuaExpression } from "$common/space_lua/ast.ts";
-import {
-  LuaEnv,
-  luaGet,
-  luaKeys,
-  type LuaStackFrame,
-} from "$common/space_lua/runtime.ts";
+import { LuaEnv, type LuaStackFrame } from "$common/space_lua/runtime.ts";
 import { evalExpression } from "$common/space_lua/eval.ts";
 import { asyncQuickSort } from "$common/space_lua/util.ts";
+import type { DataStore } from "$lib/data/datastore.ts";
 
-function buildItemEnv(item: any, env: LuaEnv, sf: LuaStackFrame): LuaEnv {
+function buildItemEnv(objectVariable: string, item: any, env: LuaEnv): LuaEnv {
   const itemEnv = new LuaEnv(env);
-  for (const key of luaKeys(item)) {
-    itemEnv.setLocal(key, luaGet(item, key, sf));
-  }
+  itemEnv.setLocal(objectVariable, item);
   return itemEnv;
 }
 
@@ -30,6 +24,7 @@ export type LuaSelect = {
  * Represents a query for a collection
  */
 export type LuaCollectionQuery = {
+  objectVariable: string;
   // The filter expression evaluated with Lua
   where?: LuaExpression;
   // The order by expression evaluated with Lua
@@ -61,71 +56,107 @@ export class ArrayQueryCollection<T> implements LuaQueryCollection {
     env: LuaEnv,
     sf: LuaStackFrame,
   ): Promise<any[]> {
-    let result: any[] = [];
+    const result: any[] = [];
 
     // Filter the array
     for (const item of this.array) {
-      const itemEnv = buildItemEnv(item, env, sf);
+      const itemEnv = buildItemEnv(query.objectVariable, item, env);
       if (query.where && !await evalExpression(query.where, itemEnv, sf)) {
         continue;
       }
       result.push(item);
     }
 
-    // Apply the select
-    if (query.select) {
-      const newResult = [];
-      for (const item of result) {
-        const itemEnv = buildItemEnv(item, env, sf);
-        const newItem: Record<string, any> = {};
-        for (const select of query.select) {
-          if (select.expr) {
-            newItem[select.name] = await evalExpression(
-              select.expr,
-              itemEnv,
-              sf,
-            );
-          } else {
-            newItem[select.name] = item[select.name];
-          }
-        }
-        newResult.push(newItem);
-      }
-      result = newResult;
-    }
-
-    // Apply the order by
-    if (query.orderBy) {
-      result = await asyncQuickSort(result, async (a, b) => {
-        // Compare each orderBy clause until we find a difference
-        for (const { expr, desc } of query.orderBy!) {
-          const aEnv = buildItemEnv(a, env, sf);
-          const bEnv = buildItemEnv(b, env, sf);
-
-          const aVal = await evalExpression(expr, aEnv, sf);
-          const bVal = await evalExpression(expr, bEnv, sf);
-
-          if (aVal < bVal) {
-            return desc ? 1 : -1;
-          }
-          if (aVal > bVal) {
-            return desc ? -1 : 1;
-          }
-          // If equal, continue to next orderBy clause
-        }
-        return 0; // All orderBy clauses were equal
-      });
-    }
-
-    // Apply the limit and offset
-    if (query.limit !== undefined && query.offset !== undefined) {
-      result = result.slice(query.offset, query.offset + query.limit);
-    } else if (query.limit !== undefined) {
-      result = result.slice(0, query.limit);
-    } else if (query.offset !== undefined) {
-      result = result.slice(query.offset);
-    }
-
-    return Promise.resolve(result);
+    return applyTransforms(result, query, env, sf);
+  }
+}
+
+async function applyTransforms(
+  result: any[],
+  query: LuaCollectionQuery,
+  env: LuaEnv,
+  sf: LuaStackFrame,
+): Promise<any[]> {
+  // Apply the select
+  if (query.select) {
+    const newResult = [];
+    for (const item of result) {
+      const itemEnv = buildItemEnv(query.objectVariable, item, env);
+      const newItem: Record<string, any> = {};
+      for (const select of query.select) {
+        if (select.expr) {
+          newItem[select.name] = await evalExpression(
+            select.expr,
+            itemEnv,
+            sf,
+          );
+        } else {
+          newItem[select.name] = item[select.name];
+        }
+      }
+      newResult.push(newItem);
+    }
+    result = newResult;
+  }
+
+  // Apply the order by
+  if (query.orderBy) {
+    result = await asyncQuickSort(result, async (a, b) => {
+      // Compare each orderBy clause until we find a difference
+      for (const { expr, desc } of query.orderBy!) {
+        const aEnv = buildItemEnv(query.objectVariable, a, env);
+        const bEnv = buildItemEnv(query.objectVariable, b, env);
+
+        const aVal = await evalExpression(expr, aEnv, sf);
+        const bVal = await evalExpression(expr, bEnv, sf);
+
+        if (aVal < bVal) {
+          return desc ? 1 : -1;
+        }
+        if (aVal > bVal) {
+          return desc ? -1 : 1;
+        }
+        // If equal, continue to next orderBy clause
+      }
+      return 0; // All orderBy clauses were equal
+    });
+  }
+
+  // Apply the limit and offset
+  if (query.limit !== undefined && query.offset !== undefined) {
+    result = result.slice(query.offset, query.offset + query.limit);
+  } else if (query.limit !== undefined) {
+    result = result.slice(0, query.limit);
+  } else if (query.offset !== undefined) {
+    result = result.slice(query.offset);
+  }
+
+  return result;
+}
+
+export class DataStoreQueryCollection implements LuaQueryCollection {
+  constructor(
+    private readonly dataStore: DataStore,
+    readonly prefix: string[],
+  ) {}
+
+  async query(
+    query: LuaCollectionQuery,
+    env: LuaEnv,
+    sf: LuaStackFrame,
+  ): Promise<any[]> {
+    const result: any[] = [];
+    for await (
+      const { value } of this.dataStore.kv.query({ prefix: this.prefix })
+    ) {
+      // Enrich
+      this.dataStore.enrichObject(value);
+      const itemEnv = buildItemEnv(query.objectVariable, value, env);
+      if (query.where && !await evalExpression(query.where, itemEnv, sf)) {
+        continue;
+      }
+      result.push(value);
+    }
+    return applyTransforms(result, query, env, sf);
   }
 }

common/space_lua/stdlib.ts

@@ -16,9 +16,13 @@ import { tableApi } from "$common/space_lua/stdlib/table.ts";
 import { osApi } from "$common/space_lua/stdlib/os.ts";
 import { jsApi } from "$common/space_lua/stdlib/js.ts";
 import { spaceLuaApi } from "$common/space_lua/stdlib/space_lua.ts";
+import type {
+  LuaCollectionQuery,
+  LuaQueryCollection,
+} from "$common/space_lua/query_collection.ts";
 
 const printFunction = new LuaBuiltinFunction(async (_sf, ...args) => {
-  console.log("[Lua]", ...(await Promise.all(args.map(luaToString))));
+  console.log("[Lua]", ...(await Promise.all(args)));
 });
 
 const assertFunction = new LuaBuiltinFunction(
@@ -123,6 +127,27 @@ const getmetatableFunction = new LuaBuiltinFunction((_sf, table: LuaTable) => {
   return table.metatable;
 });
 
+const tagFunction = new LuaBuiltinFunction(
+  (sf, tagName: LuaValue): LuaQueryCollection => {
+    const global = sf.threadLocal.get("_GLOBAL");
+    if (!global) {
+      throw new LuaRuntimeError("Global not found", sf);
+    }
+    return {
+      query: async (query: LuaCollectionQuery): Promise<any[]> => {
+        return (await global.get("datastore").get("query_lua").call(
+          sf,
+          [
+            "idx",
+            tagName,
+          ],
+          query,
+        )).asJSArray();
+      },
+    };
+  },
+);
+
 export function luaBuildStandardEnv() {
   const env = new LuaEnv();
   // Top-level builtins
@@ -143,7 +168,7 @@ export function luaBuildStandardEnv() {
   env.set("error", errorFunction);
   env.set("pcall", pcallFunction);
   env.set("xpcall", xpcallFunction);
-
+  env.set("tag", tagFunction);
   // APIs
   env.set("string", stringApi);
   env.set("table", tableApi);

common/space_lua/stdlib/table.ts

@@ -37,4 +37,7 @@ export const tableApi = new LuaTable({
   sort: new LuaBuiltinFunction((sf, tbl: LuaTable, comp?: ILuaFunction) => {
     return tbl.sort(comp, sf);
   }),
+  keys: new LuaBuiltinFunction((_sf, tbl: LuaTable) => {
+    return tbl.keys();
+  }),
 });

lib/plugos/syscalls/datastore.ts

@@ -1,13 +1,22 @@
+import {
+  DataStoreQueryCollection,
+  type LuaCollectionQuery,
+} from "$common/space_lua/query_collection.ts";
+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 } from "$common/space_lua/runtime.ts";
 
 /**
  * Exposes the datastore API to plugs, but scoping everything to a prefix based on the plug's name
  * @param ds the datastore to wrap
  * @param prefix prefix to scope all keys to to which the plug name will be appended
  */
-export function dataStoreReadSyscalls(ds: DataStore): SysCallMapping {
+export function dataStoreReadSyscalls(
+  ds: DataStore,
+  commonSystem: CommonSystem,
+): SysCallMapping {
   return {
     "datastore.batchGet": (
       _ctx,
@@ -28,6 +37,19 @@ export function dataStoreReadSyscalls(ds: DataStore): SysCallMapping {
       return ds.query(query, variables);
     },
 
+    "datastore.queryLua": (
+      _ctx,
+      prefix: string[],
+      query: LuaCollectionQuery,
+    ): Promise<KV[]> => {
+      const dsQueryCollection = new DataStoreQueryCollection(ds, prefix);
+      return dsQueryCollection.query(
+        query,
+        commonSystem.spaceLuaEnv.env,
+        LuaStackFrame.lostFrame,
+      );
+    },
+
     "datastore.listFunctions": (): string[] => {
       return Object.keys(ds.functionMap);
     },

plug-api/lib/tree.ts

@@ -215,8 +215,8 @@ export function cloneTree(tree: ParseTree): ParseTree {
 }
 
 export function parseTreeToAST(tree: ParseTree, omitTrimmable = true): AST {
-  const parseErrorNodes = collectNodesOfType(tree, "⚠");
-  if (parseErrorNodes.length > 0) {
+  if (tree.type === "⚠") {
+    console.info("Parse error", JSON.stringify(tree, null, 2));
     throw new Error(
       `Parse error in: ${renderToText(tree)}`,
     );
@@ -237,12 +237,10 @@ export function parseTreeToAST(tree: ParseTree, omitTrimmable = true): AST {
 }
 
 export function cleanTree(tree: ParseTree, omitTrimmable = true): ParseTree {
-  const parseErrorNodes = collectNodesOfType(tree, "⚠");
-  if (parseErrorNodes.length > 0) {
+  if (tree.type === "⚠") {
+    console.info("Parse error", JSON.stringify(tree, null, 2));
     throw new Error(
-      `Parse error (${parseErrorNodes[0].from}:${parseErrorNodes[0].to}): ${
-        renderToText(tree)
-      }`,
+      `Parse error in: ${renderToText(tree)}`,
     );
   }
   if (tree.text !== undefined) {

plug-api/syscalls/datastore.ts

@@ -1,3 +1,4 @@
+import type { LuaCollectionQuery } from "$common/space_lua/query_collection.ts";
 import { syscall } from "../syscall.ts";
 import type { KV, KvKey, KvQuery } from "../types.ts";
 
@@ -70,6 +71,13 @@ export function query(
   return syscall("datastore.query", query, variables);
 }
 
+export function queryLua(
+  prefix: string[],
+  query: LuaCollectionQuery,
+): Promise<KV[]> {
+  return syscall("datastore.queryLua", prefix, query);
+}
+
 /**
  * Queries the key value store and deletes all matching items
  * @param query the query to run

server/server_system.ts

@@ -135,7 +135,7 @@ export class ServerSystem extends CommonSystem {
       jsonschemaSyscalls(),
       luaSyscalls(),
       templateSyscalls(this.ds),
-      dataStoreReadSyscalls(this.ds),
+      dataStoreReadSyscalls(this.ds, this),
       codeWidgetSyscalls(codeWidgetHook),
       markdownSyscalls(),
     );

web/client_system.ts

@@ -169,7 +169,10 @@ export class ClientSystem extends CommonSystem {
         // In non-sync mode proxy to server
         : mqProxySyscalls(this.client),
       ...this.client.syncMode
-        ? [dataStoreReadSyscalls(this.ds), dataStoreWriteSyscalls(this.ds)]
+        ? [
+          dataStoreReadSyscalls(this.ds, this),
+          dataStoreWriteSyscalls(this.ds),
+        ]
         : [dataStoreProxySyscalls(this.client)],
       debugSyscalls(this.client),
       syncSyscalls(this.client),

web/syscalls/datastore.proxy.ts

@@ -10,6 +10,7 @@ export function dataStoreProxySyscalls(client: Client): SysCallMapping {
     "datastore.batchDelete",
     "datastore.batchGet",
     "datastore.query",
+    "datastore.queryLua",
     "datastore.get",
     "datastore.listFunctions",
   ]);