Implement ExecuteUpdate partial update in JSON

Against dotnet/efcore#36659

Author: roji

src/EFCore.PG/Query/Internal/NpgsqlQuerySqlGenerator.cs

@@ -159,6 +159,23 @@ protected override void GenerateTop(SelectExpression selectExpression)
         // No TOP() in PostgreSQL, see GenerateLimitOffset
     }
 
+    /// <summary>
+    ///     Generates SQL for a constant.
+    /// </summary>
+    /// <param name="sqlConstantExpression">The <see cref="SqlConstantExpression" /> for which to generate SQL.</param>
+    protected override Expression VisitSqlConstant(SqlConstantExpression sqlConstantExpression)
+    {
+        // Certain JSON functions (e.g. jsonb_set()) accept a JSONPATH argument - this is (currently) flown here as a
+        // SqlConstantExpression over IReadOnlyList<PathSegment>. Render that to a string here.
+        if (sqlConstantExpression is { Value: IReadOnlyList<PathSegment> path })
+        {
+            GenerateJsonPath(ConvertJsonPathSegments(path));
+            return sqlConstantExpression;
+        }
+
+        return base.VisitSqlConstant(sqlConstantExpression);
+    }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -1058,31 +1075,33 @@ protected virtual Expression VisitILike(PgILikeExpression likeExpression, bool n
     protected override Expression VisitJsonScalar(JsonScalarExpression jsonScalarExpression)
     {
         // TODO: Stop producing empty JsonScalarExpressions, #30768
-        var path = jsonScalarExpression.Path;
-        if (path.Count == 0)
+        var segmentsPath = jsonScalarExpression.Path;
+        if (segmentsPath.Count == 0)
         {
             Visit(jsonScalarExpression.Json);
             return jsonScalarExpression;
         }
 
+        var path = ConvertJsonPathSegments(segmentsPath);
+
         switch (jsonScalarExpression.TypeMapping)
         {
             // This case is for when a nested JSON entity is being accessed. We want the json/jsonb fragment in this case (not text),
             // so we can perform further JSON operations on it.
             case NpgsqlStructuralJsonTypeMapping:
-                GenerateJsonPath(returnsText: false);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: false, path);
                 break;
 
             // No need to cast the output when we expect a string anyway
             case StringTypeMapping:
-                GenerateJsonPath(returnsText: true);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: true, path);
                 break;
 
             // bytea requires special handling, since we encode the binary data as base64 inside the JSON, but that requires a special
             // conversion function to be extracted out to a PG bytea.
             case NpgsqlByteArrayTypeMapping:
                 Sql.Append("decode(");
-                GenerateJsonPath(returnsText: true);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: true, path);
                 Sql.Append(", 'base64')");
                 break;
 
@@ -1092,33 +1111,20 @@ protected override Expression VisitJsonScalar(JsonScalarExpression jsonScalarExp
             case NpgsqlArrayTypeMapping arrayMapping:
                 Sql.Append("(ARRAY(SELECT CAST(element AS ").Append(arrayMapping.ElementTypeMapping.StoreType)
                     .Append(") FROM jsonb_array_elements_text(");
-                GenerateJsonPath(returnsText: false);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: false, path);
                 Sql.Append(") WITH ORDINALITY AS t(element) ORDER BY ordinality))");
                 break;
 
             default:
                 Sql.Append("CAST(");
-                GenerateJsonPath(returnsText: true);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: true, path);
                 Sql.Append(" AS ");
                 Sql.Append(jsonScalarExpression.TypeMapping!.StoreType);
                 Sql.Append(")");
                 break;
         }
 
         return jsonScalarExpression;
-
-        void GenerateJsonPath(bool returnsText)
-            => this.GenerateJsonPath(
-                jsonScalarExpression.Json,
-                returnsText: returnsText,
-                jsonScalarExpression.Path.Select(
-                    s => s switch
-                    {
-                        { PropertyName: string propertyName }
-                            => new SqlConstantExpression(propertyName, _textTypeMapping ??= _typeMappingSource.FindMapping(typeof(string))),
-                        { ArrayIndex: SqlExpression arrayIndex } => arrayIndex,
-                        _ => throw new UnreachableException()
-                    }).ToList());
     }
 
     /// <summary>
@@ -1148,6 +1154,11 @@ private void GenerateJsonPath(SqlExpression expression, bool returnsText, IReadO
         // Multiple path components
         Sql.Append(returnsText ? " #>> " : " #> ");
 
+        GenerateJsonPath(path);
+    }
+
+    private void GenerateJsonPath(IReadOnlyList<SqlExpression> path)
+    {
         // Use simplified array literal syntax if all path components are constants for cleaner SQL
         if (path.All(p => p is SqlConstantExpression { Value: var pathSegment }
                 && (pathSegment is not string s || s.All(char.IsAsciiLetterOrDigit))))
@@ -1173,6 +1184,23 @@ private void GenerateJsonPath(SqlExpression expression, bool returnsText, IReadO
         }
     }
 
+    /// <summary>
+    ///     Converts the standard EF <see cref="IReadOnlyList{PathSegment}" /> to an <see cref="IReadOnlyList{SqlExpression}" />
+    ///     (the EF built-in <see cref="JsonScalarExpression" /> and <see cref="JsonQueryExpression" /> don't support non-constant
+    ///     property names, but we do via the Npgsql-specific JSON DOM support).
+    /// </summary>
+    private IReadOnlyList<SqlExpression> ConvertJsonPathSegments(IReadOnlyList<PathSegment> path)
+        => path
+            .Select(
+                s => s switch
+                {
+                    { PropertyName: string propertyName }
+                        => new SqlConstantExpression(propertyName, _textTypeMapping ??= _typeMappingSource.FindMapping(typeof(string))),
+                    { ArrayIndex: SqlExpression arrayIndex } => arrayIndex,
+                    _ => throw new UnreachableException()
+                })
+            .ToList();
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore.PG/Query/Internal/NpgsqlQueryableMethodTranslatingExpressionVisitor.cs

@@ -1046,6 +1046,8 @@ protected override bool IsNaturallyOrdered(SelectExpression selectExpression)
                     [{ Expression: ColumnExpression { Name: "ordinality", TableAlias: var orderingTableAlias } }]
                 && orderingTableAlias == unnest.Alias);
 
+    #region ExecuteUpdate
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -1100,6 +1102,119 @@ protected override bool IsValidSelectExpressionForExecuteUpdate(
         return true;
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    protected override bool TrySerializeScalarToJson(
+        JsonScalarExpression target,
+        SqlExpression value,
+        [NotNullWhen(true)] out SqlExpression? jsonValue)
+    {
+        var jsonTypeMapping = ((ColumnExpression)target.Json).TypeMapping!;
+
+        if (
+            // The base implementation doesn't handle serializing arbitrary SQL expressions to JSON, since that's
+            // database-specific. In PostgreSQL we simply do this by wrapping any expression in to_jsonb().
+            !base.TrySerializeScalarToJson(target, value, out jsonValue)
+            // In addition, for string, numeric and bool, the base implementation simply returns the value as-is, since most databases allow
+            // passing these native types directly to their JSON partial update function. In PostgreSQL, jsonb_set() always requires jsonb,
+            // so we wrap those expression with to_jsonb() as well.
+            || jsonValue.TypeMapping?.StoreType is not "jsonb" and not "json")
+        {
+            switch (value.TypeMapping!.StoreType)
+            {
+                case "jsonb" or "json":
+                    jsonValue = value;
+                    return true;
+
+                case "bytea":
+                    value = _sqlExpressionFactory.Function(
+                        "encode",
+                        [value, _sqlExpressionFactory.Constant("base64")],
+                        nullable: true,
+                        argumentsPropagateNullability: [true, true],
+                        typeof(string),
+                        _typeMappingSource.FindMapping(typeof(string))!
+                    );
+                    break;
+            }
+
+            jsonValue = _sqlExpressionFactory.Function(
+                jsonTypeMapping.StoreType switch
+                {
+                    "jsonb" => "to_jsonb",
+                    "json" => "to_json",
+                    _ => throw new UnreachableException()
+                },
+                // Make sure PG interprets constant values correctly by adding explicit typing based on the target property's type mapping.
+                // Note that we can only be here for scalar properties, for structural types we always already get a jsonb/json value
+                // and don't need to add to_jsonb/to_json.
+                [value is SqlConstantExpression ? _sqlExpressionFactory.Convert(value, target.Type, target.TypeMapping) : value],
+                nullable: true,
+                argumentsPropagateNullability: [true],
+                typeof(string),
+                jsonTypeMapping);
+        }
+
+        return true;
+    }
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    protected override SqlExpression? GenerateJsonPartialUpdateSetter(
+        Expression target,
+        SqlExpression value,
+        ref SqlExpression? existingSetterValue)
+    {
+        var (jsonColumn, path) = target switch
+        {
+            JsonScalarExpression j => ((ColumnExpression)j.Json, j.Path),
+            JsonQueryExpression j => (j.JsonColumn, j.Path),
+
+            _ => throw new UnreachableException(),
+        };
+
+        var jsonSet = _sqlExpressionFactory.Function(
+            jsonColumn.TypeMapping?.StoreType switch
+            {
+                "jsonb" => "jsonb_set",
+                "json" => "json_set",
+                _ => throw new UnreachableException()
+            },
+            arguments:
+            [
+                existingSetterValue ?? jsonColumn,
+                // Hack: Rendering of JSONPATH strings happens in value generation. We can have a special expression for modify to hold the
+                // IReadOnlyList<PathSegment> (just like Json{Scalar,Query}Expression), but instead we do the slight hack of packaging it
+                // as a constant argument; it will be unpacked and handled in SQL generation.
+                _sqlExpressionFactory.Constant(path, RelationalTypeMapping.NullMapping),
+                value
+            ],
+            nullable: true,
+            argumentsPropagateNullability: [true, true, true],
+            typeof(string),
+            jsonColumn.TypeMapping);
+
+        if (existingSetterValue is null)
+        {
+            return jsonSet;
+        }
+        else
+        {
+            existingSetterValue = jsonSet;
+            return null;
+        }
+    }
+
+    #endregion ExecuteUpdate
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in