Issue #28: Support compound primary keys

liquibase · Feb 13, 2014 · cedfa67 · cedfa67
1 parent 5e48a4f
commit cedfa67
Showing 1 changed file with 21 additions and 7 deletions.
diff --git a/src/main/java/liquibase/ext/hibernate/snapshot/TableSnapshotGenerator.java b/src/main/java/liquibase/ext/hibernate/snapshot/TableSnapshotGenerator.java
@@ -5,19 +5,16 @@
 import liquibase.snapshot.DatabaseSnapshot;
 import liquibase.snapshot.InvalidExampleException;
 import liquibase.structure.DatabaseObject;
-import liquibase.structure.core.Column;
-import liquibase.structure.core.DataType;
-import liquibase.structure.core.Schema;
-import liquibase.structure.core.Table;
+import liquibase.structure.core.*;
 import liquibase.util.StringUtils;
 import org.hibernate.cfg.Configuration;
 import org.hibernate.dialect.Dialect;
 import org.hibernate.engine.spi.Mapping;
-import org.hibernate.mapping.PrimaryKey;
 import org.hibernate.mapping.SimpleValue;
 import org.hibernate.mapping.Value;
 
 import java.util.Iterator;
+import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -43,6 +40,8 @@ protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot
         }
 
         Table table = new Table().setName(hibernateTable.getName());
+        PrimaryKey primaryKey = null;
+        int pkColumnPosition = 0;
         LOG.info("Found table " + table.getName());
 
         table.setSchema(example.getSchema());
@@ -67,9 +66,23 @@ protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot
             column.setNullable(hibernateColumn.isNullable());
             column.setCertainDataType(false);
 
-            PrimaryKey hibernatePrimaryKey = hibernateTable.getPrimaryKey();
+            org.hibernate.mapping.PrimaryKey hibernatePrimaryKey = hibernateTable.getPrimaryKey();
             if (hibernatePrimaryKey != null) {
-                if (hibernatePrimaryKey.getColumns().size() == 1 && hibernatePrimaryKey.getColumn(0).getName().equals(hibernateColumn.getName())) {
+                boolean isPrimaryKeyColumn = false;
+                for (org.hibernate.mapping.Column pkColumn : (List<org.hibernate.mapping.Column>) hibernatePrimaryKey.getColumns()) {
+                    if (pkColumn.getName().equals(hibernateColumn.getName())) {
+                        isPrimaryKeyColumn = true;
+                        break;
+                    }
+                }
+
+                if (isPrimaryKeyColumn) {
+                    if (primaryKey == null) {
+                        primaryKey = new PrimaryKey();
+                        primaryKey.setName(table.getPrimaryKey().getName());
+                    }
+                    primaryKey.addColumnName(pkColumnPosition++, column.getName());
+
                     Value value = hibernateColumn.getValue();
                     if (value instanceof SimpleValue && ((SimpleValue) value).getIdentifierGeneratorStrategy().equals("identity")) {
                         column.setAutoIncrementInformation(new Column.AutoIncrementInformation());
@@ -78,6 +91,7 @@ protected DatabaseObject snapshotObject(DatabaseObject example, DatabaseSnapshot
             }
             column.setRelation(table);
 
+            table.setPrimaryKey(primaryKey);
             table.getColumns().add(column);
 
         }