Add support for forein keys in schema generation within aggregates

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/ForeignKey.java

@@ -0,0 +1,27 @@
+package org.springframework.data.jdbc.core.mapping.schema;
+
+import java.util.Objects;
+
+/**
+ * Models a Foreign Key for generating SQL for Schema generation.
+ *
+ * @author Evgenii Koba
+ * @since 3.2
+ */
+record ForeignKey(String name, String tableName, String columnName, String referencedTableName,
+									String referencedColumnName) {
+	@Override
+	public boolean equals(Object o) {
+		if (this == o)
+			return true;
+		if (o == null || getClass() != o.getClass())
+			return false;
+		ForeignKey that = (ForeignKey) o;
+		return Objects.equals(name, that.name);
+	}
+
+	@Override
+	public int hashCode() {
+		return Objects.hash(name);
+	}
+}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/LiquibaseChangeSetWriter.java

@@ -21,8 +21,10 @@
 import liquibase.change.ColumnConfig;
 import liquibase.change.ConstraintsConfig;
 import liquibase.change.core.AddColumnChange;
+import liquibase.change.core.AddForeignKeyConstraintChange;
 import liquibase.change.core.CreateTableChange;
 import liquibase.change.core.DropColumnChange;
+import liquibase.change.core.DropForeignKeyConstraintChange;
 import liquibase.change.core.DropTableChange;
 import liquibase.changelog.ChangeLogChild;
 import liquibase.changelog.ChangeLogParameters;
@@ -52,6 +54,7 @@
 import java.util.Set;
 import java.util.function.BiPredicate;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import org.springframework.core.io.Resource;
 import org.springframework.data.mapping.context.MappingContext;
@@ -321,15 +324,15 @@ private ChangeSet createChangeSet(ChangeSetMetadata metadata, SchemaDiff differe
 	private SchemaDiff initial() {
 
 		Tables mappedEntities = Tables.from(mappingContext.getPersistentEntities().stream().filter(schemaFilter),
-				sqlTypeMapping, null);
+				sqlTypeMapping, null, mappingContext);
 		return SchemaDiff.diff(mappedEntities, Tables.empty(), nameComparator);
 	}
 
 	private SchemaDiff differenceOf(Database database) throws LiquibaseException {
 
 		Tables existingTables = getLiquibaseModel(database);
 		Tables mappedEntities = Tables.from(mappingContext.getPersistentEntities().stream().filter(schemaFilter),
-				sqlTypeMapping, database.getDefaultCatalogName());
+				sqlTypeMapping, database.getDefaultCatalogName(), mappingContext);
 
 		return SchemaDiff.diff(mappedEntities, existingTables, nameComparator);
 	}
@@ -362,6 +365,13 @@ private DatabaseChangeLog getDatabaseChangeLog(File changeLogFile, @Nullable Dat
 
 	private void generateTableAdditionsDeletions(ChangeSet changeSet, SchemaDiff difference) {
 
+		for (Table table : difference.tableDeletions()) {
+			for (ForeignKey foreignKey : table.foreignKeys()) {
+				DropForeignKeyConstraintChange dropForeignKey = dropForeignKey(foreignKey);
+				changeSet.addChange(dropForeignKey);
+			}
+		}
+
 		for (Table table : difference.tableAdditions()) {
 			CreateTableChange newTable = changeTable(table);
 			changeSet.addChange(newTable);
@@ -373,12 +383,24 @@ private void generateTableAdditionsDeletions(ChangeSet changeSet, SchemaDiff dif
 				changeSet.addChange(dropTable(table));
 			}
 		}
+
+		for (Table table : difference.tableAdditions()) {
+			for (ForeignKey foreignKey : table.foreignKeys()) {
+				AddForeignKeyConstraintChange addForeignKey = addForeignKey(foreignKey);
+				changeSet.addChange(addForeignKey);
+			}
+		}
 	}
 
 	private void generateTableModifications(ChangeSet changeSet, SchemaDiff difference) {
 
 		for (TableDiff table : difference.tableDiffs()) {
 
+			for (ForeignKey foreignKey : table.fkToDrop()) {
+				DropForeignKeyConstraintChange dropForeignKey = dropForeignKey(foreignKey);
+				changeSet.addChange(dropForeignKey);
+			}
+
 			if (!table.columnsToAdd().isEmpty()) {
 				changeSet.addChange(addColumns(table));
 			}
@@ -388,6 +410,11 @@ private void generateTableModifications(ChangeSet changeSet, SchemaDiff differen
 			if (!deletedColumns.isEmpty()) {
 				changeSet.addChange(dropColumns(table, deletedColumns));
 			}
+
+			for (ForeignKey foreignKey : table.fkToAdd()) {
+				AddForeignKeyConstraintChange addForeignKey = addForeignKey(foreignKey);
+				changeSet.addChange(addForeignKey);
+			}
 		}
 	}
 
@@ -444,12 +471,27 @@ private Tables getLiquibaseModel(Database targetDatabase) throws LiquibaseExcept
 				tableModel.columns().add(columnModel);
 			}
 
+			tableModel.foreignKeys().addAll(extractForeignKeys(table));
+
 			existingTables.add(tableModel);
 		}
 
 		return new Tables(existingTables);
 	}
 
+	private static List<ForeignKey> extractForeignKeys(liquibase.structure.core.Table table) {
+
+		return table.getOutgoingForeignKeys().stream().map(foreignKey -> {
+			String tableName = foreignKey.getForeignKeyTable().getName();
+			String columnName = foreignKey.getForeignKeyColumns().stream().findFirst()
+					.map(liquibase.structure.core.Column::getName).get();
+			String referencedTableName = foreignKey.getPrimaryKeyTable().getName();
+			String referencedColumnName = foreignKey.getPrimaryKeyColumns().stream().findFirst()
+					.map(liquibase.structure.core.Column::getName).get();
+			return new ForeignKey(foreignKey.getName(), tableName, columnName, referencedTableName, referencedColumnName);
+		}).collect(Collectors.toList());
+	}
+
 	private static AddColumnChange addColumns(TableDiff table) {
 
 		AddColumnChange addColumnChange = new AddColumnChange();
@@ -532,6 +574,25 @@ private static DropTableChange dropTable(Table table) {
 		return change;
 	}
 
+	private static AddForeignKeyConstraintChange addForeignKey(ForeignKey foreignKey) {
+
+		AddForeignKeyConstraintChange change = new AddForeignKeyConstraintChange();
+		change.setConstraintName(foreignKey.name());
+		change.setBaseTableName(foreignKey.tableName());
+		change.setBaseColumnNames(foreignKey.columnName());
+		change.setReferencedTableName(foreignKey.referencedTableName());
+		change.setReferencedColumnNames(foreignKey.referencedColumnName());
+		return change;
+	}
+
+	private static DropForeignKeyConstraintChange dropForeignKey(ForeignKey foreignKey) {
+
+		DropForeignKeyConstraintChange change = new DropForeignKeyConstraintChange();
+		change.setConstraintName(foreignKey.name());
+		change.setBaseTableName(foreignKey.tableName());
+		return change;
+	}
+
 	/**
 	 * Metadata for a ChangeSet.
 	 */

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/SchemaDiff.java

@@ -16,6 +16,7 @@
 package org.springframework.data.jdbc.core.mapping.schema;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
@@ -91,43 +92,40 @@ private static List<TableDiff> diffTable(Tables mappedEntities, Map<String, Tabl
 			TableDiff tableDiff = new TableDiff(mappedEntity);
 
 			Map<String, Column> mappedColumns = createMapping(mappedEntity.columns(), Column::name, nameComparator);
-			mappedEntity.keyColumns().forEach(it -> mappedColumns.put(it.name(), it));
-
 			Map<String, Column> existingColumns = createMapping(existingTable.columns(), Column::name, nameComparator);
-			existingTable.keyColumns().forEach(it -> existingColumns.put(it.name(), it));
-
 			// Identify deleted columns
-			Map<String, Column> toDelete = new TreeMap<>(nameComparator);
-			toDelete.putAll(existingColumns);
-			mappedColumns.keySet().forEach(toDelete::remove);
-
-			tableDiff.columnsToDrop().addAll(toDelete.values());
-
-			// Identify added columns
-			Map<String, Column> addedColumns = new TreeMap<>(nameComparator);
-			addedColumns.putAll(mappedColumns);
-
-			existingColumns.keySet().forEach(addedColumns::remove);
-
-			// Add columns in order. This order can interleave with existing columns.
-			for (Column column : mappedEntity.keyColumns()) {
-				if (addedColumns.containsKey(column.name())) {
-					tableDiff.columnsToAdd().add(column);
-				}
-			}
-
+			tableDiff.columnsToDrop().addAll(findDiffs(mappedColumns, existingColumns, nameComparator));
+			// Identify added columns and add columns in order. This order can interleave with existing columns.
+			List<Column> addedColumns = new ArrayList<>(findDiffs(existingColumns, mappedColumns, nameComparator));
 			for (Column column : mappedEntity.columns()) {
-				if (addedColumns.containsKey(column.name())) {
+				if (addedColumns.contains(column)) {
 					tableDiff.columnsToAdd().add(column);
 				}
 			}
 
+			Map<String, ForeignKey> mappedForeignKeys = createMapping(mappedEntity.foreignKeys(), ForeignKey::name,
+					nameComparator);
+			Map<String, ForeignKey> existingForeignKeys = createMapping(existingTable.foreignKeys(), ForeignKey::name,
+					nameComparator);
+			// Identify deleted columns
+			tableDiff.fkToDrop().addAll(findDiffs(mappedForeignKeys, existingForeignKeys, nameComparator));
+			// Identify added columns
+			tableDiff.fkToAdd().addAll(findDiffs(existingForeignKeys, mappedForeignKeys, nameComparator));
+
 			tableDiffs.add(tableDiff);
 		}
 
 		return tableDiffs;
 	}
 
+	private static <T> Collection<T> findDiffs(Map<String, T> baseMapping, Map<String, T> toCompareMapping,
+			Comparator<String> nameComparator) {
+		Map<String, T> diff = new TreeMap<>(nameComparator);
+		diff.putAll(toCompareMapping);
+		baseMapping.keySet().forEach(diff::remove);
+		return diff.values();
+	}
+
 	private static <T> SortedMap<String, T> createMapping(List<T> items, Function<T, String> keyFunction,
 			Comparator<String> nameComparator) {
 

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/Table.java

@@ -27,7 +27,7 @@
  * @author Kurt Niemi
  * @since 3.2
  */
-record Table(@Nullable String schema, String name, List<Column> keyColumns, List<Column> columns) {
+record Table(@Nullable String schema, String name, List<Column> columns, List<ForeignKey> foreignKeys) {
 
 	public Table(@Nullable String schema, String name) {
 		this(schema, name, new ArrayList<>(), new ArrayList<>());

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/TableDiff.java

@@ -25,10 +25,11 @@
  * @author Kurt Niemi
  * @since 3.2
  */
-record TableDiff(Table table, List<Column> columnsToAdd, List<Column> columnsToDrop) {
+record TableDiff(Table table, List<Column> columnsToAdd, List<Column> columnsToDrop, List<ForeignKey> fkToAdd,
+								 List<ForeignKey> fkToDrop) {
 
 	public TableDiff(Table table) {
-		this(table, new ArrayList<>(), new ArrayList<>());
+		this(table, new ArrayList<>(), new ArrayList<>(), new ArrayList<>(), new ArrayList<>());
 	}
 
 }

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/mapping/schema/Tables.java

@@ -15,14 +15,19 @@
  */
 package org.springframework.data.jdbc.core.mapping.schema;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import org.springframework.data.annotation.Id;
+import org.springframework.data.mapping.context.MappingContext;
+import org.springframework.data.relational.core.mapping.MappedCollection;
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
@@ -37,15 +42,16 @@
 record Tables(List<Table> tables) {
 
 	public static Tables from(RelationalMappingContext context) {
-		return from(context.getPersistentEntities().stream(), new DefaultSqlTypeMapping(), null);
+		return from(context.getPersistentEntities().stream(), new DefaultSqlTypeMapping(), null, context);
 	}
 
-	// TODO: Add support (i.e. create tickets) to support mapped collections, entities, embedded properties, and aggregate
-	// references.
+	// TODO: Add support (i.e. create tickets) to support entities, embedded properties, and aggregate references.
 
 	public static Tables from(Stream<? extends RelationalPersistentEntity<?>> persistentEntities,
-			SqlTypeMapping sqlTypeMapping, @Nullable String defaultSchema) {
+			SqlTypeMapping sqlTypeMapping, @Nullable String defaultSchema,
+			MappingContext<? extends RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> context) {
 
+		Map<String, List<ColumnWithForeignKey>> colAndFKByTableName = new HashMap<>();
 		List<Table> tables = persistentEntities
 				.filter(it -> it.isAnnotationPresent(org.springframework.data.relational.core.mapping.Table.class)) //
 				.map(entity -> {
@@ -54,26 +60,85 @@ public static Tables from(Stream<? extends RelationalPersistentEntity<?>> persis
 
 					Set<RelationalPersistentProperty> identifierColumns = new LinkedHashSet<>();
 					entity.getPersistentProperties(Id.class).forEach(identifierColumns::add);
+					collectForeignKeysInfo(entity, context, colAndFKByTableName, sqlTypeMapping);
 
 					for (RelationalPersistentProperty property : entity) {
 
 						if (property.isEntity() && !property.isEmbedded()) {
 							continue;
 						}
 
-						String columnType = sqlTypeMapping.getRequiredColumnType(property);
-
 						Column column = new Column(property.getColumnName().getReference(), sqlTypeMapping.getColumnType(property),
-								sqlTypeMapping.isNullable(property), identifierColumns.contains(property));
+												   sqlTypeMapping.isNullable(property), identifierColumns.contains(property));
 						table.columns().add(column);
 					}
 					return table;
 				}).collect(Collectors.toList());
 
+		applyForeignKeys(tables, colAndFKByTableName);
+
 		return new Tables(tables);
 	}
 
 	public static Tables empty() {
 		return new Tables(Collections.emptyList());
 	}
+
+	private static void applyForeignKeys(List<Table> tables,
+			Map<String, List<ColumnWithForeignKey>> colAndFKByTableName) {
+
+		colAndFKByTableName.forEach(
+				(tableName, colsAndFK) -> tables.stream().filter(table -> table.name().equals(tableName)).forEach(table -> {
+
+					colsAndFK.forEach(colAndFK -> {
+						if (!table.columns().contains(colAndFK.column())) {
+							table.columns().add(colAndFK.column());
+						}
+					});
+
+					colsAndFK.forEach(colAndFK -> table.foreignKeys().add(colAndFK.foreignKey()));
+				}));
+	}
+
+	private static void collectForeignKeysInfo(RelationalPersistentEntity<?> entity,
+			MappingContext<? extends RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> context,
+			Map<String, List<ColumnWithForeignKey>> keyColumnsByTableName, SqlTypeMapping sqlTypeMapping) {
+
+		RelationalPersistentProperty identifierColumn = entity.getPersistentProperty(Id.class);
+
+		entity.getPersistentProperties(MappedCollection.class).forEach(property -> {
+			if (property.isEntity()) {
+				property.getPersistentEntityTypeInformation().forEach(typeInformation -> {
+
+					String tableName = context.getRequiredPersistentEntity(typeInformation).getTableName().getReference();
+					String columnName = property.getReverseColumnName(entity).getReference();
+					String referencedTableName = entity.getTableName().getReference();
+					String referencedColumnName = identifierColumn.getColumnName().getReference();
+
+					ForeignKey foreignKey = new ForeignKey(getForeignKeyName(referencedTableName, referencedColumnName),
+							tableName, columnName, referencedTableName, referencedColumnName);
+					Column column = new Column(columnName, sqlTypeMapping.getColumnType(identifierColumn), true, false);
+
+					ColumnWithForeignKey columnWithForeignKey = new ColumnWithForeignKey(column, foreignKey);
+					keyColumnsByTableName.compute(
+							context.getRequiredPersistentEntity(typeInformation).getTableName().getReference(), (key, value) -> {
+								if (value == null) {
+									return new ArrayList<>(List.of(columnWithForeignKey));
+								} else {
+									value.add(columnWithForeignKey);
+									return value;
+								}
+							});
+				});
+			}
+		});
+	}
+
+	//TODO should we place it in BasicRelationalPersistentProperty/BasicRelationalPersistentEntity and generate using NamingStrategy?
+	private static String getForeignKeyName(String referencedTableName, String referencedColumnName) {
+		return String.format("%s_%s_fk", referencedTableName, referencedColumnName);
+	}
+
+	private record ColumnWithForeignKey(Column column, ForeignKey foreignKey) {
+	}
 }