Fixes for migrating partitions with foreign keys in PostgreSQL

src/Weasel.Core/ISchemaObject.cs

@@ -19,6 +19,15 @@ public SchemaMigrationException(AutoCreate autoCreate, IEnumerable<object> inval
     }
 }
 
+/// <summary>
+/// Schema objects that may need to analyze other schema objects in order to correctly
+/// generate their own model. Originally introduced for PostgreSQL partitioning with foreign keys
+/// </summary>
+public interface ISchemaObjectWithPostProcessing : ISchemaObject
+{
+    void PostProcess(ISchemaObject[] allObjects);
+}
+
 /// <summary>
 ///     Responsible for the desired configuration of a single database object like
 ///     a table, sequence, of function.

src/Weasel.Core/Migrations/DatabaseBase.cs

@@ -113,6 +113,8 @@ public string ToDatabaseScript()
             .ToArray();
         var writer = new StringWriter();
 
+        applyPostProcessingIfAny();
+
         Migrator.WriteScript(writer, (m, w) =>
         {
             m.WriteSchemaCreationSql(schemaNames, writer);
@@ -123,6 +125,15 @@ public string ToDatabaseScript()
         return writer.ToString();
     }
 
+    private void applyPostProcessingIfAny()
+    {
+        var objects = AllObjects().ToArray();
+        foreach (var postProcessing in objects.OfType<ISchemaObjectWithPostProcessing>().ToArray())
+        {
+            postProcessing.PostProcess(objects);
+        }
+    }
+
     public async Task WriteCreationScriptToFileAsync(string filename, CancellationToken ct = default)
     {
         var directory = Path.GetDirectoryName(filename);
@@ -152,6 +163,8 @@ private async Task initializeSchemaWithNewConnection(CancellationToken ct)
     /// <param name="directory"></param>
     public async Task WriteScriptsByTypeAsync(string directory, CancellationToken ct = default)
     {
+        applyPostProcessingIfAny();
+
         FileSystem.CleanDirectory(directory);
 
         await initializeSchemaWithNewConnection(ct).ConfigureAwait(false);
@@ -195,6 +208,7 @@ await Migrator.WriteTemplatedFile(directory.AppendPath(scriptName), (m, w) =>
 
     public async Task<SchemaMigration> CreateMigrationAsync(CancellationToken ct = default)
     {
+        applyPostProcessingIfAny();
         var objects = AllObjects().ToArray();
 
         await using var conn = CreateConnection();
@@ -261,6 +275,11 @@ public async Task<SchemaPatchDifference> ApplyAllConfiguredChangesToDatabaseAsyn
             Migrator.AssertValidIdentifier(objectName.Name);
         }
 
+        foreach (var postProcessing in objects.OfType<ISchemaObjectWithPostProcessing>().ToArray())
+        {
+            postProcessing.PostProcess(objects);
+        }
+
         TConnection? conn = null;
         try
         {
@@ -400,10 +419,11 @@ private async ValueTask ensureStorageExistsAsync(
             await ensureStorageExistsAsync(types, dependentType, token).ConfigureAwait(false);
         }
 
-        await generateOrUpdateFeature(featureType, feature, token).ConfigureAwait(false);
+        await generateOrUpdateFeature(featureType, feature, token, false).ConfigureAwait(false);
     }
 
-    protected async ValueTask generateOrUpdateFeature(Type featureType, IFeatureSchema feature, CancellationToken token)
+    protected async ValueTask generateOrUpdateFeature(Type featureType, IFeatureSchema feature, CancellationToken token,
+        bool skipPostProcessing)
     {
         if (_checks.ContainsKey(featureType))
         {
@@ -413,12 +433,20 @@ protected async ValueTask generateOrUpdateFeature(Type featureType, IFeatureSche
 
         var schemaObjects = feature.Objects;
 
-
         foreach (var objectName in schemaObjects.SelectMany(x => x.AllNames()))
         {
             Migrator.AssertValidIdentifier(objectName.Name);
         }
 
+        if (!skipPostProcessing)
+        {
+            var allObjects = AllObjects().ToArray();
+            foreach (var processing in schemaObjects.OfType<ISchemaObjectWithPostProcessing>().ToArray())
+            {
+                processing.PostProcess(allObjects);
+            }
+        }
+
         using (await _migrateLocker.Lock(5.Seconds(), token).ConfigureAwait(false))
         {
             if (_checks.ContainsKey(featureType))

src/Weasel.Postgresql/Tables/ForeignKey.cs

@@ -156,4 +156,35 @@ public void WriteDropStatement(Table parent, TextWriter writer)
     {
         writer.WriteLine($"ALTER TABLE {parent.Identifier} DROP CONSTRAINT IF EXISTS {Name};");
     }
+
+    public void TryToCorrectForLink(Table parentTable, Table linkedTable)
+    {
+        // Depends on "id" always being first in Marten world
+        LinkedNames = linkedTable.PrimaryKeyColumns.ToArray();
+        if (ColumnNames.Length != LinkedNames.Length)
+        {
+            // Leave the first column alone!
+            for (int i = 1; i < LinkedNames.Length; i++)
+            {
+                var columnName = LinkedNames[i];
+                var matching = parentTable.ColumnFor(columnName);
+                if (matching != null)
+                {
+                    ColumnNames = ColumnNames.Concat([columnName]).ToArray();
+                }
+                else
+                {
+                    throw new InvalidForeignKeyException(
+                        $"Cannot make a foreign key relationship from {parentTable.Identifier}({ColumnNames.Join(", ")}) to {linkedTable.Identifier}({LinkedNames.Join(", ")}) ");
+                }
+            }
+        }
+    }
+}
+
+public class InvalidForeignKeyException: Exception
+{
+    public InvalidForeignKeyException(string? message) : base(message)
+    {
+    }
 }

src/Weasel.Postgresql/Tables/Table.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics;
 using System.Globalization;
 using JasperFx.Core;
 using Npgsql;
@@ -7,7 +8,7 @@
 
 namespace Weasel.Postgresql.Tables;
 
-public partial class Table: ISchemaObject
+public partial class Table: ISchemaObjectWithPostProcessing
 {
     private readonly List<TableColumn> _columns = new();
 
@@ -163,6 +164,20 @@ public IEnumerable<DbObjectName> AllNames()
         foreach (var fk in ForeignKeys) yield return new PostgresqlObjectName(Identifier.Schema, fk.Name);
     }
 
+    public void PostProcess(ISchemaObject[] allObjects)
+    {
+        if (!ForeignKeys.Any()) return;
+
+        foreach (var key in ForeignKeys)
+        {
+            var matching = allObjects.OfType<Table>().FirstOrDefault(x => x.Identifier.Equals(key.LinkedTable));
+            if (matching != null)
+            {
+                key.TryToCorrectForLink(this, matching);
+            }
+        }
+    }
+
     /// <summary>
     ///     Mutate this table to change the identifier to being in a different schema
     /// </summary>