diff --git a/rewrite-core/src/main/java/org/openrewrite/DataTableStore.java b/rewrite-core/src/main/java/org/openrewrite/DataTableStore.java index b337fe202b..ca7071f1c9 100644 --- a/rewrite-core/src/main/java/org/openrewrite/DataTableStore.java +++ b/rewrite-core/src/main/java/org/openrewrite/DataTableStore.java @@ -63,6 +63,17 @@ static DataTableStore noop() { */ Stream getRows(String dataTableName, @Nullable String group); + /** + * Stream all rows for a given data table name, merging across all + * groups and instance names. + * + * @param dataTableName the fully qualified class name of the data table + * @return a stream of all rows across all buckets with this name + */ + default Stream getAllRows(String dataTableName) { + return getRows(dataTableName, null); + } + /** * Get the set of {@link DataTable} instances that have received rows. * diff --git a/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java b/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java index 5ce6612fa4..61bd45652a 100644 --- a/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java +++ b/rewrite-core/src/main/java/org/openrewrite/InMemoryDataTableStore.java @@ -78,6 +78,19 @@ private static Stream snapshotRows(Bucket bucket) { return snapshot.stream(); } + @Override + public Stream getAllRows(String dataTableName) { + List allRows = new ArrayList<>(); + for (Bucket bucket : buckets.values()) { + if (bucket.dataTable.getName().equals(dataTableName)) { + synchronized (bucket.rows) { + allRows.addAll(bucket.rows); + } + } + } + return allRows.stream(); + } + @Override public Collection> getDataTables() { List> result = new ArrayList<>(buckets.size());