chore: implement self in schemav2 (#2887)

Author: tstirrat15

pkg/schema/arrows.go

@@ -180,6 +180,9 @@ func (as *ArrowSet) collectArrowInformationForSetOperation(ctx context.Context,
 		case *core.SetOperation_Child_XNil:
 			// Nothing to do
 
+		case *core.SetOperation_Child_XSelf:
+			// Nothing to do
+
 		default:
 			return fmt.Errorf("unknown set operation child `%T` in addArrowRelationsInSetOperation", child)
 		}

pkg/schema/full_reachability.go

@@ -220,6 +220,9 @@ func setOperationReferencesRelation(ctx context.Context, so *core.SetOperation,
 		case *core.SetOperation_Child_XNil:
 			// Nothing to do
 
+		case *core.SetOperation_Child_XSelf:
+			// Nothing to do
+
 		default:
 			return false, fmt.Errorf("unknown set operation child `%T` in setOperationReferencesRelation", child)
 		}

pkg/schema/v2/flatten.go

@@ -185,6 +185,10 @@ func flattenOperation(op Operation, def *Definition, baseName string, options Fl
 		// Nil reference is a leaf node, no flattening needed
 		return o, nil, nil
 
+	case *SelfReference:
+		// Self reference is a leaf node, no flattening needed
+		return o, nil, nil
+
 	case *UnionOperation:
 		// Flatten children first
 		flattenedChildren := make([]Operation, len(o.children))

pkg/schema/v2/flatten_test.go

@@ -32,6 +32,17 @@ definition user {}`,
 
 definition user {}`,
 		},
+		{
+			name: "self permission",
+			schemaString: `use self
+definition user {
+	permission view = self
+			}`,
+			expectedString: `use self
+definition user {
+	permission view = self
+			}`,
+		},
 		{
 			name: "nested union in permission",
 			schemaString: `definition document {

pkg/schema/v2/resolved.go

@@ -151,6 +151,10 @@ func resolveOperation(op Operation, def *Definition) (Operation, error) {
 		// NilReference is not replaced during resolution
 		return o, nil
 
+	case *SelfReference:
+		// SelfReference is not replaced during resolution
+		return o, nil
+
 	default:
 		return nil, fmt.Errorf("unknown operation type: %T", op)
 	}

pkg/schema/v2/resolved_test.go

@@ -29,6 +29,35 @@ func TestResolveSchema_Empty(t *testing.T) {
 	require.Empty(t, resolved.Schema().definitions)
 }
 
+func TestResolveSchema_SelfReference(t *testing.T) {
+	schema := &Schema{
+		definitions: map[string]*Definition{
+			"user": {
+				name: "user",
+				permissions: map[string]*Permission{
+					"view": {
+						name:      "view",
+						operation: &SelfReference{},
+					},
+				},
+			},
+		},
+	}
+
+	resolved, err := ResolveSchema(schema)
+	require.NoError(t, err)
+	require.NotNil(t, resolved)
+
+	// Check that the original schema is unchanged
+	originalPerm := schema.definitions["user"].permissions["view"]
+	require.IsType(t, &SelfReference{}, originalPerm.operation)
+
+	// Check that the resolved schema has ResolvedRelationReference
+	resolvedDef := resolved.Schema().definitions["user"]
+	resolvedPerm := resolvedDef.permissions["view"]
+	require.IsType(t, &SelfReference{}, resolvedPerm.operation)
+}
+
 func TestResolveSchema_SimpleRelationReference(t *testing.T) {
 	rel := &Relation{
 		name:          "viewer",

pkg/schema/v2/tocorev1.go

@@ -344,6 +344,19 @@ func operationToUsersetRewrite(op Operation) (*core.UsersetRewrite, error) {
 			},
 		}, nil
 
+	case *SelfReference:
+		return &core.UsersetRewrite{
+			RewriteOperation: &core.UsersetRewrite_Union{
+				Union: &core.SetOperation{
+					Child: []*core.SetOperation_Child{
+						{
+							ChildType: &core.SetOperation_Child_XSelf{},
+						},
+					},
+				},
+			},
+		}, nil
+
 	default:
 		return nil, fmt.Errorf("unknown operation type: %T", op)
 	}
@@ -462,6 +475,11 @@ func operationToChild(op Operation) (*core.SetOperation_Child, error) {
 			ChildType: &core.SetOperation_Child_XNil{},
 		}, nil
 
+	case *SelfReference:
+		return &core.SetOperation_Child{
+			ChildType: &core.SetOperation_Child_XSelf{},
+		}, nil
+
 	default:
 		return nil, fmt.Errorf("unknown operation type: %T", op)
 	}

pkg/schema/v2/tocorev1_test.go

@@ -119,6 +119,21 @@ func TestAsCompiledSchema(t *testing.T) {
 			expectedDefinitionNames: []string{"user", "organization", "resource"},
 			expectedCaveatNames:     []string{},
 		},
+		{
+			name: "schema with self",
+			schemaText: `
+			use self
+
+			definition user {
+				relation viewer: user
+				permission view = viewer + self
+			}
+		`,
+			expectedNumDefinitions:  1,
+			expectedNumCaveats:      0,
+			expectedDefinitionNames: []string{"user"},
+			expectedCaveatNames:     []string{},
+		},
 	}
 
 	for _, tc := range tcs {
@@ -172,11 +187,16 @@ func TestAsCompiledSchemaRoundTrip(t *testing.T) {
 	t.Parallel()
 
 	schemaText := `
+		use self
+
 		caveat ip_check(ip string) {
 			ip == "192.168.1.1"
 		}
 
-		definition user {}
+		definition user {
+			relation viewer: user
+			permission view = viewer + self
+		}
 
 		definition organization {
 			relation member: user