From b02d1bb11b00de0aac23543b719ce08b0ec5ec70 Mon Sep 17 00:00:00 2001
From: Jonathan Schneider <jkschneider@gmail.com>
Date: Fri, 3 Apr 2026 15:58:53 -0400
Subject: [PATCH] Fix InMemoryDataTableStore.getRows returning only first
 matching bucket

getRows returned early on the first (name, group) match instead of
concatenating rows from all matching buckets. This caused ungrouped
tables with different instance names to lose rows.
---
 .../org/openrewrite/InMemoryDataTableStore.java | 16 +++++-----------
 .../org/openrewrite/DataTableStoreTest.java     | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java b/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java
index 5ce6612fa4..ab55da41c5 100644
--- a/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java
+++ b/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java
@@ -60,22 +60,16 @@ public <Row> void insertRow(DataTable<Row> dataTable, ExecutionContext ctx, Row
 
     @Override
     public Stream<?> getRows(String dataTableName, @Nullable String group) {
-        // Scan for matching bucket
+        List<Object> allRows = new ArrayList<>();
         for (Bucket bucket : buckets.values()) {
             if (bucket.dataTable.getName().equals(dataTableName) &&
                 java.util.Objects.equals(bucket.dataTable.getGroup(), group)) {
-                return snapshotRows(bucket);
+                synchronized (bucket.rows) {
+                    allRows.addAll(bucket.rows);
+                }
             }
         }
-        return Stream.empty();
-    }
-
-    private static Stream<Object> snapshotRows(Bucket bucket) {
-        List<Object> snapshot;
-        synchronized (bucket.rows) {
-            snapshot = new ArrayList<>(bucket.rows);
-        }
-        return snapshot.stream();
+        return allRows.stream();
     }
 
     @Override
diff --git a/rewrite-core/src/test/java/org/openrewrite/DataTableStoreTest.java b/rewrite-core/src/test/java/org/openrewrite/DataTableStoreTest.java
index 4728adba2b..c3c54ec812 100644
--- a/rewrite-core/src/test/java/org/openrewrite/DataTableStoreTest.java
+++ b/rewrite-core/src/test/java/org/openrewrite/DataTableStoreTest.java
@@ -114,6 +114,23 @@ void ungroupedTablesGetSeparateBuckets() {
         assertThat(store.getDataTables()).hasSize(2);
     }
 
+    @Test
+    void getRowsAggregatesAcrossUngroupedBuckets() {
+        InMemoryDataTableStore store = new InMemoryDataTableStore();
+        ExecutionContext ctx = ctx();
+
+        // Two ungrouped tables with different instance names end up in separate buckets
+        TestTable table1 = new TestTable(Recipe.noop()).withInstanceName(() -> "instance-a");
+        TestTable table2 = new TestTable(Recipe.noop()).withInstanceName(() -> "instance-b");
+
+        store.insertRow(table1, ctx, new TestTable.Row("from-a"));
+        store.insertRow(table2, ctx, new TestTable.Row("from-b"));
+
+        // getRows should return rows from ALL matching buckets, not just the first
+        List<?> rows = store.getRows(TestTable.class.getName(), null).collect(Collectors.toList());
+        assertThat(rows).hasSize(2);
+    }
+
     // =========================================================================
     // DataTable identity
     // =========================================================================