Repository: geldata/gel
Pull request: add `std::json_set` function
Closes #3826 

Problem statement:
Update parts of JSON properties with a json_set function
Let's say initially `myMovie.json = {"a": 10, "b": 20}`

Instead of `update Movie set {myJson := to_json('{"a": 11, "b": 20}')};`, make it possible to do something like: `update Movie set {myJson := json_set(.myJson, 'a', .myJson['a'] + 1)}`.

It could just be syntactic sugar for cloning the whole JSON object and updating a deep field.

The syntax could be similar to `json_get` where the first argument specifies the JSON property, then the middle arguments are a path to a location within the JSON object to update, and the last argument is the value to set.

What if the path is invalid, e.g. `json_set(.myJson, 'a', 'x', 'y', 11)`? Maybe create singleton objects along the way to make it valid like `{"a": {"x": {"y": 11}}}`.

If there's an array along the path, could use numbers to specify indexes, and error if the path part isn't a number.

The command can potentially have a merge flag, so if `myMovie.myJson` initially equals `{"a": {"x": 5, "y": 6}}`, then `json_set(.myJson, 'a', to_json('{"y": 7, "z": 8}'))` could return `{"a": {"x": 5, "y": 7, "z": 8}}`.

Hints from issue discussion:
I propose the following signature:

```
json_set(
  target: json,
  path: array<str>,  # We can't use variadic here, because new_value follows.
  new_value: optional json,
  named only create_if_missing: bool = false,
  named only empty_treatment: JsonSetEmptyTreatment = JsonSetEmptyTreatment.ReturnEmpty,
) -> optional json
```

Where `create_if_missing=true` causes intermediate elements to be created if missing, and `empty_treatment` is an enum determining what to do if *new_value* is an empty set:

```
scalar type JsonSetEmptyTreatment extending enum<
  ReturnEmpty,     # Return empty set, default 
  ReturnTarget,    # Return *target* unmodified
  RaiseException,  # Raise a ValueError
  UseJsonNull,     # Use a 'null` JSON value
  DeleteKey,       # Delete the object key
>;
```

This is basically the equivalent of PostgreSQL's `jsonb_set` and `jsonb_set_lax` functions.
Looks about right.

`path` probably needs to be something like `array<str|int32>` to index into JSON arrays along the path.

I guess you're thinking `ReturnEmpty` would cause json_set to always be a no-op, i.e. doesn't make sense to call it this way in a command by itself, but sometimes useful as part of a larger expression?

Maybe the idea about being able to merge objects together at the end of the path, rather than replacing the whole leaf object, could be yet another named enum argument. Or a separate function `json_update` instead of `json_set`.
> `path` probably needs to be something like `array<str|int32>` to index into JSON arrays along the path.

Not really, in JavaScript `foo["1"]` and `foo[1]` are equivalent, so it's safe to use strings to index into arrays (`json_get` uses a `variadic` `str`).

> Maybe the idea about being able to merge objects together at the end of the path

This is where you should probably just use `++` instead.

Commit message: rename enum for function behaviour for empty json

Patch to review:
diff --git a/docs/stdlib/json.rst b/docs/stdlib/json.rst
index 77325ebc123..2443b505b9c 100644
--- a/docs/stdlib/json.rst
+++ b/docs/stdlib/json.rst
@@ -38,6 +38,9 @@ JSON
     * - :eql:func:`json_get`
       - :eql:func-desc:`json_get`
 
+    * - :eql:func:`json_set`
+      - :eql:func-desc:`json_set`
+
     * - :eql:func:`json_array_unpack`
       - :eql:func-desc:`json_array_unpack`
 
@@ -144,7 +147,7 @@ possible to cast a JSON value directly into a :eql:type:`tuple`.
     .. code-block:: edgeql-repl
 
         db> select <str>to_json('["a", "b", "c"]');
-        InternalServerError: expected json string, null; got json array
+        InvalidValueError: expected json string or null; got JSON array
 
     Use the :eql:func:`to_json` and :eql:func:`to_str`
     functions to dump or parse a :eql:type:`json` value to or
@@ -353,6 +356,105 @@ possible to cast a JSON value directly into a :eql:type:`tuple`.
 ----------
 
 
+.. eql:function:: std::json_set( \
+                    target: json, \
+                    variadic path: str, \
+                    named only value: optional json, \
+                    named only create_if_missing: bool = true, \
+                    named only empty_treatment: JsonEmpty = \
+                      JsonEmpty.ReturnEmpty) \
+                  -> optional json
+
+    Return an updated JSON target with a new value.
+
+    .. warning::
+
+      This function is only available in EdgeDB 2.0 or later.
+
+    .. code-block:: edgeql-repl
+
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>true,
+        ... );
+        {'{"a": true, "b": 20}'}
+        db> select json_set(
+        ...   to_json('{"a": {"b": {}}}'),
+        ...   'a', 'b', 'c',
+        ...   value := <json>42,
+        ... );
+        {'{"a": {"b": {"c": 42}}}'}
+
+    If ``create_if_missing`` is set to ``false``, a new path for the value
+    won't be created.
+
+    .. code-block:: edgeql-repl
+
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'с',
+        ...   value := <json>42,
+        ... );
+        {'{"a": 10, "b": 20, "c": 42}'}
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'с',
+        ...   value := <json>42,
+        ...   create_if_missing := false,
+        ... );
+        {'{"a": 10, "b": 20}'}
+
+    ``empty_treatment`` is an enumeration responsible for the behavior of the
+    function if an empty set is passed to ``new_value``. It contains one of
+    the following values:
+
+    - ``ReturnEmpty``: return empty set, default
+    - ``ReturnTarget``: return ``target`` unmodified
+    - ``Error``: raise an ``InvalidValueError``
+    - ``UseNull``: use a ``null`` JSON value
+    - ``DeleteKey``: delete the object key
+
+    .. code-block:: edgeql-repl
+
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>{}
+        ... );
+        {}
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>{},
+        ...   empty_treatment := JsonEmpty.ReturnTarget,
+        ... );
+        {'{"a": 10, "b": 20}'}
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>{},
+        ...   empty_treatment := JsonEmpty.Error,
+        ... );
+        InvalidValueError: invalid empty JSON value
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>{},
+        ...   empty_treatment := JsonEmpty.UseNull,
+        ... );
+        {'{"a": null, "b": 20}'}
+        db> select json_set(
+        ...   to_json('{"a": 10, "b": 20}'),
+        ...   'a',
+        ...   value := <json>{},
+        ...   empty_treatment := JsonEmpty.DeleteKey,
+        ... );
+        {'{"b": 20}'}
+
+----------
+
+
 .. eql:function:: std::json_object_unpack(json: json) -> \
                   set of tuple<str, json>
 
diff --git a/edb/buildmeta.py b/edb/buildmeta.py
index 5257c6538c0..914ba3d72b5 100644
--- a/edb/buildmeta.py
+++ b/edb/buildmeta.py
@@ -44,7 +44,7 @@
 
 
 # Increment this whenever the database layout or stdlib changes.
-EDGEDB_CATALOG_VERSION = 2022_07_22_00_00
+EDGEDB_CATALOG_VERSION = 2022_07_25_00_00
 EDGEDB_MAJOR_VERSION = 3
 
 
diff --git a/edb/lib/std/30-jsonfuncs.edgeql b/edb/lib/std/30-jsonfuncs.edgeql
index 502c35880bf..20f46875f44 100644
--- a/edb/lib/std/30-jsonfuncs.edgeql
+++ b/edb/lib/std/30-jsonfuncs.edgeql
@@ -19,6 +19,7 @@
 
 ## JSON functions and operators.
 
+CREATE SCALAR TYPE std::JsonEmpty EXTENDING enum<ReturnEmpty, ReturnTarget, Error, UseNull, DeleteKey>;
 
 CREATE FUNCTION
 std::json_typeof(json: std::json) -> std::str
@@ -74,6 +75,42 @@ std::json_get(
     $$;
 };
 
+CREATE FUNCTION
+std::json_set(
+    target: std::json,
+    VARIADIC path: std::str,
+    NAMED ONLY value: OPTIONAL std::json,
+    NAMED ONLY create_if_missing: std::bool = true,
+    NAMED ONLY empty_treatment: std::JsonEmpty = std::JsonEmpty.ReturnEmpty,
+) -> OPTIONAL std::json
+{
+    CREATE ANNOTATION std::description :=
+        'Return an updated JSON target with a new value.';
+    SET volatility := 'Immutable';
+    USING SQL $$
+    SELECT (
+        CASE
+        WHEN "value" IS NULL AND "empty_treatment" = 'ReturnEmpty' THEN
+            NULL
+        WHEN "value" IS NULL AND "empty_treatment" = 'ReturnTarget' THEN
+            "target"
+        WHEN "value" IS NULL AND "empty_treatment" = 'Error' THEN
+            edgedb.raise(
+                NULL::jsonb,
+                'invalid_parameter_value',
+                msg => 'invalid empty JSON value'
+            )
+        WHEN "value" IS NULL AND "empty_treatment" = 'UseNull' THEN
+            jsonb_set("target", "path", 'null'::jsonb, "create_if_missing")
+        WHEN "value" IS NULL AND "empty_treatment" = 'DeleteKey' THEN
+            "target" #- "path"
+        ELSE
+            jsonb_set("target", "path", "value", "create_if_missing")
+        END
+    )
+    $$;
+};
+
 CREATE INFIX OPERATOR
 std::`=` (l: std::json, r: std::json) -> std::bool {
     CREATE ANNOTATION std::identifier := 'eq';
diff --git a/tests/test_edgeql_json.py b/tests/test_edgeql_json.py
index 204d3f09d2f..b773fa91a1d 100644
--- a/tests/test_edgeql_json.py
+++ b/tests/test_edgeql_json.py
@@ -1094,6 +1094,132 @@ async def test_edgeql_json_get_04(self):
             ['"42!"']
         )
 
+    async def test_edgeql_json_set_01(self):
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(JT0.j_object, 'a', value := <json>42);
+            ''',
+            [{"a": 42, "b": 2}],
+            ['{"a": 42, "b": 2}'],
+        )
+
+        # by default create_if_missing should create new value
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(JT0.j_object, 'c', value := <json>42);
+            ''',
+            [{"a": 1, "b": 2, "c": 42}],
+            ['{"a": 1, "b": 2, "c": 42}'],
+        )
+
+        # create_if_missing should also work inside of nested objects
+        # by variadic path
+        await self.assert_query_result(
+            r'''
+                WITH j_object := to_json('{"a": {"b": {}}}')
+                SELECT json_set(j_object, 'a', 'b', 'c', value := <json>42);
+            ''',
+            [{"a": {"b": {"c": 42}}}],
+            ['{"a": {"b": {"c": 42}}}'],
+        )
+
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(
+                    JT0.j_object,
+                    'c',
+                    value := <json>42,
+                    create_if_missing := false,
+                );
+            ''',
+            [{"a": 1, "b": 2}],
+            ['{"a": 1, "b": 2}'],
+        )
+
+        # by default empty_treatment should return empty set
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(JT0.j_object, 'a', value := <json>{});
+            ''',
+            [],
+            [],
+        )
+
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(
+                    JT0.j_object,
+                    'a',
+                    value := <json>{},
+                    empty_treatment := JsonEmpty.ReturnEmpty,
+                );
+            ''',
+            [],
+            [],
+        )
+
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(
+                    JT0.j_object,
+                    'a',
+                    value := <json>{},
+                    empty_treatment := JsonEmpty.ReturnTarget,
+                );
+            ''',
+            [{"a": 1, "b": 2}],
+            ['{"a": 1, "b": 2}'],
+        )
+
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(
+                    JT0.j_object,
+                    'a',
+                    value := <json>{},
+                    empty_treatment := JsonEmpty.UseNull,
+                );
+            ''',
+            [{"a": None, "b": 2}],
+            ['{"a": null, "b": 2}'],
+        )
+
+        await self.assert_query_result(
+            r'''
+                WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                SELECT json_set(
+                    JT0.j_object,
+                    'a',
+                    value := <json>{},
+                    empty_treatment := JsonEmpty.DeleteKey,
+                );
+            ''',
+            [{"b": 2}],
+            ['{"b": 2}'],
+        )
+
+        with self.assertRaisesRegex(
+                edgedb.InvalidValueError,
+                r"invalid empty JSON value"):
+            await self.con.execute(
+                r'''
+                    WITH JT0 := (SELECT JSONTest FILTER .number = 0)
+                    SELECT json_set(
+                        JT0.j_object,
+                        'a',
+                        value := <json>{},
+                        empty_treatment := JsonEmpty.Error,
+                    );
+                ''',
+            )
+
     async def test_edgeql_json_cast_object_to_json_01(self):
         res = await self.con.query("""
             WITH MODULE schema


Code review comments:
diff_hunk: @@ -44,7 +44,7 @@
 
 
 # Increment this whenever the database layout or stdlib changes.
-EDGEDB_CATALOG_VERSION = 2022_07_22_00_00
+EDGEDB_CATALOG_VERSION = 2022_07_25_00_00
path: edb/buildmeta.py
text: @user1:
```suggestion
EDGEDB_CATALOG_VERSION = 2022_07_26_00_00
```