From c3fc7df332b827a156934ccf94944f0754c3019c Mon Sep 17 00:00:00 2001 From: "Jeremy D. Miller" Date: Mon, 23 Dec 2024 12:11:26 -0600 Subject: [PATCH] Fixes for migrating partitions with foreign keys in PostgreSQL --- src/Weasel.Core/ISchemaObject.cs | 9 ++++++ src/Weasel.Core/Migrations/DatabaseBase.cs | 34 ++++++++++++++++++++-- src/Weasel.Postgresql/Tables/ForeignKey.cs | 31 ++++++++++++++++++++ src/Weasel.Postgresql/Tables/Table.cs | 17 ++++++++++- 4 files changed, 87 insertions(+), 4 deletions(-) diff --git a/src/Weasel.Core/ISchemaObject.cs b/src/Weasel.Core/ISchemaObject.cs index 9087da5b..b08d5f96 100644 --- a/src/Weasel.Core/ISchemaObject.cs +++ b/src/Weasel.Core/ISchemaObject.cs @@ -19,6 +19,15 @@ public SchemaMigrationException(AutoCreate autoCreate, IEnumerable inval } } +/// +/// 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 +/// +public interface ISchemaObjectWithPostProcessing : ISchemaObject +{ + void PostProcess(ISchemaObject[] allObjects); +} + /// /// Responsible for the desired configuration of a single database object like /// a table, sequence, of function. diff --git a/src/Weasel.Core/Migrations/DatabaseBase.cs b/src/Weasel.Core/Migrations/DatabaseBase.cs index e8259c20..37d64de7 100644 --- a/src/Weasel.Core/Migrations/DatabaseBase.cs +++ b/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().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) /// 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 CreateMigrationAsync(CancellationToken ct = default) { + applyPostProcessingIfAny(); var objects = AllObjects().ToArray(); await using var conn = CreateConnection(); @@ -261,6 +275,11 @@ public async Task ApplyAllConfiguredChangesToDatabaseAsyn Migrator.AssertValidIdentifier(objectName.Name); } + foreach (var postProcessing in objects.OfType().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().ToArray()) + { + processing.PostProcess(allObjects); + } + } + using (await _migrateLocker.Lock(5.Seconds(), token).ConfigureAwait(false)) { if (_checks.ContainsKey(featureType)) diff --git a/src/Weasel.Postgresql/Tables/ForeignKey.cs b/src/Weasel.Postgresql/Tables/ForeignKey.cs index c4e92112..8ce80357 100644 --- a/src/Weasel.Postgresql/Tables/ForeignKey.cs +++ b/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) + { + } } diff --git a/src/Weasel.Postgresql/Tables/Table.cs b/src/Weasel.Postgresql/Tables/Table.cs index 35796666..bd5887f4 100644 --- a/src/Weasel.Postgresql/Tables/Table.cs +++ b/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 _columns = new(); @@ -163,6 +164,20 @@ public IEnumerable 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().FirstOrDefault(x => x.Identifier.Equals(key.LinkedTable)); + if (matching != null) + { + key.TryToCorrectForLink(this, matching); + } + } + } + /// /// Mutate this table to change the identifier to being in a different schema ///