From e92b40c48597eb583292e874f562861cb45f95ca Mon Sep 17 00:00:00 2001
From: "Jeremy D. Miller" <jeremydmiller@yahoo.com>
Date: Tue, 7 Jan 2025 09:02:44 -0600
Subject: [PATCH] Always putting tenant_id first in all foreign keys

---
 .../partitioning_and_foreign_keys.cs          | 34 +++++++++++++++++++
 .../Marten.CommandLine.csproj                 |  2 +-
 src/Marten/Marten.csproj                      |  2 +-
 src/Marten/Schema/DocumentMapping.cs          | 18 ++++++++++
 src/Marten/Storage/DocumentTable.cs           |  7 ++++
 5 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/src/CoreTests/Partitioning/partitioning_and_foreign_keys.cs b/src/CoreTests/Partitioning/partitioning_and_foreign_keys.cs
index f7056ee030..19917d17d7 100644
--- a/src/CoreTests/Partitioning/partitioning_and_foreign_keys.cs
+++ b/src/CoreTests/Partitioning/partitioning_and_foreign_keys.cs
@@ -1,4 +1,7 @@
+using System.Linq;
 using System.Threading.Tasks;
+using Marten.Schema;
+using Marten.Storage;
 using Marten.Testing.Documents;
 using Marten.Testing.Harness;
 using Shouldly;
@@ -63,4 +66,35 @@ await Should.ThrowAsync<InvalidForeignKeyException>(async () =>
             await theStore.Storage.ApplyAllConfiguredChangesToDatabaseAsync();
         });
     }
+
+    [Fact]
+    public async Task partitioned_by_tenant_id_to_partitioned_to_tenant_id_and_tenant_id_is_sorted_first()
+    {
+        StoreOptions(opts =>
+        {
+            opts.Schema.For<Issue>()
+                .ForeignKey<User>(x => x.AssigneeId);
+
+            opts.Schema.For<User>();
+
+            opts.Policies.AllDocumentsAreMultiTenantedWithPartitioning(partitioning =>
+            {
+                partitioning.ByHash("one", "two", "three");
+            });
+        });
+
+        // Just smoke test that it works
+        await theStore.Storage.ApplyAllConfiguredChangesToDatabaseAsync();
+
+        await theStore.Storage.Database.AssertDatabaseMatchesConfigurationAsync();
+
+        var mapping = (DocumentMapping)theStore.Options.Storage.FindMapping(typeof(Issue));
+        var table = new DocumentTable(mapping);
+
+        var fk = table.ForeignKeys.Single();
+        fk.ColumnNames.ShouldBe(["tenant_id", "assignee_id"]);
+        fk.LinkedNames.ShouldBe(["tenant_id", "id"]);
+
+        fk.Name.ShouldBe("mt_doc_issue_tenant_id_assignee_id_fkey");
+    }
 }
diff --git a/src/Marten.CommandLine/Marten.CommandLine.csproj b/src/Marten.CommandLine/Marten.CommandLine.csproj
index def02289c6..bfbd92be2c 100644
--- a/src/Marten.CommandLine/Marten.CommandLine.csproj
+++ b/src/Marten.CommandLine/Marten.CommandLine.csproj
@@ -36,7 +36,7 @@
         <PackageReference Include="JasperFx.CodeGeneration.Commands" Version="3.7.1" />
         <PackageReference Include="Microsoft.SourceLink.GitHub" Version="8.0.0" PrivateAssets="All"/>
         <PackageReference Include="Oakton" Version="6.3.0" />
-        <PackageReference Include="Weasel.CommandLine" Version="7.13.1" />
+        <PackageReference Include="Weasel.CommandLine" Version="7.13.2" />
     </ItemGroup>
     <Import Project="../../Analysis.Build.props"/>
 </Project>
diff --git a/src/Marten/Marten.csproj b/src/Marten/Marten.csproj
index 2c468cf8e1..d15feab36d 100644
--- a/src/Marten/Marten.csproj
+++ b/src/Marten/Marten.csproj
@@ -61,7 +61,7 @@
         <PackageReference Include="Polly.Core" Version="8.5.0" />
         <PackageReference Include="System.Threading.Tasks.Dataflow" Version="8.0.0" />
         <PackageReference Include="System.Text.Json" Version="8.0.5" />
-        <PackageReference Include="Weasel.Postgresql" Version="7.13.1" />
+        <PackageReference Include="Weasel.Postgresql" Version="7.13.2" />
     </ItemGroup>
 
     <!--SourceLink specific settings-->
diff --git a/src/Marten/Schema/DocumentMapping.cs b/src/Marten/Schema/DocumentMapping.cs
index 395db9a4d0..aa9cb5c692 100644
--- a/src/Marten/Schema/DocumentMapping.cs
+++ b/src/Marten/Schema/DocumentMapping.cs
@@ -15,6 +15,7 @@
 using Marten.Schema.Indexing.FullText;
 using Marten.Schema.Indexing.Unique;
 using Marten.Storage;
+using Marten.Storage.Metadata;
 using NpgsqlTypes;
 using Weasel.Core;
 using Weasel.Postgresql;
@@ -961,3 +962,20 @@ public DocumentCodeGen(DocumentMapping mapping)
     public string AccessId { get; }
     public string ParameterValue { get; }
 }
+
+internal static class ForeignKeyExtensions
+{
+    public static void TryMoveTenantIdFirst(this ForeignKey foreignKey, DocumentMapping mapping)
+    {
+        // Guard clause, do nothing if this document is not tenanted
+        if (mapping.TenancyStyle == TenancyStyle.Single) return;
+
+        foreignKey.ColumnNames = new string[] { TenantIdColumn.Name }
+            .Concat(foreignKey.ColumnNames.Where(x => x != TenantIdColumn.Name)).ToArray();
+
+        foreignKey.LinkedNames = new string[] { TenantIdColumn.Name }
+            .Concat(foreignKey.LinkedNames.Where(x => x != TenantIdColumn.Name)).ToArray();
+
+        foreignKey.Name = $"{mapping.TableName.Name}_{foreignKey.ColumnNames.Join("_")}_fkey";
+    }
+}
diff --git a/src/Marten/Storage/DocumentTable.cs b/src/Marten/Storage/DocumentTable.cs
index 720ea8e792..3f74f8fbbf 100644
--- a/src/Marten/Storage/DocumentTable.cs
+++ b/src/Marten/Storage/DocumentTable.cs
@@ -79,6 +79,13 @@ public DocumentTable(DocumentMapping mapping): base(mapping.TableName)
         }
 
         Indexes.AddRange(mapping.Indexes);
+
+        // tenant_id should always be first
+        foreach (var foreignKey in mapping.ForeignKeys)
+        {
+            foreignKey.TryMoveTenantIdFirst(mapping);
+        }
+
         ForeignKeys.AddRange(mapping.ForeignKeys);
 
         Partitioning = mapping.Partitioning;