Skip to content

Commit 9f73d87

Browse files
tstirrat15tstirrat15
authored
chore: fix query plan alias iterator LR behavior (#2826)
1 parent 93dd1c7 commit 9f73d87

File tree

2 files changed

+96
-8
lines changed

2 files changed

+96
-8
lines changed

‎pkg/query/alias.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,40 @@ func (a *Alias) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (Path
122122
return nil, err
123123
}
124124

125+
// If the relation on the alias iterator is the same
126+
// as the subject relation on the subject, we have
127+
// a self edge and want to yield it accordingly.
128+
if a.relation == subject.Relation {
129+
selfPath := Path{
130+
Resource: GetObject(subject),
131+
Relation: a.relation,
132+
Subject: subject,
133+
}
134+
135+
combined := func(yield func(Path, error) bool) {
136+
// Yield the self path first
137+
if !yield(selfPath, nil) {
138+
return
139+
}
140+
141+
for subPath, err := range subSeq {
142+
if err != nil {
143+
yield(Path{}, err)
144+
return
145+
}
146+
147+
// Then yield remaining paths with the rewrite
148+
subPath.Relation = a.relation
149+
if !yield(subPath, nil) {
150+
return
151+
}
152+
}
153+
}
154+
155+
return DeduplicatePathSeq(combined), nil
156+
}
157+
158+
// else we don't have a subpath and do the normal logic
125159
return func(yield func(Path, error) bool) {
126160
for path, err := range subSeq {
127161
if err != nil {

‎pkg/query/alias_test.go‎

Lines changed: 62 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,12 @@ import (
99
func TestAliasIterator(t *testing.T) {
1010
t.Parallel()
1111

12-
require := require.New(t)
13-
1412
// Create test context
1513
ctx := NewLocalContext(t.Context())
1614

1715
t.Run("Check_BasicRelationRewriting", func(t *testing.T) {
1816
t.Parallel()
17+
require := require.New(t)
1918

2019
// Create a sub-iterator with document relations
2120
subIt := NewDocumentAccessFixedIterator()
@@ -46,6 +45,7 @@ func TestAliasIterator(t *testing.T) {
4645

4746
t.Run("Check_SelfEdgeDetection", func(t *testing.T) {
4847
t.Parallel()
48+
require := require.New(t)
4949

5050
// Create an empty sub-iterator since we only want to test self-edge detection
5151
subIt := NewEmptyFixedIterator()
@@ -77,6 +77,7 @@ func TestAliasIterator(t *testing.T) {
7777

7878
t.Run("Check_NoSelfEdge", func(t *testing.T) {
7979
t.Parallel()
80+
require := require.New(t)
8081

8182
// Create a sub-iterator
8283
subIt := NewSingleUserFixedIterator("bob")
@@ -100,6 +101,7 @@ func TestAliasIterator(t *testing.T) {
100101

101102
t.Run("Check_MultipleResources", func(t *testing.T) {
102103
t.Parallel()
104+
require := require.New(t)
103105

104106
subIt := NewDocumentAccessFixedIterator()
105107
aliasIt := NewAlias("access", subIt)
@@ -126,6 +128,7 @@ func TestAliasIterator(t *testing.T) {
126128

127129
t.Run("IterSubjects_RelationRewriting", func(t *testing.T) {
128130
t.Parallel()
131+
require := require.New(t)
129132

130133
subIt := NewDocumentAccessFixedIterator()
131134
aliasIt := NewAlias("permission", subIt)
@@ -146,6 +149,7 @@ func TestAliasIterator(t *testing.T) {
146149

147150
t.Run("IterResources_RelationRewriting", func(t *testing.T) {
148151
t.Parallel()
152+
require := require.New(t)
149153

150154
subIt := NewDocumentAccessFixedIterator()
151155
aliasIt := NewAlias("can_view", subIt)
@@ -162,8 +166,41 @@ func TestAliasIterator(t *testing.T) {
162166
}
163167
})
164168

169+
t.Run("IterResources_SelfEdgeDetection", func(t *testing.T) {
170+
t.Parallel()
171+
require := require.New(t)
172+
173+
// Create an empty sub-iterator since we only want to test self-edge detection
174+
subIt := NewEmptyFixedIterator()
175+
176+
// Create an alias iterator that rewrites to "admin"
177+
aliasIt := NewAlias("admin", subIt)
178+
179+
// Check for a self-edge: user:alice#admin@user:alice#admin
180+
subject := NewObjectAndRelation("alice", "user", "admin")
181+
pathSeq, err := ctx.IterResources(aliasIt, subject)
182+
require.NoError(err)
183+
184+
rels, err := CollectAll(pathSeq)
185+
require.NoError(err)
186+
187+
// Should find exactly one relation (the self-edge)
188+
require.Len(rels, 1, "should find exactly one self-edge relation")
189+
190+
// Verify it's the correct self-edge relation
191+
rel := rels[0]
192+
expectedResource := NewObjectAndRelation("alice", "user", "admin")
193+
require.Equal(expectedResource.ObjectID, rel.Resource.ObjectID)
194+
require.Equal(expectedResource.ObjectType, rel.Resource.ObjectType)
195+
require.Equal(expectedResource.Relation, rel.Relation)
196+
require.Equal(subject.ObjectID, rel.Subject.ObjectID)
197+
require.Equal(subject.ObjectType, rel.Subject.ObjectType)
198+
require.Equal(subject.Relation, rel.Subject.Relation)
199+
})
200+
165201
t.Run("Check_EmptySubIterator", func(t *testing.T) {
166202
t.Parallel()
203+
require := require.New(t)
167204

168205
subIt := NewEmptyFixedIterator()
169206
aliasIt := NewAlias("empty_alias", subIt)
@@ -178,6 +215,7 @@ func TestAliasIterator(t *testing.T) {
178215

179216
t.Run("Check_SelfEdgeWithEmptySubIterator", func(t *testing.T) {
180217
t.Parallel()
218+
require := require.New(t)
181219

182220
// Create empty sub-iterator
183221
subIt := NewEmptyFixedIterator()
@@ -198,6 +236,7 @@ func TestAliasIterator(t *testing.T) {
198236

199237
t.Run("Check_SelfEdgeExactMatch", func(t *testing.T) {
200238
t.Parallel()
239+
require := require.New(t)
201240

202241
subIt := NewEmptyFixedIterator()
203242
aliasIt := NewAlias("owner", subIt)
@@ -249,10 +288,9 @@ func TestAliasIteratorClone(t *testing.T) {
249288
func TestAliasIteratorExplain(t *testing.T) {
250289
t.Parallel()
251290

252-
require := require.New(t)
253-
254291
t.Run("ExplainWithSubIterator", func(t *testing.T) {
255292
t.Parallel()
293+
require := require.New(t)
256294

257295
subIt := NewDocumentAccessFixedIterator()
258296
aliasIt := NewAlias("test_relation", subIt)
@@ -268,6 +306,7 @@ func TestAliasIteratorExplain(t *testing.T) {
268306

269307
t.Run("ExplainWithEmptySubIterator", func(t *testing.T) {
270308
t.Parallel()
309+
require := require.New(t)
271310

272311
subIt := NewEmptyFixedIterator()
273312
aliasIt := NewAlias("empty_test", subIt)
@@ -281,13 +320,12 @@ func TestAliasIteratorExplain(t *testing.T) {
281320
func TestAliasIteratorErrorHandling(t *testing.T) {
282321
t.Parallel()
283322

284-
require := require.New(t)
285-
286323
// Create test context
287324
ctx := NewLocalContext(t.Context())
288325

289326
t.Run("Check_SubIteratorError", func(t *testing.T) {
290327
t.Parallel()
328+
require := require.New(t)
291329

292330
// Create a faulty sub-iterator
293331
faultyIt := NewFaultyIterator(true, false)
@@ -299,6 +337,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
299337

300338
t.Run("Check_SubIteratorCollectionError", func(t *testing.T) {
301339
t.Parallel()
340+
require := require.New(t)
302341

303342
// Create a faulty sub-iterator that fails during collection
304343
faultyIt := NewFaultyIterator(false, true)
@@ -314,6 +353,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
314353

315354
t.Run("IterSubjects_SubIteratorError", func(t *testing.T) {
316355
t.Parallel()
356+
require := require.New(t)
317357

318358
// Create a faulty sub-iterator that fails on IterSubjectsImpl
319359
faultyIt := NewFaultyIterator(true, false)
@@ -326,6 +366,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
326366

327367
t.Run("IterSubjects_SubIteratorCollectionError", func(t *testing.T) {
328368
t.Parallel()
369+
require := require.New(t)
329370

330371
// Create a faulty sub-iterator that fails during collection
331372
faultyIt := NewFaultyIterator(false, true)
@@ -342,6 +383,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
342383

343384
t.Run("IterResources_SubIteratorError", func(t *testing.T) {
344385
t.Parallel()
386+
require := require.New(t)
345387

346388
// Create a faulty sub-iterator that fails on IterResourcesImpl
347389
faultyIt := NewFaultyIterator(true, false)
@@ -354,6 +396,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
354396

355397
t.Run("IterResources_SubIteratorCollectionError", func(t *testing.T) {
356398
t.Parallel()
399+
require := require.New(t)
357400

358401
// Create a faulty sub-iterator that fails during collection
359402
faultyIt := NewFaultyIterator(false, true)
@@ -370,6 +413,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
370413

371414
t.Run("Check_SelfEdgeWithSubIteratorError", func(t *testing.T) {
372415
t.Parallel()
416+
require := require.New(t)
373417

374418
// Create a faulty sub-iterator that errors on CheckImpl
375419
faultyIt := NewFaultyIterator(true, false)
@@ -385,6 +429,7 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
385429

386430
t.Run("Check_SelfEdgeWithSubIteratorCollectionError", func(t *testing.T) {
387431
t.Parallel()
432+
require := require.New(t)
388433

389434
// Create a faulty sub-iterator that fails during collection
390435
faultyIt := NewFaultyIterator(false, true)
@@ -405,13 +450,12 @@ func TestAliasIteratorErrorHandling(t *testing.T) {
405450
func TestAliasIteratorAdvancedScenarios(t *testing.T) {
406451
t.Parallel()
407452

408-
require := require.New(t)
409-
410453
// Create test context
411454
ctx := NewLocalContext(t.Context())
412455

413456
t.Run("Check_MultipleResourcesSelfEdgeWithSubResults", func(t *testing.T) {
414457
t.Parallel()
458+
require := require.New(t)
415459

416460
// Create sub-iterator with real data
417461
subIt := NewDocumentAccessFixedIterator()
@@ -451,6 +495,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
451495

452496
t.Run("Check_MultipleResourcesMultipleSelfEdges", func(t *testing.T) {
453497
t.Parallel()
498+
require := require.New(t)
454499

455500
// Create empty sub-iterator to isolate self-edge logic
456501
subIt := NewEmptyFixedIterator()
@@ -478,6 +523,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
478523

479524
t.Run("IterSubjects_EmptySubIterator", func(t *testing.T) {
480525
t.Parallel()
526+
require := require.New(t)
481527

482528
subIt := NewEmptyFixedIterator()
483529
aliasIt := NewAlias("empty_relation", subIt)
@@ -492,6 +538,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
492538

493539
t.Run("IterResources_EmptySubIterator", func(t *testing.T) {
494540
t.Parallel()
541+
require := require.New(t)
495542

496543
subIt := NewEmptyFixedIterator()
497544
aliasIt := NewAlias("empty_relation", subIt)
@@ -506,6 +553,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
506553

507554
t.Run("Check_LargeResultSet", func(t *testing.T) {
508555
t.Parallel()
556+
require := require.New(t)
509557

510558
// Use iterator with many results
511559
subIt := NewLargeFixedIterator()
@@ -529,6 +577,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
529577

530578
t.Run("Clone_Independence", func(t *testing.T) {
531579
t.Parallel()
580+
require := require.New(t)
532581

533582
subIt := NewDocumentAccessFixedIterator()
534583
original := NewAlias("original", subIt)
@@ -564,6 +613,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
564613

565614
t.Run("Check_EarlyReturnBasicRewriting", func(t *testing.T) {
566615
t.Parallel()
616+
require := require.New(t)
567617

568618
subIt := NewDocumentAccessFixedIterator()
569619
aliasIt := NewAlias("early_test", subIt)
@@ -586,6 +636,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
586636

587637
t.Run("Check_EarlyReturnSelfEdgeWithSubResults", func(t *testing.T) {
588638
t.Parallel()
639+
require := require.New(t)
589640

590641
subIt := NewDocumentAccessFixedIterator()
591642
aliasIt := NewAlias("admin", subIt)
@@ -610,6 +661,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
610661

611662
t.Run("IterSubjects_EarlyReturn", func(t *testing.T) {
612663
t.Parallel()
664+
require := require.New(t)
613665

614666
subIt := NewDocumentAccessFixedIterator()
615667
aliasIt := NewAlias("early_subjects", subIt)
@@ -632,6 +684,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
632684

633685
t.Run("IterResources_EarlyReturn", func(t *testing.T) {
634686
t.Parallel()
687+
require := require.New(t)
635688

636689
subIt := NewDocumentAccessFixedIterator()
637690
aliasIt := NewAlias("early_resources", subIt)
@@ -654,6 +707,7 @@ func TestAliasIteratorAdvancedScenarios(t *testing.T) {
654707

655708
t.Run("Check_SelfEdgeWithSubIteratorPathRewriting", func(t *testing.T) {
656709
t.Parallel()
710+
require := require.New(t)
657711

658712
// Use sub-iterator that will return paths that need rewriting
659713
subIt := NewDocumentAccessFixedIterator()

0 commit comments

Comments
 (0)