Fixed the combination of using an Include() + Select() where the sele…

…ct uses the Id of the original document. Closes GH-3096

src/LinqTests/Bugs/Bug_2223_list_contains_any_generates_invalid_sql.cs

@@ -64,12 +64,12 @@ public async Task should_be_able_to_query_with_multiple_list_items_and_have_incl
             var otherTestEntityLookup = new Dictionary<Guid, OtherTestEntity>();
             var entities = await session.Query<TestEntity>()
                 .Include(x => x.OtherIds, otherTestEntityLookup)
-                .Where(x => Enumerable.Any<Guid>(x.OtherIds, id => otherIdsQuery.Contains(id)))
+                .Where(x => x.OtherIds.Any<Guid>(id => otherIdsQuery.Contains(id)))
                 .ToListAsync();
 
             entities.Count.ShouldBe(1);
-            ShouldBeTestExtensions.ShouldBe(entities[0].OtherIds.Count, 2);
-            ShouldBeEnumerableTestExtensions.ShouldContain(entities[0].OtherIds, otherEntityTestId);
+            entities[0].OtherIds.Count.ShouldBe(2);
+            entities[0].OtherIds.ShouldContain(otherEntityTestId);
 
             otherTestEntityLookup.Count.ShouldBe(2);
             otherTestEntityLookup.ShouldContainKey(otherEntityTestId);

src/LinqTests/Bugs/Bug_3096_include_where_select.cs

@@ -0,0 +1,53 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Marten;
+using Marten.Testing.Documents;
+using Marten.Testing.Harness;
+using Shouldly;
+using Xunit.Abstractions;
+
+namespace LinqTests.Bugs;
+
+public class Bug_3096_include_where_select : IntegrationContext
+{
+    private readonly ITestOutputHelper _output;
+
+    public Bug_3096_include_where_select(DefaultStoreFixture fixture, ITestOutputHelper output) : base(fixture)
+    {
+        _output = output;
+    }
+
+    [Fact]
+    public void include_to_dictionary_with_where_and_projection()
+    {
+        var user1 = new User();
+        var user2 = new User();
+
+        var issue1 = new Issue { AssigneeId = user1.Id, Title = "Garage Door is ok" };
+        var issue2 = new Issue { AssigneeId = user2.Id, Title = "Garage Door is busted" };
+        var issue3 = new Issue { AssigneeId = user2.Id, Title = "Garage Door is busted" };
+
+        using var session = theStore.IdentitySession();
+        session.Store(user1, user2);
+        session.Store(issue1, issue2, issue3);
+        session.SaveChanges();
+
+        using var query = theStore.QuerySession();
+        query.Logger = new TestOutputMartenLogger(_output);
+
+        var dict = new Dictionary<Guid, User>();
+
+        var issues = query
+            .Query<Issue>()
+            .Where(i => i.Title.Contains("ok"))
+            .Include(x => x.AssigneeId, dict)
+            .Select(i => new { i.Id, i.Title, })
+            .ToArray();
+
+        issues.Length.ShouldBe(1);
+
+        dict.Count.ShouldBe(1);
+        dict.ContainsKey(user1.Id).ShouldBeTrue();
+    }
+}

src/LinqTests/Bugs/Bug_490_hierarchy_and_include.cs

@@ -68,7 +68,7 @@ public void load_abstract_type_with_include()
 
         using (var session = theStore.QuerySession())
         {
-            List<Account> accounts = new List<Account>();
+            var accounts = new List<Account>();
             session.Query<Activity>()
                 .Include(a => a.AccountId, accounts)
                 .ToList()
@@ -110,4 +110,4 @@ public async Task load_abstract_type_with_include_async()
             accounts.First().Id.ShouldBe(1);
         }
     }
-}
+}

src/LinqTests/Includes/end_to_end_query_with_include.cs

@@ -750,6 +750,8 @@ public void include_many_to_list()
 
         using (var query = theStore.QuerySession())
         {
+            query.Logger = new TestOutputMartenLogger(_output);
+
             var list = new List<User>();
 
             query.Query<Group>()

src/Marten/Internal/Storage/DataAndIdSelectClause.cs

@@ -0,0 +1,48 @@
+using System;
+using System.Linq;
+using JasperFx.Core;
+using Marten.Linq;
+using Marten.Linq.QueryHandlers;
+using Marten.Linq.Selectors;
+using Marten.Linq.SqlGeneration;
+using Weasel.Postgresql;
+using Weasel.Postgresql.SqlGeneration;
+
+namespace Marten.Internal.Storage;
+
+internal class DataAndIdSelectClause<T>: ISelectClause, IModifyableFromObject
+{
+    private readonly IDocumentStorage<T> _inner;
+
+    public DataAndIdSelectClause(IDocumentStorage<T> inner)
+    {
+        _inner = inner;
+        FromObject = inner.FromObject;
+    }
+
+    public void Apply(ICommandBuilder builder)
+    {
+        builder.Append($"select {_inner.SelectFields().Concat(["d.id"]).Join(", ")} from ");
+        builder.Append(FromObject);
+        builder.Append(" as d");
+    }
+
+    public string FromObject { get; set; }
+    public Type SelectedType => typeof(T);
+    public string[] SelectFields() => ["d.data", "d.id"];
+
+    public ISelector BuildSelector(IMartenSession session)
+    {
+        return _inner.BuildSelector(session);
+    }
+
+    public IQueryHandler<T1> BuildHandler<T1>(IMartenSession session, ISqlFragment topStatement, ISqlFragment currentStatement)
+    {
+        return _inner.BuildHandler<T1>(session, topStatement, currentStatement);
+    }
+
+    public ISelectClause UseStatistics(QueryStatistics statistics)
+    {
+        return _inner.UseStatistics(statistics);
+    }
+}

src/Marten/Internal/Storage/QueryOnlyDocumentStorage.cs

@@ -4,16 +4,27 @@
 using System.Threading.Tasks;
 using Marten.Internal.CodeGeneration;
 using Marten.Linq.Selectors;
+using Marten.Linq.SqlGeneration;
 using Marten.Schema;
 
 namespace Marten.Internal.Storage;
 
-public abstract class QueryOnlyDocumentStorage<T, TId>: DocumentStorage<T, TId>
+internal interface IQueryOnlyDocumentStorage: IDocumentStorage
+{
+    ISelectClause SelectClauseForIncludes();
+}
+
+public abstract class QueryOnlyDocumentStorage<T, TId>: DocumentStorage<T, TId>, IQueryOnlyDocumentStorage
 {
     public QueryOnlyDocumentStorage(DocumentMapping document): base(StorageStyle.QueryOnly, document)
     {
     }
 
+    public ISelectClause SelectClauseForIncludes()
+    {
+        return new DataAndIdSelectClause<T>(this);
+    }
+
     public sealed override void Store(IMartenSession session, T document)
     {
     }

src/Marten/Linq/Includes/TemporaryTableStatement.cs

@@ -1,4 +1,6 @@
+using System.Linq;
 using Marten.Internal;
+using Marten.Internal.Storage;
 using Marten.Linq.SqlGeneration;
 using Weasel.Postgresql;
 
@@ -9,8 +11,15 @@ public class TemporaryTableStatement: Statement
     public TemporaryTableStatement(Statement inner, IMartenSession session)
     {
         Inner = inner;
+        var selectorStatement = Inner.SelectorStatement();
+        selectorStatement.Mode = StatementMode.Inner;
 
-        Inner.SelectorStatement().Mode = StatementMode.Inner;
+        // This is ugly, but you need to pick up the id column *just* in case there's a Select()
+        // clause that needs it.
+        if (selectorStatement.SelectClause is IQueryOnlyDocumentStorage s)
+        {
+            selectorStatement.SelectClause = s.SelectClauseForIncludes();
+        }
 
         ExportName = session.NextTempTableName();
     }

src/Marten/Linq/SqlGeneration/SelectDataSelectClause.cs

@@ -9,7 +9,7 @@
 
 namespace Marten.Linq.SqlGeneration;
 
-internal class SelectDataSelectClause<T>: ISelectClause, IScalarSelectClause
+internal class SelectDataSelectClause<T>: ISelectClause, IScalarSelectClause, IModifyableFromObject
 {
     public SelectDataSelectClause(string from, ISqlFragment selector)
     {
@@ -21,7 +21,7 @@ public SelectDataSelectClause(string from, ISqlFragment selector)
 
     public Type SelectedType => typeof(T);
 
-    public string FromObject { get; }
+    public string FromObject { get; set; }
 
     public void Apply(ICommandBuilder sql)
     {