Add support for forein keys in schema generation within aggregates

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

@@ -63,4 +63,9 @@ public DefaultSqlTypeMapping() {
 	public String getColumnType(RelationalPersistentProperty property) {
 		return typeMap.get(ClassUtils.resolvePrimitiveIfNecessary(property.getActualType()));
 	}
+
+	@Override
+	public String getColumnTypeByClass(Class clazz) {
+		return typeMap.get(clazz);
+	}
 }

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

@@ -0,0 +1,30 @@
+package org.springframework.data.jdbc.core.mapping.schema;
+
+import java.util.List;
+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, List<String> columnNames, String referencedTableName,
+									List<String> referencedColumnNames) {
+	@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(tableName, that.tableName) && Objects.equals(columnNames, that.columnNames) && Objects.equals(
+				referencedTableName, that.referencedTableName) && Objects.equals(referencedColumnNames,
+				that.referencedColumnNames);
+	}
+
+	@Override
+	public int hashCode() {
+		return Objects.hash(tableName, columnNames, referencedTableName, referencedColumnNames);
+	}
+}

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();
+			List<String> columnNames = foreignKey.getForeignKeyColumns().stream()
+					.map(liquibase.structure.core.Column::getName).toList();
+			String referencedTableName = foreignKey.getPrimaryKeyTable().getName();
+			List<String> referencedColumnNames = foreignKey.getPrimaryKeyColumns().stream()
+					.map(liquibase.structure.core.Column::getName).toList();
+			return new ForeignKey(foreignKey.getName(), tableName, columnNames, referencedTableName, referencedColumnNames);
+		}).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(String.join(",", foreignKey.columnNames()));
+		change.setReferencedTableName(foreignKey.referencedTableName());
+		change.setReferencedColumnNames(String.join(",", foreignKey.referencedColumnNames()));
+		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.
+			Collection<Column> addedColumns = 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 foreign keys
+			tableDiff.fkToDrop().addAll(findDiffs(mappedForeignKeys, existingForeignKeys, nameComparator));
+			// Identify added foreign keys
+			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/SqlTypeMapping.java

@@ -40,6 +40,18 @@ public interface SqlTypeMapping {
 	@Nullable
 	String getColumnType(RelationalPersistentProperty property);
 
+	/**
+	 * Determines a column type for Class.
+	 *
+	 * @param clazz class for which the type should be determined.
+	 * @return the SQL type to use, such as {@code VARCHAR} or {@code NUMERIC}. Can be {@literal null} if the strategy
+	 *         cannot provide a column type.
+	 */
+	@Nullable
+	default String getColumnTypeByClass(Class clazz) {
+		return null;
+	}
+
 	/**
 	 * Returns the required column type for a persistent property or throws {@link IllegalArgumentException} if the type
 	 * cannot be determined.

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

@@ -18,6 +18,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import java.util.stream.Collectors;
 import org.springframework.lang.Nullable;
 import org.springframework.util.ObjectUtils;
 
@@ -27,7 +28,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<>());
@@ -37,6 +38,10 @@ public Table(String name) {
 		this(null, name);
 	}
 
+	public List<Column> getIdColumns() {
+		return columns().stream().filter(Column::identity).collect(Collectors.toList());
+	}
+
 	@Override
 	public boolean equals(Object o) {
 

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<>());
 	}
 
 }