Identify WITH as generally composable in FromSql

src/EFCore.Relational/Query/QuerySqlGenerator.cs

@@ -444,81 +444,64 @@ protected override Expression VisitFromSql(FromSqlExpression fromSqlExpression)
         /// <exception cref="InvalidOperationException">The given SQL isn't composable.</exception>
         protected virtual void CheckComposableSql(string sql)
         {
-            Check.NotNull(sql, nameof(sql));
-
-            var pos = -1;
-            char c;
+            var span = sql.AsSpan().TrimStart();
 
             while (true)
             {
-                c = NextChar();
-
-                if (char.IsWhiteSpace(c))
-                {
-                    continue;
-                }
-
                 // SQL -- comment
-                if (c == '-')
+                if (span.StartsWith("--"))
                 {
-                    if (NextChar() != '-')
-                    {
-                        throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
-                    }
-
-                    while (NextChar() != '\n')
-                    {
-                    }
-
+                    var i = span.IndexOf('\n');
+                    span = i > 0
+                        ? span.Slice(i + 1).TrimStart()
+                        : throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
                     continue;
                 }
 
                 // SQL /* */ comment
-                if (c == '/')
+                if (span.StartsWith("/*"))
                 {
-                    if (NextChar() != '*')
-                    {
-                        throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
-                    }
-
-                    while (true)
-                    {
-                        while (NextChar() != '*')
-                        {
-                        }
-
-                        if (NextChar() == '/')
-                        {
-                            break;
-                        }
-                    }
-
+                    var i = span.IndexOf("*/");
+                    span = i > 0
+                        ? span.Slice(i + 2).TrimStart()
+                        : throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
                     continue;
                 }
 
-                if (char.ToLowerInvariant(c) == 's'
-                    && char.ToLowerInvariant(NextChar()) == 'e'
-                    && char.ToLowerInvariant(NextChar()) == 'l'
-                    && char.ToLowerInvariant(NextChar()) == 'e'
-                    && char.ToLowerInvariant(NextChar()) == 'c'
-                    && char.ToLowerInvariant(NextChar()) == 't')
-                {
-                    var (c1, c2) = (NextChar(), NextChar());
-                    if (char.IsWhiteSpace(c1)
-                        || c1 == '-' && c2 == '-'
-                        || c1 == '/' && c2 == '*')
-                    {
-                        return;
-                    }
-                }
+                break;
+            }
 
+            CheckComposableSqlTrimmed(span);
+        }
+
+        /// <summary>
+        ///     Checks whether a given SQL string is composable, i.e. can be embedded as a subquery within a
+        ///     larger SQL query. The provided <paramref name="sql" /> is already trimmed for whitespace and comments.
+        /// </summary>
+        /// <param name="sql">An trimmed SQL string to be checked for composability.</param>
+        /// <exception cref="InvalidOperationException">The given SQL isn't composable.</exception>
+        protected virtual void CheckComposableSqlTrimmed(ReadOnlySpan<char> sql)
+        {
+            if (sql.StartsWith("SELECT", StringComparison.OrdinalIgnoreCase))
+            {
+                sql = sql.Slice("SELECT".Length);
+            }
+            else if (sql.StartsWith("WITH", StringComparison.OrdinalIgnoreCase))
+            {
+                sql = sql.Slice("WITH".Length);
+            }
+            else
+            {
                 throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
             }
 
-            char NextChar()
-                => ++pos < sql.Length
-                    ? sql[pos]
-                    : throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
+            if (sql.Length > 0
+                && (char.IsWhiteSpace(sql[0]) || sql.StartsWith("--") || sql.StartsWith("/*")))
+            {
+                return;
+            }
+
+            throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
         }
 
         /// <inheritdoc />

src/EFCore.SqlServer/Query/Internal/SqlServerQuerySqlGenerator.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.Query;
 using Microsoft.EntityFrameworkCore.Query.SqlExpressions;
@@ -163,5 +164,16 @@ protected override Expression VisitExtension(Expression extensionExpression)
 
             return base.VisitExtension(extensionExpression);
         }
+
+        /// <inheritdoc />
+        protected override void CheckComposableSqlTrimmed(ReadOnlySpan<char> sql)
+        {
+            base.CheckComposableSqlTrimmed(sql);
+
+            if (sql.StartsWith("WITH", StringComparison.OrdinalIgnoreCase))
+            {
+                throw new InvalidOperationException(RelationalStrings.FromSqlNonComposable);
+            }
+        }
     }
 }

test/EFCore.Relational.Specification.Tests/Query/FromSqlQueryTestBase.cs

@@ -1502,6 +1502,28 @@ public virtual async Task FromSqlRaw_with_dbParameter_mixed_in_subquery(bool asy
             Assert.Equal(26, actual.Length);
         }
 
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual async Task FromSqlRaw_composed_with_common_table_expression(bool async)
+        {
+            using var context = CreateContext();
+            var query = context.Set<Customer>()
+                .FromSqlRaw(
+                    NormalizeDelimitersInRawString(
+                        @"WITH [Customers2] AS (
+    SELECT * FROM [Customers]
+)
+SELECT * FROM [Customers2]"))
+                .Where(c => c.ContactName.Contains("z"));
+
+            var actual = async
+                ? await query.ToArrayAsync()
+                : query.ToArray();
+
+            Assert.Equal(14, actual.Length);
+            Assert.Equal(14, context.ChangeTracker.Entries().Count());
+        }
+
         protected string NormalizeDelimitersInRawString(string sql)
             => Fixture.TestStore.NormalizeDelimitersInRawString(sql);
 

test/EFCore.Relational.Tests/Query/Internal/QuerySqlGeneratorTest.cs

@@ -21,11 +21,18 @@ public class QuerySqlGeneratorTest
         [InlineData("")]
         [InlineData("--SELECT")]
         public void CheckComposableSql_throws(string sql)
-            => Assert.Equal(
+        {
+            Assert.Equal(
                 RelationalStrings.FromSqlNonComposable,
                 Assert.Throws<InvalidOperationException>(
                     () => CreateDummyQuerySqlGenerator().CheckComposableSql(sql)).Message);
 
+            Assert.Equal(
+                RelationalStrings.FromSqlNonComposable,
+                Assert.Throws<InvalidOperationException>(
+                    () => CreateDummyQuerySqlGenerator().CheckComposableSql(sql.Replace("SELECT", "WITH"))).Message);
+        }
+
         [Theory]
         [InlineData("SELECT something")]
         [InlineData("   SELECT something")]
@@ -36,7 +43,11 @@ public void CheckComposableSql_throws(string sql)
         [InlineData("  /* multi\n*line\r\n * comment */ \nSELECT--\n1")]
         [InlineData("SELECT/* comment */1")]
         public void CheckComposableSql_does_not_throw(string sql)
-            => CreateDummyQuerySqlGenerator().CheckComposableSql(sql);
+        {
+            CreateDummyQuerySqlGenerator().CheckComposableSql(sql);
+
+            CreateDummyQuerySqlGenerator().CheckComposableSql(sql.Replace("SELECT", "WITH"));
+        }
 
         private DummyQuerySqlGenerator CreateDummyQuerySqlGenerator()
             => new(

test/EFCore.SqlServer.FunctionalTests/Query/FromSqlQuerySqlServerTest.cs

@@ -1,10 +1,12 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Data.Common;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.Data.SqlClient;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.TestUtilities;
 using Xunit;
 using Xunit.Abstractions;
@@ -752,6 +754,14 @@ SELECT 1
     WHERE [m].[CustomerID] = [o].[CustomerID])");
         }
 
+        public override async Task FromSqlRaw_composed_with_common_table_expression(bool async)
+        {
+            var exception =
+                await Assert.ThrowsAsync<InvalidOperationException>(() => base.FromSqlRaw_composed_with_common_table_expression(async));
+
+            Assert.Equal(RelationalStrings.FromSqlNonComposable, exception.Message);
+        }
+
         protected override DbParameter CreateDbParameter(string name, object value)
             => new SqlParameter { ParameterName = name, Value = value };
 

test/EFCore.Sqlite.FunctionalTests/Query/FromSqlQuerySqliteTest.cs

@@ -75,7 +75,25 @@ public override Task Bad_data_error_handling_invalid_cast_no_tracking(bool async
             return Task.CompletedTask;
         }
 
+        public override async Task FromSqlRaw_composed_with_common_table_expression(bool async)
+        {
+            await base.FromSqlRaw_composed_with_common_table_expression(async);
+
+            AssertSql(
+                @"SELECT ""m"".""CustomerID"", ""m"".""Address"", ""m"".""City"", ""m"".""CompanyName"", ""m"".""ContactName"", ""m"".""ContactTitle"", ""m"".""Country"", ""m"".""Fax"", ""m"".""Phone"", ""m"".""PostalCode"", ""m"".""Region""
+FROM (
+    WITH ""Customers2"" AS (
+        SELECT * FROM ""Customers""
+    )
+    SELECT * FROM ""Customers2""
+) AS ""m""
+WHERE ('z' = '') OR (instr(""m"".""ContactName"", 'z') > 0)");
+        }
+
         protected override DbParameter CreateDbParameter(string name, object value)
             => new SqliteParameter { ParameterName = name, Value = value };
+
+        private void AssertSql(params string[] expected)
+            => Fixture.TestSqlLoggerFactory.AssertBaseline(expected);
     }
 }