Fix updating of two single links with linkprops at once (#8932)

Mutating a single link with link properties works by compiling the
actual expression that defines the value into the link table update,
and then copying that value into the table itself by (ab/re)using the
rewrite machinery.

The rewrite machinery uses a "policy_ctx", which uses one of the input
tables as part of the overlay instead of an output table. The
"policy_ctx" has its overlay tables copied before being updated, but
it does the copy "just in time", which means that you can sometimes
get away with depending on things added to the context *after* the
policy_ctx was created.

Of course, this only worked *once*, after which the copy had been
made. As a result there was a bug where an 'except' overlay wasn't
being applied when populating a single link.

Be more explicit about all this by populating *both* overlay tables.

Also, it turns out that we do depend on incorporating unrelated
changes made to the overlay tables during compilation (like if an
INSERT is happening nested inside other DML), so we need to explicitly
reapply the policy overlays to the main one. (There was an existing bug
caused by this that this change made appear more easily.)

Possibly also fix some other policy related issues, since the code is
shared, but I haven't tracked it down too much.

Author: msullivan

edb/pgsql/compiler/dml.py

@@ -810,7 +810,7 @@ def process_insert_body(
         # use later. We do this in advance so that the link update code
         # can populate overlay fields in it.
         with ctx.new() as pol_ctx:
-            pass
+            pol_ctx.rel_overlays = context.RelOverlays()
 
     needs_insert_on_conflict = bool(
         ir_stmt.on_conflict and not on_conflict_fake_iterator)
@@ -861,6 +861,11 @@ def _update_link_tables(inp_cte: pgast.CommonTableExpr) -> None:
         assert not needs_insert_on_conflict
 
         assert pol_ctx
+        # Now that all the compilation for the INSERT has been done,
+        # apply the tweaked policy overlays.
+        pol_ctx.rel_overlays = update_overlay(
+            ctx.rel_overlays, pol_ctx.rel_overlays
+        )
         with pol_ctx.reenter(), pol_ctx.newrel() as rctx:
             # Pull in ptr rel overlays, so we can see the pointers
             merge_overlays_globally((ir_stmt,), ctx=rctx)
@@ -1192,6 +1197,25 @@ def merge_overlays_globally(
     ctx.rel_overlays.ptr = ctx.rel_overlays.ptr.set(None, ptr_overlay)
 
 
+def update_overlay(
+    left: context.RelOverlays,
+    right: context.RelOverlays,
+) -> context.RelOverlays:
+    '''Make an updated copy of the left overlay using right.'''
+    left = left.copy()
+
+    for ir_stmt, tm in right.type.items():
+        ltm = left.type.get(ir_stmt, immu.Map())
+        ltm = ltm.update(tm)
+        left.type = left.type.set(ir_stmt, ltm)
+    for ir_stmt, pm in right.ptr.items():
+        lpm = left.ptr.get(ir_stmt, immu.Map())
+        lpm = lpm.update(pm)
+        left.ptr = left.ptr.set(ir_stmt, lpm)
+
+    return left
+
+
 def compile_policy_check(
     dml_cte: pgast.CommonTableExpr,
     ir_stmt: irast.MutatingStmt,
@@ -1752,7 +1776,7 @@ def process_update_body(
         # use later. We do this in advance so that the link update code
         # can populate overlay fields in it.
         with ctx.new() as pol_ctx:
-            pass
+            pol_ctx.rel_overlays = context.RelOverlays()
 
     no_update = not values and not rewrites and not single_external
     if no_update:
@@ -1806,6 +1830,11 @@ def process_update_body(
         if rewrites or single_external:
             rewrites = rewrites or {}
             assert pol_ctx
+            # Now that all the compilation for the UPDATE has been done,
+            # apply the tweaked policy overlays.
+            pol_ctx.rel_overlays = update_overlay(
+                ctx.rel_overlays, pol_ctx.rel_overlays
+            )
             with pol_ctx.reenter(), pol_ctx.new() as rctx:
                 merge_overlays_globally((ir_stmt,), ctx=rctx)
                 contents_cte, contents_rvar, values = process_update_rewrites(
@@ -2704,14 +2733,20 @@ def process_link_update(
         # Record the effect of this removal in the relation overlay
         # context to ensure that references to the link in the result
         # of this DML statement yield the expected results.
-        relctx.add_ptr_rel_overlay(
-            mptrref,
-            context.OverlayOp.EXCEPT,
-            delcte,
-            path_id=path_id.ptr_path(),
-            dml_stmts=ctx.dml_stmt_stack,
-            ctx=ctx
+        except_overlay = (lambda octx:
+            relctx.add_ptr_rel_overlay(
+                mptrref,
+                context.OverlayOp.EXCEPT,
+                delcte,
+                path_id=path_id.ptr_path(),
+                dml_stmts=ctx.dml_stmt_stack,
+                ctx=octx,
+            )
         )
+        except_overlay(ctx)
+        if policy_ctx:
+            except_overlay(policy_ctx)
+
         toplevel.append_cte(delcte)
     else:
         delqry = None
@@ -2918,7 +2953,6 @@ def register_overlays(
         )
 
     if policy_ctx:
-        policy_ctx.rel_overlays = policy_ctx.rel_overlays.copy()
         register_overlays(data_cte, policy_ctx)
 
     register_overlays(update, ctx)

tests/schemas/updates.esdl

@@ -44,6 +44,9 @@ type UpdateTest {
     link annotated_status -> Status {
         property note -> str;
     }
+    link annotated_status2 -> Status {
+        property note -> str;
+    }
 
     # for testing links to sets
     multi link tags -> Tag;

tests/test_edgeql_policies.py

@@ -794,6 +794,18 @@ async def test_edgeql_policies_11(self):
             };
             ''')
 
+        async with self.assertRaisesRegexTx(
+                edgedb.InvalidValueError,
+                r"access policy violation on insert of default::Issue"):
+            await self.con.query('''
+            insert Issue {
+                name := '', body := '',
+                watchers := {},
+                status := (select Status filter .name = 'Open'), number := '',
+                owner := (insert User { name := "???" }),
+            };
+            ''')
+
     async def test_edgeql_policies_12(self):
         await self.con.query('''
             create global cur_user_obj := (

tests/test_edgeql_update.py

@@ -1794,6 +1794,32 @@ async def test_edgeql_update_props_10(self):
             ]
         )
 
+    async def test_edgeql_update_props_11(self):
+        # Check that we can update a link property on a specific link.
+
+        # Setup some multi links
+        await self.con.execute(
+            r"""
+                UPDATE UpdateTest
+                FILTER .name = 'update-test1'
+                SET {
+                    annotated_status := (select Status limit 1),
+                    annotated_status2 := (select Status limit 1),
+                };
+            """,
+        )
+
+        await self.con.execute(
+            r"""
+                UPDATE UpdateTest
+                FILTER .name = 'update-test1'
+                SET {
+                    annotated_status := (select Status limit 1),
+                    annotated_status2 := (select Status limit 1),
+                };
+            """,
+        )
+
     async def test_edgeql_update_for_01(self):
         await self.assert_query_result(
             r"""