Merge branch 'josecollazzi-master'

aaberg · Mar 14, 2024 · b69a570 · b69a570
2 parents a4164c0 + 0036563
commit b69a570
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git 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
@@ -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() {