From dc6bd251aa4bb9111f5da3c5f45a8c32a70cea99 Mon Sep 17 00:00:00 2001
From: Aaron Morton <aaron.morton@datastax.com>
Date: Fri, 15 Nov 2024 08:13:09 +1300
Subject: [PATCH] Ajm/fix 1720 bad order inserted return (#1722)

Co-authored-by: Mahesh Rajamani <99678631+maheshrajamani@users.noreply.github.com>
---
 .../operation/tables/TableInsertAttempt.java  |  3 +--
 .../tables/WriteableTableRowBuilder.java      | 19 ++++++++++++++++---
 .../schema/tables/ApiColumnDefContainer.java  |  2 +-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableInsertAttempt.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableInsertAttempt.java
index fc380a223..83fdc0e5e 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableInsertAttempt.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableInsertAttempt.java
@@ -44,8 +44,7 @@ public Optional<ColumnsDescContainer> schemaDescription() {
       return Optional.empty();
     }
 
-    var apiColumns =
-        schemaObject.apiTableDef().allColumns().filterBy(row.keyColumns().getIdentifiers());
+    var apiColumns = schemaObject.apiTableDef().primaryKeys();
     if (!apiColumns.filterByUnsupported().isEmpty()) {
       throw new IllegalStateException(
           "Unsupported columns primary key: %s" + apiColumns.filterByUnsupported());
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/WriteableTableRowBuilder.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/WriteableTableRowBuilder.java
index a7548a13e..8be6e81a7 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/WriteableTableRowBuilder.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/WriteableTableRowBuilder.java
@@ -79,13 +79,26 @@ public WriteableTableRow build(JsonNamedValueContainer source) {
     // now need to split the columns into key and non key columns
     var keyColumns = new CqlNamedValueContainer();
     var nonKeyColumns = new CqlNamedValueContainer();
-    for (var cqlNamedValue : decoded.values()) {
-      if (tableMetadata.getPrimaryKey().contains(cqlNamedValue.name())) {
-        keyColumns.put(cqlNamedValue);
+
+    // Get the primary keys out of the new values in the order they are defined on the table
+    for (var keyMetadata : tableMetadata.getPrimaryKey()) {
+      if (decoded.containsKey(keyMetadata)) {
+        keyColumns.put(decoded.get(keyMetadata));
       } else {
+        // the primary keys have been checked above
+        throw new IllegalStateException(
+            String.format(
+                "build: primary key column not found in decoded values, column=%s", keyMetadata));
+      }
+    }
+
+    // any column in decoded that is now not in keyColumns is a non key column
+    for (var cqlNamedValue : decoded.values()) {
+      if (!keyColumns.containsKey(cqlNamedValue.name())) {
         nonKeyColumns.put(cqlNamedValue);
       }
     }
+
     return new WriteableTableRow(tableSchemaObject, keyColumns, nonKeyColumns);
   }
 
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/schema/tables/ApiColumnDefContainer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/schema/tables/ApiColumnDefContainer.java
index 55b7537f0..20e812ed4 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/schema/tables/ApiColumnDefContainer.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/schema/tables/ApiColumnDefContainer.java
@@ -106,7 +106,7 @@ public ApiColumnDefContainer filterBy(ApiTypeName type) {
 
   public ApiColumnDefContainer filterBy(Collection<CqlIdentifier> identifiers) {
     return new ApiColumnDefContainer(
-        identifiers.stream().map(this::get).filter(Objects::nonNull).toList());
+        values().stream().filter(columnDef -> identifiers.contains(columnDef.name())).toList());
   }
 
   public ApiColumnDefContainer filterByUnsupported() {