From 00365630f609d69012d02bcc919f87f2e420fa71 Mon Sep 17 00:00:00 2001
From: jose <jose.collazzi@manywho.com>
Date: Wed, 2 Nov 2016 15:07:00 +0000
Subject: [PATCH] resolve #249 getKey failing in postgresql

---
 core/src/main/java/org/sql2o/Connection.java  |  4 ++-
 .../extensions/postgres/PostgresTest.java     | 26 +++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/sql2o/Connection.java b/core/src/main/java/org/sql2o/Connection.java
index 12caa6c1..6c96afb1 100644
--- a/core/src/main/java/org/sql2o/Connection.java
+++ b/core/src/main/java/org/sql2o/Connection.java
@@ -187,7 +187,9 @@ void setKeys(ResultSet rs) throws SQLException {
         }
         this.keys = new ArrayList<Object>();
         while(rs.next()){
-            this.keys.add(rs.getObject(1));
+            for (int column=1; column <= rs.getMetaData().getColumnCount(); column++) {
+                this.keys.add(rs.getObject(column));
+            }
         }
     }
 
diff --git a/extensions/postgres/src/test/java/org/sql2o/extensions/postgres/PostgresTest.java b/extensions/postgres/src/test/java/org/sql2o/extensions/postgres/PostgresTest.java
index de623380..b5afb7f0 100644
--- a/extensions/postgres/src/test/java/org/sql2o/extensions/postgres/PostgresTest.java
+++ b/extensions/postgres/src/test/java/org/sql2o/extensions/postgres/PostgresTest.java
@@ -148,6 +148,32 @@ public void testKeyKeyOnSerial() {
         }
     }
 
+    @Test
+    public void testKeysKeyOnSerial() {
+        Connection connection = null;
+
+        try {
+            connection = sql2o.beginTransaction();
+
+            String createTableSql = "create table test_serial_table (val varchar(20), id serial primary key)";
+            connection.createQuery(createTableSql).executeUpdate();
+
+            String insertSql = "insert into test_serial_table(val) values ('something')";
+            Object[] key = connection.createQuery(insertSql, true).executeUpdate().getKeys();
+
+            assertThat((String) key[0], equalTo("something"));
+            assertThat((Integer) key[1], equalTo(1));
+
+            key = connection.createQuery(insertSql, true).executeUpdate().getKeys();
+            assertThat((String) key[0], equalTo("something"));
+            assertThat((Integer) key[1], equalTo(2));
+        } finally {
+            if (connection != null) {
+                connection.rollback();
+            }
+        }
+    }
+
     @Test
     public void testUUID() {