Repository: edgedb/edgedb

Issue:
InternalServerError: cannot cast type uuid to jsonb
EdgeDB Version: 3.2+75dd258
EdgeDB CLI Version: EdgeDB CLI 3.4.0+97cad0e
OS Version: Mac OS 13.0.1 (22A400)

Steps to Reproduce:

1. edgedb migration create succeeds
2. edgedb migrate fails with "edb.errors.InternalServerError: cannot cast type uuid to jsonb"

```
type User { 
    email:str;
}
function getUser(idx:str) -> set of json {
    using (
        select <json> (
            select User {
                email
            } filter .id = <uuid>idx
        )
    )
};
```


Discussion:

Resolving patch:
diff --git a/edb/pgsql/compiler/context.py b/edb/pgsql/compiler/context.py
--- a/edb/pgsql/compiler/context.py
+++ b/edb/pgsql/compiler/context.py
@@ -536,8 +536,11 @@ def output_format(
     output_format: OutputFormat,
 ) -> Generator[None, None, None]:
     original_output_format = ctx.env.output_format
+    original_ignore_object_shapes = ctx.env.ignore_object_shapes
     ctx.env.output_format = output_format
+    ctx.env.ignore_object_shapes = False
     try:
         yield
     finally:
         ctx.env.output_format = original_output_format
+        ctx.env.ignore_object_shapes = original_ignore_object_shapes
diff --git a/edb/pgsql/compiler/relgen.py b/edb/pgsql/compiler/relgen.py
--- a/edb/pgsql/compiler/relgen.py
+++ b/edb/pgsql/compiler/relgen.py
@@ -2110,24 +2110,24 @@ def process_set_as_type_cast(
     inner_set = expr.expr
     is_json_cast = expr.to_type.id == s_obj.get_known_type_id('std::json')
 
-    with ctx.new() as subctx:
+    # Are we casting by compiling the innards in json mode?
+    implicit_cast = (
+        is_json_cast
+        and not irtyputils.is_range(inner_set.typeref)
+        and (irtyputils.is_collection(inner_set.typeref)
+             or irtyputils.is_object(inner_set.typeref))
+    )
+    fmt_ctx = (
+        context.output_format(ctx, context.OutputFormat.JSONB) if implicit_cast
+        else contextlib.nullcontext()
+    )
+
+    with fmt_ctx, ctx.new() as subctx:
         pathctx.put_path_id_map(ctx.rel, ir_set.path_id, inner_set.path_id)
 
-        if (is_json_cast
-                and not irtyputils.is_range(inner_set.typeref)
-                and (irtyputils.is_collection(inner_set.typeref)
-                     or irtyputils.is_object(inner_set.typeref))):
+        if implicit_cast:
             subctx.expr_exposed = True
-            # XXX: this is necessary until pathctx is converted
-            #      to use context levels instead of using env
-            #      directly.
-            orig_output_format = subctx.env.output_format
-            subctx.env.output_format = context.OutputFormat.JSONB
-            implicit_cast = True
-        else:
-            implicit_cast = False
 
-        if implicit_cast:
             set_expr = dispatch.compile(inner_set, ctx=subctx)
 
             serialized: Optional[pgast.BaseExpr] = (
@@ -2148,9 +2148,8 @@ def process_set_as_type_cast(
                 pathctx.put_path_serialized_var(
                     stmt, inner_set.path_id, serialized, force=True
                 )
-
-            subctx.env.output_format = orig_output_format
         else:
+            # Rely on the simple implementation of TypeCast
             set_expr = dispatch.compile(expr, ctx=subctx)
 
             # A proper path var mapping way would be to wrap


Test patch:
diff --git a/tests/test_edgeql_json.py b/tests/test_edgeql_json.py
--- a/tests/test_edgeql_json.py
+++ b/tests/test_edgeql_json.py
@@ -1437,6 +1437,28 @@ async def test_edgeql_json_cast_object_to_json_06(self):
 
         self.assertTrue('id' not in data)
 
+    async def test_edgeql_json_cast_object_to_json_07(self):
+        await self.con.execute('''
+            create function _get(idx: int64) -> set of json {
+                using (
+                    select <json> (
+                        select JSONTest {
+                            number, edb_string
+                        } filter .number = idx
+                    )
+                )
+            };
+        ''')
+        await self.assert_query_result(
+            r"""
+                SELECT _get(0)
+            """,
+            [
+                {'number': 0, 'edb_string': 'jumps'},
+            ],
+            json_only=True,
+        )
+
     async def test_edgeql_json_cast_tuple_to_json_01(self):
         res = await self.con.query("""
             WITH MODULE schema
diff --git a/tests/test_edgeql_volatility.py b/tests/test_edgeql_volatility.py
--- a/tests/test_edgeql_volatility.py
+++ b/tests/test_edgeql_volatility.py
@@ -1663,6 +1663,15 @@ async def test_edgeql_volatility_shape_array_02(self):
             for obj in res:
                 self.assertEqual(obj[0], obj[1])
 
+    async def test_edgeql_volatility_in_func_01(self):
+        await self.con.execute('''
+            create function foo() -> float64 using (
+              with Z := Obj { v := random() },
+              select sum(Z.v)
+            );
+        ''')
+        await self.con.query('select foo()')
+
     async def test_edgeql_volatility_errors_01(self):
         async with self._run_and_rollback():
             with self.assertRaisesRegex(