Skip to content

Commit cc3ede7

Browse files
seawindeseawinde
committed
[test](fe) Restore auth state in CTE privilege tests
### What problem does this PR solve? Issue Number: None Related PR: #62339 Problem Summary: Restore the original AccessControllerManager and test user after CTE privilege tests replace process-global FE auth state, preventing order-dependent leakage across TestWithFeService test methods. ### Release note None ### Check List (For Author) - Test: No test (not run locally) - Behavior changed: No - Does this need documentation: No
1 parent 2d91b04 commit cc3ede7

1 file changed

Lines changed: 95 additions & 85 deletions

File tree

‎fe/fe-core/src/test/java/org/apache/doris/nereids/privileges/TestCheckPrivileges.java‎

Lines changed: 95 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -132,95 +132,102 @@ public void testPrivilegesAndPolicies() throws Exception {
132132
.addDataMasking(user, "id", "concat(id, '_****_', id)")
133133
);
134134

135-
AccessControllerManager accessManager = Env.getCurrentEnv().getAccessManager();
135+
Env currentEnv = Env.getCurrentEnv();
136+
AccessControllerManager accessManager = currentEnv.getAccessManager();
137+
AccessControllerManager originalAccessManager = accessManager;
136138
CatalogAccessController catalogAccessController = accessManager.getAccessControllerOrDefault(catalog);
137139
AccessControllerManager spyAccessManager = Mockito.spy(accessManager);
138140
Mockito.doReturn(catalogAccessController).when(spyAccessManager)
139141
.getAccessControllerOrDefault("internal");
140-
Deencapsulation.setField(Env.getCurrentEnv(), "accessManager", spyAccessManager);
141-
142-
withPrivileges(privileges, () -> {
143-
// test base table
144-
{
145-
// has table privilege
146-
query("select * from custom_catalog.test_db.test_tbl1");
147-
148-
// has id column privilege
149-
query("select id from custom_catalog.test_db.test_tbl2");
150-
151-
// no name column privilege, throw exception:
152-
//
153-
// Permission denied: user ['test_nereids_privilege_user'@'%'] does not have privilege for
154-
// [priv predicate: OR, Admin_priv Select_priv ] command on
155-
// [custom_catalog].[test_db].[test_tbl2].[name]
156-
Assertions.assertThrows(AnalysisException.class, () ->
157-
query("select * from custom_catalog.test_db.test_tbl2")
158-
);
159-
160-
// no table privilege
161-
Assertions.assertThrows(AnalysisException.class, () ->
162-
query("select * from custom_catalog.test_db.test_tbl3")
163-
);
164-
}
142+
Deencapsulation.setField(currentEnv, "accessManager", spyAccessManager);
165143

166-
// test row policy with data masking
167-
{
168-
Function<NamedExpression, Boolean> checkId = (NamedExpression ne) -> {
169-
if (!(ne instanceof Alias) || !ne.getName().equals("id")) {
170-
return false;
171-
}
172-
return ne.child(0).toSql().equals("'1_****_1'");
173-
};
174-
PlanChecker.from(connectContext)
175-
.parse("select id,"
176-
+ " test_tbl4.id,"
177-
+ " test_db.test_tbl4.id, "
178-
+ " custom_catalog.test_db.test_tbl4.id, "
179-
+ " * "
180-
+ "from custom_catalog.test_db.test_tbl4")
181-
.analyze()
182-
.rewrite()
183-
.matches(logicalProject(
184-
logicalFilter(
185-
logicalTestScan()
186-
).when(f -> {
187-
EqualTo predicate = (EqualTo) f.getPredicate();
188-
return predicate.left() instanceof Slot
189-
&& predicate.right().equals(new IntegerLiteral((byte) 1));
190-
})
191-
).when(p -> {
192-
List<NamedExpression> projects = p.getProjects();
193-
if (!checkId.apply(projects.get(0)) || !checkId.apply(projects.get(1))
194-
|| !checkId.apply(projects.get(2)) || !checkId.apply(projects.get(3))
195-
|| !checkId.apply(projects.get(4))) {
196-
return false;
197-
}
198-
return projects.get(5) instanceof Slot && projects.get(5).getName().equals("name");
199-
}));
200-
201-
PlanChecker.from(connectContext)
202-
.parse("select id, t.id, *"
203-
+ "from custom_catalog.test_db.test_tbl4 t")
204-
.analyze()
205-
.rewrite()
206-
.matches(logicalProject(
207-
logicalFilter(
144+
try {
145+
withPrivileges(privileges, () -> {
146+
// test base table
147+
{
148+
// has table privilege
149+
query("select * from custom_catalog.test_db.test_tbl1");
150+
151+
// has id column privilege
152+
query("select id from custom_catalog.test_db.test_tbl2");
153+
154+
// no name column privilege, throw exception:
155+
//
156+
// Permission denied: user ['test_nereids_privilege_user'@'%'] does not have privilege for
157+
// [priv predicate: OR, Admin_priv Select_priv ] command on
158+
// [custom_catalog].[test_db].[test_tbl2].[name]
159+
Assertions.assertThrows(AnalysisException.class, () ->
160+
query("select * from custom_catalog.test_db.test_tbl2")
161+
);
162+
163+
// no table privilege
164+
Assertions.assertThrows(AnalysisException.class, () ->
165+
query("select * from custom_catalog.test_db.test_tbl3")
166+
);
167+
}
168+
169+
// test row policy with data masking
170+
{
171+
Function<NamedExpression, Boolean> checkId = (NamedExpression ne) -> {
172+
if (!(ne instanceof Alias) || !ne.getName().equals("id")) {
173+
return false;
174+
}
175+
return ne.child(0).toSql().equals("'1_****_1'");
176+
};
177+
PlanChecker.from(connectContext)
178+
.parse("select id,"
179+
+ " test_tbl4.id,"
180+
+ " test_db.test_tbl4.id, "
181+
+ " custom_catalog.test_db.test_tbl4.id, "
182+
+ " * "
183+
+ "from custom_catalog.test_db.test_tbl4")
184+
.analyze()
185+
.rewrite()
186+
.matches(logicalProject(
187+
logicalFilter(
208188
logicalTestScan()
209-
).when(f -> {
210-
EqualTo predicate = (EqualTo) f.getPredicate();
211-
return predicate.left() instanceof Slot
212-
&& predicate.right().equals(new IntegerLiteral((byte) 1));
213-
})
214-
).when(p -> {
215-
List<NamedExpression> projects = p.getProjects();
216-
if (!checkId.apply(projects.get(0)) || !checkId.apply(projects.get(1))
217-
|| !checkId.apply(projects.get(2))) {
218-
return false;
219-
}
220-
return projects.get(3) instanceof Slot && projects.get(3).getName().equals("name");
221-
}));
222-
}
223-
});
189+
).when(f -> {
190+
EqualTo predicate = (EqualTo) f.getPredicate();
191+
return predicate.left() instanceof Slot
192+
&& predicate.right().equals(new IntegerLiteral((byte) 1));
193+
})
194+
).when(p -> {
195+
List<NamedExpression> projects = p.getProjects();
196+
if (!checkId.apply(projects.get(0)) || !checkId.apply(projects.get(1))
197+
|| !checkId.apply(projects.get(2)) || !checkId.apply(projects.get(3))
198+
|| !checkId.apply(projects.get(4))) {
199+
return false;
200+
}
201+
return projects.get(5) instanceof Slot && projects.get(5).getName().equals("name");
202+
}));
203+
204+
PlanChecker.from(connectContext)
205+
.parse("select id, t.id, *"
206+
+ "from custom_catalog.test_db.test_tbl4 t")
207+
.analyze()
208+
.rewrite()
209+
.matches(logicalProject(
210+
logicalFilter(
211+
logicalTestScan()
212+
).when(f -> {
213+
EqualTo predicate = (EqualTo) f.getPredicate();
214+
return predicate.left() instanceof Slot
215+
&& predicate.right().equals(new IntegerLiteral((byte) 1));
216+
})
217+
).when(p -> {
218+
List<NamedExpression> projects = p.getProjects();
219+
if (!checkId.apply(projects.get(0)) || !checkId.apply(projects.get(1))
220+
|| !checkId.apply(projects.get(2))) {
221+
return false;
222+
}
223+
return projects.get(3) instanceof Slot && projects.get(3).getName().equals("name");
224+
}));
225+
}
226+
});
227+
} finally {
228+
Deencapsulation.setField(currentEnv, "accessManager", originalAccessManager);
229+
useUser("root");
230+
}
224231
}
225232

226233
@Test
@@ -254,17 +261,19 @@ public void testCtePrivilegeCheck() throws Exception {
254261
addUser(user, true);
255262
useUser(user);
256263

264+
Env currentEnv = Env.getCurrentEnv();
265+
AccessControllerManager originalAccessManager = currentEnv.getAccessManager();
257266
try {
258267
List<MakeTablePrivileges> privileges = ImmutableList.of(
259268
MakePrivileges.table("internal", cteDb, "allowed_tbl").allowSelectTable(user)
260269
);
261270

262-
AccessControllerManager accessManager = Env.getCurrentEnv().getAccessManager();
271+
AccessControllerManager accessManager = currentEnv.getAccessManager();
263272
CatalogAccessController catalogAccessController = accessManager.getAccessControllerOrDefault(catalog);
264273
AccessControllerManager spyAccessManager = Mockito.spy(accessManager);
265274
Mockito.doReturn(catalogAccessController).when(spyAccessManager)
266275
.getAccessControllerOrDefault("internal");
267-
Deencapsulation.setField(Env.getCurrentEnv(), "accessManager", spyAccessManager);
276+
Deencapsulation.setField(currentEnv, "accessManager", spyAccessManager);
268277

269278
withPrivileges(privileges, () -> {
270279
// CTE with authorized table should succeed
@@ -290,6 +299,7 @@ public void testCtePrivilegeCheck() throws Exception {
290299
);
291300
});
292301
} finally {
302+
Deencapsulation.setField(currentEnv, "accessManager", originalAccessManager);
293303
useUser("root");
294304
}
295305
}

0 commit comments

Comments
 (0)