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 void insertRow(DataTable dataTable, ExecutionContext ctx, Row @Override public Stream getRows(String dataTableName, @Nullable String group) { - // Scan for matching bucket + List 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 snapshotRows(Bucket bucket) { - List 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 // =========================================================================