diff --git a/ext/store/bigquery/resource_urn.go b/ext/store/bigquery/resource_urn.go index d3e8b90674..3d5790fa9f 100644 --- a/ext/store/bigquery/resource_urn.go +++ b/ext/store/bigquery/resource_urn.go @@ -73,13 +73,14 @@ type ResourceURNWithUpstreamsList []*ResourceURNWithUpstreams func (rs ResourceURNWithUpstreamsList) FlattenUnique() []*ResourceURNWithUpstreams { var output []*ResourceURNWithUpstreams - for _, u := range rs { - if u == nil { + for _, r := range rs { + if r == nil { continue } - nested := ResourceURNWithUpstreamsList(u.Upstreams).FlattenUnique() - u.Upstreams = nil - output = append(output, u) + newResource := *r + newResource.Upstreams = nil + nested := ResourceURNWithUpstreamsList(r.Upstreams).FlattenUnique() + output = append(output, &newResource) output = append(output, nested...) } @@ -88,13 +89,15 @@ func (rs ResourceURNWithUpstreamsList) FlattenUnique() []*ResourceURNWithUpstrea func (rs ResourceURNWithUpstreamsList) unique() ResourceURNWithUpstreamsList { mapUnique := map[ResourceURN]*ResourceURNWithUpstreams{} - for _, u := range rs { - mapUnique[u.ResourceURN] = u + for _, r := range rs { + mapUnique[r.ResourceURN] = r } - var output []*ResourceURNWithUpstreams + output := make([]*ResourceURNWithUpstreams, len(mapUnique)) + i := 0 for _, u := range mapUnique { - output = append(output, u) + output[i] = u + i++ } return output } diff --git a/ext/store/bigquery/resource_urn_test.go b/ext/store/bigquery/resource_urn_test.go index 369d3a1684..6fd6424e09 100644 --- a/ext/store/bigquery/resource_urn_test.go +++ b/ext/store/bigquery/resource_urn_test.go @@ -39,3 +39,52 @@ func TestURN(t *testing.T) { assert.Equal(t, "bigquery://project:dataset.table", resourceURN.URN()) }) } + +func TestFlattenUnique(t *testing.T) { + t.Run("should return flattened upstream in the form of resource", func(t *testing.T) { + resources := []*bigquery.ResourceURNWithUpstreams{ + { + ResourceURN: newResourceURN("project_test_1", "dataset_test_1", "name_test_1"), + }, + { + ResourceURN: newResourceURN("project_test_2", "dataset_test_2", "name_test_2"), + Upstreams: []*bigquery.ResourceURNWithUpstreams{ + { + ResourceURN: newResourceURN("project_test_3", "dataset_test_3", "name_test_3"), + }, + { + ResourceURN: newResourceURN("project_test_4", "dataset_test_4", "name_test_4"), + }, + { + ResourceURN: newResourceURN("project_test_1", "dataset_test_1", "name_test_1"), + }, + nil, + }, + }, + } + + expectedResources := []*bigquery.ResourceURNWithUpstreams{ + { + ResourceURN: newResourceURN("project_test_1", "dataset_test_1", "name_test_1"), + }, + { + ResourceURN: newResourceURN("project_test_2", "dataset_test_2", "name_test_2"), + }, + { + ResourceURN: newResourceURN("project_test_3", "dataset_test_3", "name_test_3"), + }, + { + ResourceURN: newResourceURN("project_test_4", "dataset_test_4", "name_test_4"), + }, + } + + actualResources := bigquery.ResourceURNWithUpstreamsList(resources).FlattenUnique() + + assert.ElementsMatch(t, expectedResources, actualResources) + }) +} + +func newResourceURN(project, dataset, name string) bigquery.ResourceURN { + resourceURN, _ := bigquery.NewResourceURN(project, dataset, name) + return resourceURN +}