Revamp support for primitive collections within JSON

Closes #3122

Author: roji

src/EFCore.PG/Extensions/Internal/NpgsqlShapedQueryExpressionExtensions.cs

@@ -56,7 +56,7 @@ public static bool TryExtractArray(
                 Offset: null
             } select
             && (ignorePredicate || select.Predicate is null)
-            // We can only apply the indexing if the JSON array is ordered by its natural ordered, i.e. by the "ordinality" column that
+            // We can only apply the indexing if the array is ordered by its natural ordered, i.e. by the "ordinality" column that
             // we created in TranslatePrimitiveCollection. For example, if another ordering has been applied (e.g. by the array elements
             // themselves), we can no longer simply index into the original array.
             && (ignoreOrderings
@@ -76,6 +76,78 @@ public static bool TryExtractArray(
         return false;
     }
 
+    /// <summary>
+    ///     If the given <paramref name="source" /> wraps a JSON-array-returning expression without any additional clauses (e.g. filter,
+    ///     ordering...), returns that expression.
+    /// </summary>
+    /// <remarks>
+    ///     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.
+    /// </remarks>
+    public static bool TryExtractJsonArray(
+        this ShapedQueryExpression source,
+        [NotNullWhen(true)] out SqlExpression? jsonArray,
+        [NotNullWhen(true)] out SqlExpression? projectedElement,
+        out bool isElementNullable,
+        bool ignoreOrderings = false,
+        bool ignorePredicate = false)
+    {
+        if (source.QueryExpression is SelectExpression
+            {
+                Tables:
+                [
+                    TableValuedFunctionExpression
+                {
+                    Name: "jsonb_array_elements_text" or "json_array_elements_text",
+                    Arguments: [var json]
+                } tvf
+                ],
+                GroupBy: [],
+                Having: null,
+                IsDistinct: false,
+                Limit: null,
+                Offset: null
+            } select
+            && (ignorePredicate || select.Predicate is null)
+            // We can only apply the indexing if the array is ordered by its natural ordered, i.e. by the "ordinality" column that
+            // we created in TranslatePrimitiveCollection. For example, if another ordering has been applied (e.g. by the array elements
+            // themselves), we can no longer simply index into the original array.
+            && (ignoreOrderings
+                || select.Orderings is []
+                || (select.Orderings is [{ Expression: ColumnExpression { Name: "ordinality", TableAlias: var orderingTableAlias } }]
+                    && orderingTableAlias == tvf.Alias))
+            && TryGetScalarProjection(source, out var projectedScalar))
+        {
+            jsonArray = json;
+
+            // The projected ColumnExpression is wrapped in a Convert to apply the element type mapping - unless it happens to be text.
+            switch (projectedScalar)
+            {
+                case SqlUnaryExpression
+                {
+                    OperatorType: ExpressionType.Convert,
+                    Operand: ColumnExpression { IsNullable: var isNullable }
+                } convert:
+                    projectedElement = convert;
+                    isElementNullable = isNullable;
+                    return true;
+                case ColumnExpression { IsNullable: var isNullable } column:
+                    projectedElement = column;
+                    isElementNullable = isNullable;
+                    return true;
+                default:
+                    throw new UnreachableException();
+            }
+        }
+
+        jsonArray = null;
+        projectedElement = null;
+        isElementNullable = false;
+        return false;
+    }
+
     /// <summary>
     ///     If the given <paramref name="source" /> wraps a <see cref="ValuesExpression" /> without any additional clauses (e.g. filter,
     ///     ordering...), converts that to a <see cref="NewArrayExpression" /> and returns that.
@@ -86,7 +158,7 @@ public static bool TryExtractArray(
     ///     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.
     /// </remarks>
-    public static bool TryConvertValuesToArray(
+    public static bool TryConvertToArray(
         this ShapedQueryExpression source,
         [NotNullWhen(true)] out SqlExpression? array,
         bool ignoreOrderings = false,
@@ -128,13 +200,28 @@ private static bool IsPostgresArray(SqlExpression expression)
         {
             { TypeMapping: NpgsqlArrayTypeMapping } => true,
             { TypeMapping: NpgsqlMultirangeTypeMapping } => false,
+            { TypeMapping: NpgsqlJsonTypeMapping } => false,
             { Type: var type } when type.IsMultirange() => false,
             _ => true
         };
 
     private static bool TryGetProjectedColumn(
         ShapedQueryExpression shapedQueryExpression,
         [NotNullWhen(true)] out ColumnExpression? projectedColumn)
+    {
+        if (TryGetScalarProjection(shapedQueryExpression, out var scalar) && scalar is ColumnExpression column)
+        {
+            projectedColumn = column;
+            return true;
+        }
+
+        projectedColumn = null;
+        return false;
+    }
+
+    private static bool TryGetScalarProjection(
+        ShapedQueryExpression shapedQueryExpression,
+        [NotNullWhen(true)] out SqlExpression? projectedScalar)
     {
         var shaperExpression = shapedQueryExpression.ShaperExpression;
         if (shaperExpression is UnaryExpression { NodeType: ExpressionType.Convert } unaryExpression
@@ -146,13 +233,13 @@ private static bool TryGetProjectedColumn(
 
         if (shaperExpression is ProjectionBindingExpression projectionBindingExpression
             && shapedQueryExpression.QueryExpression is SelectExpression selectExpression
-            && selectExpression.GetProjection(projectionBindingExpression) is ColumnExpression c)
+            && selectExpression.GetProjection(projectionBindingExpression) is SqlExpression scalar)
         {
-            projectedColumn = c;
+            projectedScalar = scalar;
             return true;
         }
 
-        projectedColumn = null;
+        projectedScalar = null;
         return false;
     }
 }

src/EFCore.PG/Metadata/Conventions/NpgsqlConventionSetBuilder.cs

@@ -1,4 +1,5 @@
 using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure.Internal;
+using Npgsql.EntityFrameworkCore.PostgreSQL.Storage.Internal;
 
 namespace Npgsql.EntityFrameworkCore.PostgreSQL.Metadata.Conventions;
 
@@ -17,7 +18,7 @@ namespace Npgsql.EntityFrameworkCore.PostgreSQL.Metadata.Conventions;
 /// </remarks>
 public class NpgsqlConventionSetBuilder : RelationalConventionSetBuilder
 {
-    private readonly IRelationalTypeMappingSource _typeMappingSource;
+    private readonly NpgsqlTypeMappingSource _typeMappingSource;
     private readonly Version _postgresVersion;
     private readonly IReadOnlyList<EnumDefinition> _enumDefinitions;
 
@@ -35,7 +36,7 @@ public NpgsqlConventionSetBuilder(
         INpgsqlSingletonOptions npgsqlSingletonOptions)
         : base(dependencies, relationalDependencies)
     {
-        _typeMappingSource = typeMappingSource;
+        _typeMappingSource = (NpgsqlTypeMappingSource)typeMappingSource;
         _postgresVersion = npgsqlSingletonOptions.PostgresVersion;
         _enumDefinitions = npgsqlSingletonOptions.EnumDefinitions;
     }

src/EFCore.PG/Metadata/Conventions/NpgsqlPostgresModelFinalizingConvention.cs

@@ -1,4 +1,6 @@
 using Npgsql.EntityFrameworkCore.PostgreSQL.Infrastructure.Internal;
+using Npgsql.EntityFrameworkCore.PostgreSQL.Storage.Internal;
+using Npgsql.EntityFrameworkCore.PostgreSQL.Storage.Internal.Mapping;
 
 namespace Npgsql.EntityFrameworkCore.PostgreSQL.Metadata.Conventions;
 
@@ -8,24 +10,10 @@ namespace Npgsql.EntityFrameworkCore.PostgreSQL.Metadata.Conventions;
 /// <remarks>
 ///     See <see href="https://aka.ms/efcore-docs-conventions">Model building conventions</see>.
 /// </remarks>
-public class NpgsqlPostgresModelFinalizingConvention : IModelFinalizingConvention
+public class NpgsqlPostgresModelFinalizingConvention(
+    NpgsqlTypeMappingSource typeMappingSource,
+    IReadOnlyList<EnumDefinition> enumDefinitions) : IModelFinalizingConvention
 {
-    private readonly IRelationalTypeMappingSource _typeMappingSource;
-    private readonly IReadOnlyList<EnumDefinition> _enumDefinitions;
-
-    /// <summary>
-    ///     Creates a new instance of <see cref="NpgsqlPostgresModelFinalizingConvention" />.
-    /// </summary>
-    /// <param name="typeMappingSource">The type mapping source to use.</param>
-    /// <param name="enumDefinitions"></param>
-    public NpgsqlPostgresModelFinalizingConvention(
-        IRelationalTypeMappingSource typeMappingSource,
-        IReadOnlyList<EnumDefinition> enumDefinitions)
-    {
-        _typeMappingSource = typeMappingSource;
-        _enumDefinitions = enumDefinitions;
-    }
-
     /// <inheritdoc />
     public virtual void ProcessModelFinalizing(IConventionModelBuilder modelBuilder, IConventionContext<IConventionModelBuilder> context)
     {
@@ -34,7 +22,7 @@ public virtual void ProcessModelFinalizing(IConventionModelBuilder modelBuilder,
             foreach (var property in entityType.GetDeclaredProperties())
             {
                 var typeMapping = (RelationalTypeMapping?)property.FindTypeMapping()
-                    ?? _typeMappingSource.FindMapping((IProperty)property);
+                    ?? typeMappingSource.FindMapping((IProperty)property);
 
                 if (typeMapping is not null)
                 {
@@ -52,7 +40,7 @@ public virtual void ProcessModelFinalizing(IConventionModelBuilder modelBuilder,
     /// </summary>
     protected virtual void SetupEnums(IConventionModelBuilder modelBuilder)
     {
-        foreach (var enumDefinition in _enumDefinitions)
+        foreach (var enumDefinition in enumDefinitions)
         {
             modelBuilder.HasPostgresEnum(
                 enumDefinition.StoreTypeSchema,

src/EFCore.PG/Query/Expressions/Internal/PgTableValuedFunctionExpression.cs

@@ -55,7 +55,7 @@ public PgTableValuedFunctionExpression(
         string alias,
         string name,
         IReadOnlyList<SqlExpression> arguments,
-        IReadOnlyList<ColumnInfo>? columnInfos,
+        IReadOnlyList<ColumnInfo>? columnInfos = null,
         bool withOrdinality = true)
         : base(alias, name, schema: null, builtIn: true, arguments)
     {

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

@@ -1086,12 +1086,19 @@ protected override Expression VisitJsonScalar(JsonScalarExpression jsonScalarExp
 
         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.
+            // Nested JSON structural type. We want the json/jsonb fragment in this case (not text), so we can perform further JSON
+            // operations on it.
             case NpgsqlStructuralJsonTypeMapping:
                 GenerateJsonPath(jsonScalarExpression.Json, returnsText: false, path);
                 break;
 
+            // Scalar collection mapped to JSON (either because it's nested within a JSON document, or because the user explicitly
+            // opted for this rather than the default PG array mapping).
+            case NpgsqlJsonTypeMapping typeMapping:
+                Check.DebugAssert(typeMapping.ElementTypeMapping is not null);
+                GenerateJsonPath(jsonScalarExpression.Json, returnsText: false, path);
+                break;
+
             // No need to cast the output when we expect a string anyway
             case StringTypeMapping:
                 GenerateJsonPath(jsonScalarExpression.Json, returnsText: true, path);
@@ -1105,15 +1112,10 @@ protected override Expression VisitJsonScalar(JsonScalarExpression jsonScalarExp
                 Sql.Append(", 'base64')");
                 break;
 
-            // Arrays require special handling; we cannot simply cast a JSON array (as text) to a PG array ([1,2,3] isn't a valid PG array
-            // representation). We use jsonb_array_elements_text to extract the array elements as a set, cast them to their PG element type
-            // and then build an array from that.
-            case NpgsqlArrayTypeMapping arrayMapping:
-                Sql.Append("(ARRAY(SELECT CAST(element AS ").Append(arrayMapping.ElementTypeMapping.StoreType)
-                    .Append(") FROM jsonb_array_elements_text(");
-                GenerateJsonPath(jsonScalarExpression.Json, returnsText: false, path);
-                Sql.Append(") WITH ORDINALITY AS t(element) ORDER BY ordinality))");
-                break;
+            // We should never have an NpgsqlArrayTypeMapping within a JSON document; scalar collections should be represented as an
+            // NpgsqlJsonTypeMapping with the appropriate ElementTypeMapping, just like in other providers.
+            case NpgsqlArrayTypeMapping:
+                throw new UnreachableException();
 
             default:
                 Sql.Append("CAST(");