4MB memory limit on QuickJS execution, closes #4

README.md

@@ -17,7 +17,7 @@ datasette install datasette-enrichments-quickjs
 
 This enrichment allows you to select rows from a table and specify a custom JavaScript function to use to generate a value for each of those rows, storing that value in a specified column and creating that column if it does not exist.
 
-Code runs in a [QuickJS sandbox](https://github.com/PetterS/quickjs) with a 0.1s time limit for the execution of each function.
+Code runs in a [QuickJS sandbox](https://github.com/PetterS/quickjs) with a 0.1s time limit for the execution of each function and a 4MB memory limit.
 
 Enrichment JavaScript functions look like this:
 

datasette_enrichments_quickjs/__init__.py

@@ -81,9 +81,14 @@ async def enrich_batch(
     ):
         function = Function("enrich", config["javascript"])
         function.set_time_limit(0.1)  # 0.1s
+        function.set_memory_limit(4 * 1024 * 1024)  # 4MB
         output_column = config["output_column"]
         for row in rows:
-            output = function(row)
+            try:
+                output = function(row)
+            except Exception as ex:
+                print(ex, repr(ex))
+                raise
             await db.execute_write(
                 "update [{table}] set [{output_column}] = ? where {wheres}".format(
                     table=table,

pyproject.toml

@@ -27,7 +27,7 @@ CI = "https://github.com/datasette/datasette-enrichments-quickjs/actions"
 enrichments_quickjs = "datasette_enrichments_quickjs"
 
 [project.optional-dependencies]
-test = ["pytest", "pytest-asyncio"]
+test = ["pytest", "pytest-asyncio", "pytest-timeout"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "strict"

tests/test_enrichments_quickjs.py

@@ -4,11 +4,11 @@
 import sqlite_utils
 
 
-async def _cookies(datasette):
+async def _cookies(datasette, path="/-/enrich/data/items/quickjs"):
     cookies = {"ds_actor": datasette.sign({"a": {"id": "root"}}, "actor")}
-    csrftoken = (
-        await datasette.client.get("/-/enrich/data/items/quickjs", cookies=cookies)
-    ).cookies["ds_csrftoken"]
+    csrftoken = (await datasette.client.get(path, cookies=cookies)).cookies[
+        "ds_csrftoken"
+    ]
     cookies["ds_csrftoken"] = csrftoken
     return cookies
 
@@ -62,3 +62,87 @@ async def test_enrichment(tmpdir):
             "description_length": 10,
         },
     ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(5)
+@pytest.mark.parametrize(
+    "javascript,expected_error",
+    [
+        # Deeply nested object should run out of memory
+        (
+            """
+    function enrich() {
+        let obj = {};
+        for (let i = 0; i < 100000; i++) {
+            obj = {nested: obj};
+        }
+        return obj;
+    }""",
+            "null\n",
+        ),
+        # Long running operation should return interrupted error
+        (
+            """
+    function enrich() {
+        let start = Date.now();
+        while (Date.now() - start < 500);
+        return 'Task completed';
+    }""",
+            "InternalError: interrupted\n    at enrich (<input>:3)\n",
+        ),
+        # Should work
+        (
+            """
+    function enrich() {
+        return 1;
+    }""",
+            None,
+        ),
+    ],
+)
+async def test_time_and_memory_limit(javascript, expected_error):
+    ds = Datasette()
+    db = ds.add_memory_database("test_time_and_memory_limit")
+    await db.execute_write("create table if not exists foo (id integer primary key)")
+    await db.execute_write("insert or replace into foo (id) values (1)")
+    try:
+        # In-memory DB persists between runs, so clear it
+        await db.execute_write("delete from _enrichment_jobs")
+        await db.execute_write("delete from _enrichment_errors")
+    except Exception:
+        # Table does not exist yet
+        pass
+    cookies = await _cookies(
+        ds, path="/-/enrich/test_time_and_memory_limit/foo/quickjs"
+    )
+
+    rows = [
+        {"id": 1, "name": "One", "description": "First item"},
+        {"id": 2, "name": "Two", "description": "Second item"},
+        {"id": 3, "name": "Three", "description": "Third item"},
+    ]
+    post = {
+        "javascript": javascript,
+        "output_column": "description_length",
+        "output_column_type": "integer",
+    }
+    post["csrftoken"] = cookies["ds_csrftoken"]
+    response = await ds.client.post(
+        "/-/enrich/test_time_and_memory_limit/foo/quickjs",
+        data=post,
+        cookies=cookies,
+    )
+    assert response.status_code == 302
+    await asyncio.sleep(0.3)
+    jobs = await db.execute("select * from _enrichment_jobs")
+    row = dict(jobs.rows[0])
+    assert row["status"] == "finished"
+
+    if expected_error:
+        assert row["error_count"] == 1
+        errors = (await db.execute("select * from _enrichment_errors")).rows
+        error = dict(errors[0])["error"]
+        assert error == expected_error
+    else:
+        assert row["error_count"] == 0