From 79aad5fc5915804ad3c855bbcb585a2163faa20d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20Westk=C3=A4mper?= <timo.westkamper@gmail.com>
Date: Sat, 20 Sep 2014 21:51:20 +0300
Subject: [PATCH] Better key handling

---
 .../mysema/query/sql/dml/SQLMergeClause.java   | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/querydsl-sql/src/main/java/com/mysema/query/sql/dml/SQLMergeClause.java b/querydsl-sql/src/main/java/com/mysema/query/sql/dml/SQLMergeClause.java
index 1a733e9e5..84a474bdb 100644
--- a/querydsl-sql/src/main/java/com/mysema/query/sql/dml/SQLMergeClause.java
+++ b/querydsl-sql/src/main/java/com/mysema/query/sql/dml/SQLMergeClause.java
@@ -97,6 +97,16 @@ public class SQLMergeClause extends AbstractSQLClause<SQLMergeClause> implements
         metadata.addFlag(new QueryFlag(position, flag));
         return this;
     }
+    
+    private List<? extends Path<?>> getKeys() {
+        if (!keys.isEmpty()) {
+            return keys;
+        } else if (entity.getPrimaryKey() != null) {
+            return entity.getPrimaryKey().getLocalColumns();
+        } else {
+            throw new IllegalStateException("No keys where defined, invoke keys(..) to add keys");
+        }
+    }
 
     /**
      * Add the current state of bindings as a batch item
@@ -243,7 +253,7 @@ public class SQLMergeClause extends AbstractSQLClause<SQLMergeClause> implements
                     // update
                     SQLUpdateClause update = new SQLUpdateClause(connection, configuration.getTemplates(), entity);
                     populate(update);
-                    update.where(ExpressionUtils.in((Expression)keys.get(0),ids));
+                    update.where(ExpressionUtils.in((Expression)getKeys().get(0),ids));
                     return EmptyResultSet.DEFAULT;
                 } else {
                     // insert
@@ -296,7 +306,7 @@ public class SQLMergeClause extends AbstractSQLClause<SQLMergeClause> implements
                 query.where(ExpressionUtils.eq(columns.get(i),(Expression)values.get(i)));
             }
         }
-        return query.list(keys.get(0));
+        return query.list(getKeys().get(0));
     }
 
     @SuppressWarnings("unchecked")
@@ -306,7 +316,7 @@ public class SQLMergeClause extends AbstractSQLClause<SQLMergeClause> implements
             // update
             SQLUpdateClause update = new SQLUpdateClause(connection, configuration.getTemplates(), entity);
             populate(update);
-            update.where(ExpressionUtils.in((Expression)keys.get(0),ids));
+            update.where(ExpressionUtils.in((Expression)getKeys().get(0),ids));
             return update.execute();
         } else {
             // insert
@@ -411,7 +421,7 @@ public class SQLMergeClause extends AbstractSQLClause<SQLMergeClause> implements
         if (withKeys) {
             String[] target = new String[keys.size()];
             for (int i = 0; i < target.length; i++) {
-                target[i] = ColumnMetadata.getName(keys.get(i));
+                target[i] = ColumnMetadata.getName(getKeys().get(i));
             }
             stmt = connection.prepareStatement(queryString, target);
         } else {