DO NOT MERGE: feat(sharding): add label-driven predicate routing

[mlwelles]

DO NOT MERGE

POC/Experiment for review / discussion purposes only.

This is a proof-of-concept implementation for label-based predicate sharding and program-based data compartmentalization. It is intended for review and testing purposes only.

---

Summary

Two complementary security features that together enable multi-level data isolation in Dgraph:

1. Label-based predicate routing (@label directive) — Routes predicates to specific alpha groups, creating physically isolated sub-clusters.
2. Program-based data compartmentalization (dgraph.programs facets) — Applies need-to-know access control at the posting level, filtering data based on user program membership.

---

Feature 1: Label-Based Predicate Routing

Concept

Add @label(name) directive to predicates, allowing them to be pinned to specific alpha groups that share the same label. This creates implicit sub-clusters for data isolation.

Changes

Protocol Buffers (protos/pb.proto, protos/pb/pb.pb.go)

• Add label field to Member (field 15), Tablet (field 12), and SchemaUpdate (field 16)
• New protos/pb/labeled.go with nil-safe IsLabeled() helpers for Tablet, Member, and SchemaUpdate

Schema Parsing (schema/parse.go, schema/schema.go)

• Parse @label(X) directive in schema definitions
• GetLabel() accessor for retrieving a predicate's label from schema state

Alpha CLI (dgraph/cmd/alpha/run.go, x/config.go)

• New --label flag on alpha, stored in WorkerOptions.Label
• Alpha registers its label with Zero via Member.Label

Zero Routing (dgraph/cmd/zero/zero.go, raft.go, tablet.go)

• Inform() and ShouldServe(): Route labeled predicates to matching labeled groups
• Prevent unlabeled predicates from being assigned to labeled groups
• Allow tablet re-routing when labels change (e.g., after DropAll + re-schema)
• chooseTablet(): Skip labeled predicates during rebalancing (pinned)
• movePredicate(): Block moves when destination group label doesn't match
• Helper functions: groupLabel(), labelGroup(), isLabeledGroup(), firstUnlabeledGroup()

Worker Routing (worker/groups.go, mutation.go, proposal.go, predicate_move.go, online_restore.go)

• Tablet(key, label), BelongsTo(key, label), ForceTablet(key, label) now accept label parameter
• ServesTablet() auto-resolves label from stored schema
• GetTabletLabel() returns cached tablet label for authorization
• populateMutationMap() passes label from schema for correct group routing
• runSchemaMutation() records schema metadata for labeled predicates served by other groups
• proposeAndWait() adds checkSchemaTablet() that allows labeled predicates to be served by different groups
• Restore operations (online_restore.go) pass labels when force-creating tablets

Label Routing Flow

1. Alpha with --label=X registers with Zero, Member.Label is set
2. Schema with @label(X) is parsed, SchemaUpdate.Label is set
3. Zero routes tablet to group where Member.Label matches
4. Rebalancer skips tablets with non-empty Label (pinned)
5. /moveTablet validates label compatibility before moving

Example Usage

# Start alphas with labels:
#   alpha1: (no label) - serves unlabeled predicates
#   alpha2: --label=secret
#   alpha3: --label=top_secret

# Schema:
name: string @index(term) .                    # -> group 1 (unlabeled)
codename: string @index(term) @label(secret) . # -> group 2 (secret)
alias: string @index(term) @label(top_secret) .# -> group 3 (top_secret)

---

Feature 2: Program-Based Data Compartmentalization

Concept

Implements need-to-know compartmentalization using dgraph.programs facets. When a user has programs in their auth context, those programs are automatically attached to every mutation (as a facet). During queries, postings are filtered so users only see data matching their program membership.

Changes

Auth Context (x/auth/)

New package providing security context extraction and authorization logic:

File	Purpose
context.go	AuthContext struct (UserID, Namespace, Groups, Level, Programs, IsNil), context helpers
extractor.go	ExtractOrNil() — extracts auth from context or returns nil-token sentinel
levels.go	CanAccess(userLevel, requiredLabel) — hierarchical level comparison (top_secret > secret > confidential > unclassified)
programs.go	HasAnyProgram() — OR-logic program membership check
transport.go	UnaryInterceptor / StreamInterceptor — gRPC interceptors that extract x-dgraph-level and x-dgraph-programs headers into AuthContext
mock.go	MockAuthContext() for testing
*_test.go	Unit tests for levels and programs

Mutation Path (edgraph/server.go)

• injectProgramFacets() — Auto-attaches dgraph.programs facet to every SET edge based on auth context
• checkMutationProgramAuth() — Pre-Raft authorization: checks if user can modify existing program-protected data
• programFacetCache — Write-through cache for recently committed program facets (addresses MemoryLayer read-through limitation)
• updateProgramFacetCache() — Updates cache after successful CommitOverNetwork

Query Path (edgraph/access.go)

• authorizeLabeledPreds() — Filters predicates by security level (hierarchical), runs independently of ACL
• authorizeQuery() — Extended to merge label-blocked and ACL-blocked predicates; both checks must pass

Query Execution (worker/task.go)

• checkProgramAuthorization() — Core posting-level filter: checks dgraph.programs facet against user's programs (OR logic)
• Applied at every query path: handleValuePostings, handleUidPostings, handleHasFunction, handleHasWithOrderFunction, handleRegexFunction, handleCompareFunction, handleMatchFunction
• facetsFilterUidPostingList() and facetsFilterValuePostingList() — Extended with program checks
• funcArgs struct gains ctx field for threading auth context through query execution

Posting Store (posting/lists.go)

• GetNoStoreSafe() — Nil-safe variant of GetNoStore for use before posting package initialization

Keys & Config (x/keys.go)

• ProgramFacetKey = "dgraph.programs" — Reserved facet key constant
• IsReservedFacetKey() for validation

Program Authorization Flow

1. gRPC interceptor extracts x-dgraph-programs header → AuthContext.Programs
2. WRITE: injectProgramFacets() adds dgraph.programs facet to edges
3. WRITE: checkMutationProgramAuth() verifies user can modify existing protected data
4. READ: checkProgramAuthorization() filters each posting by program facet
5. Result: Users only see/modify data within their program compartments

---

Testing

Test Cluster Configuration

Docker cluster in systest/label/docker-compose.yml:

Container	Group	Label	Purpose
zero1	-	-	Cluster coordinator
alpha1	1	(none)	Serves unlabeled predicates
alpha2	2	secret	Serves @label(secret) predicates
alpha3	3	top_secret	Serves @label(top_secret) predicates

Label Routing Tests (systest/label/label_test.go)

Test	What It Verifies
TestLabeledAlphaRegistration	Alphas with --label=X register label with Zero
TestLabeledPredicateRouting	@label(X) predicates route to matching group
TestLabeledPredicateDataIsolation	Mutations and queries work across labeled groups
TestLabeledPredicateCannotBeMoved	/moveTablet validates label compatibility
TestUnlabeledPredicateNotOnLabeledGroup	Unlabeled predicates never go to labeled groups
TestMissingLabelGroupError	Clear error when no alpha has requested label

Program Authorization Tests (systest/label/label_auth_test.go)

Test	What It Verifies
TestLabelQueryAuthorization	Level-based query filtering (top_secret > secret > unclassified)
TestProgramInsertFacetInjection	Programs auto-injected as facets during mutations
TestProgramQueryFiltering	Query results filtered by program membership
TestProgramMutationAuthorization	Update/delete authorization for program-protected data
TestProgramQueryFilteringOnUnlabeledPredicates	Programs work on predicates without @label

Run Tests

# Build and run all label tests
cd t && go build . && ./t --pkg=systest/label

# Verbose output
./t --pkg=systest/label -v

# Keep containers for debugging
./t --pkg=systest/label --keep

# Single test
./t --pkg=systest/label --test=TestProgramQueryFiltering

---

Files Changed (from main)

Area	Files	Lines
Proto & helpers	protos/pb.proto, protos/pb/pb.pb.go, protos/pb/labeled.go	~830
Schema	schema/parse.go, schema/schema.go	34
Alpha CLI	dgraph/cmd/alpha/run.go, x/config.go	7
Zero routing	dgraph/cmd/zero/zero.go, raft.go, tablet.go	~170
Auth package	x/auth/*.go (7 files)	~477
Query auth	edgraph/access.go, edgraph/server.go	~301
Worker routing	worker/groups.go, mutation.go, proposal.go, predicate_move.go, online_restore.go	~110
Query execution	worker/task.go	~250
Posting	posting/lists.go	9
Keys	x/keys.go	10
Tests	systest/label/ (3 files)	~1510
Test utils	testutil/zero.go	2
Config	.gitignore	4
Total	33 files	~3700 lines