8282import org .apache .iotdb .db .queryengine .plan .statement .metadata .ShowVariablesStatement ;
8383import org .apache .iotdb .db .queryengine .plan .statement .metadata .UnSetTTLStatement ;
8484import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .CreateModelStatement ;
85+ import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .CreateTrainingStatement ;
8586import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .DropModelStatement ;
87+ import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .LoadModelStatement ;
88+ import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .ShowAIDevicesStatement ;
8689import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .ShowAINodesStatement ;
90+ import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .ShowLoadedModelsStatement ;
8791import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .ShowModelsStatement ;
92+ import org .apache .iotdb .db .queryengine .plan .statement .metadata .model .UnloadModelStatement ;
8893import org .apache .iotdb .db .queryengine .plan .statement .metadata .pipe .AlterPipeStatement ;
8994import org .apache .iotdb .db .queryengine .plan .statement .metadata .pipe .CreatePipePluginStatement ;
9095import org .apache .iotdb .db .queryengine .plan .statement .metadata .pipe .CreatePipeStatement ;
123128import org .apache .iotdb .db .queryengine .plan .statement .sys .AuthorStatement ;
124129import org .apache .iotdb .db .queryengine .plan .statement .sys .ClearCacheStatement ;
125130import org .apache .iotdb .db .queryengine .plan .statement .sys .ExplainAnalyzeStatement ;
131+ import org .apache .iotdb .db .queryengine .plan .statement .sys .ExplainStatement ;
126132import org .apache .iotdb .db .queryengine .plan .statement .sys .FlushStatement ;
127133import org .apache .iotdb .db .queryengine .plan .statement .sys .KillQueryStatement ;
134+ import org .apache .iotdb .db .queryengine .plan .statement .sys .LoadConfigurationStatement ;
128135import org .apache .iotdb .db .queryengine .plan .statement .sys .SetConfigurationStatement ;
129136import org .apache .iotdb .db .queryengine .plan .statement .sys .SetSqlDialectStatement ;
130137import org .apache .iotdb .db .queryengine .plan .statement .sys .SetSystemStatusStatement ;
@@ -161,10 +168,7 @@ public class TreeAccessCheckVisitor extends StatementVisitor<TSStatus, TreeAcces
161168
162169 @ Override
163170 public TSStatus visitNode (StatementNode node , TreeAccessCheckContext context ) {
164- if (AuthorityChecker .SUPER_USER .equals (context .getUsername ())) {
165- return SUCCEED ;
166- }
167- return AuthorityChecker .getTSStatus (false , "Only the admin user can perform this operation" );
171+ throw new IllegalStateException ("Each operation should have permission check." );
168172 }
169173
170174 @ Override
@@ -229,7 +233,7 @@ private TSStatus checkTemplateShowRelated(
229233 return SUCCEED ;
230234 }
231235 // own SYSTEM can see all, otherwise can only see PATHS that user has READ_SCHEMA auth
232- if (!AuthorityChecker . checkSystemPermission (context .getUsername (), PrivilegeType .SYSTEM )) {
236+ if (!checkHasGlobalAuth (context .getUsername (), PrivilegeType .SYSTEM )) {
233237 statement .setCanSeeAll (false );
234238 return visitAuthorityInformation (statement , context );
235239 } else {
@@ -289,11 +293,7 @@ public TSStatus visitAlterSchemaTemplate(
289293 if (AuthorityChecker .SUPER_USER .equals (context .getUsername ())) {
290294 return SUCCEED ;
291295 }
292- return AuthorityChecker .getTSStatus (
293- AuthorityChecker .checkSystemPermission (context .getUsername (), PrivilegeType .SYSTEM )
294- || AuthorityChecker .checkSystemPermission (
295- context .getUsername (), PrivilegeType .EXTEND_TEMPLATE ),
296- PrivilegeType .SYSTEM );
296+ return checkGlobalAuth (context .getUsername (), PrivilegeType .EXTEND_TEMPLATE );
297297 }
298298
299299 // ============================= timeseries view related ===============
@@ -455,8 +455,7 @@ public TSStatus visitAuthor(AuthorStatement statement, TreeAccessCheckContext co
455455 }
456456
457457 case LIST_ROLE :
458- if (AuthorityChecker .checkSystemPermission (
459- context .getUsername (), PrivilegeType .MANAGE_ROLE )) {
458+ if (checkHasGlobalAuth (context .getUsername (), PrivilegeType .MANAGE_ROLE )) {
460459 return SUCCEED ;
461460 }
462461 // list roles of other user is not allowed
@@ -572,6 +571,36 @@ public TSStatus visitDropModel(DropModelStatement statement, TreeAccessCheckCont
572571 return checkModelManagement (context .getUsername ());
573572 }
574573
574+ @ Override
575+ public TSStatus visitCreateTraining (
576+ CreateTrainingStatement createTrainingStatement , TreeAccessCheckContext context ) {
577+ return checkModelManagement (context .getUsername ());
578+ }
579+
580+ @ Override
581+ public TSStatus visitUnloadModel (
582+ UnloadModelStatement unloadModelStatement , TreeAccessCheckContext context ) {
583+ return checkModelManagement (context .getUsername ());
584+ }
585+
586+ @ Override
587+ public TSStatus visitLoadModel (
588+ LoadModelStatement loadModelStatement , TreeAccessCheckContext context ) {
589+ return checkModelManagement (context .getUsername ());
590+ }
591+
592+ @ Override
593+ public TSStatus visitShowAIDevices (
594+ ShowAIDevicesStatement showAIDevicesStatement , TreeAccessCheckContext context ) {
595+ return checkModelManagement (context .getUsername ());
596+ }
597+
598+ @ Override
599+ public TSStatus visitShowLoadedModels (
600+ ShowLoadedModelsStatement showLoadedModelsStatement , TreeAccessCheckContext context ) {
601+ return SUCCEED ;
602+ }
603+
575604 @ Override
576605 public TSStatus visitShowModels (ShowModelsStatement statement , TreeAccessCheckContext context ) {
577606 return SUCCEED ;
@@ -691,10 +720,7 @@ private TSStatus checkTriggerManagement(String userName) {
691720 if (AuthorityChecker .SUPER_USER .equals (userName )) {
692721 return SUCCEED ;
693722 }
694- return AuthorityChecker .getTSStatus (
695- AuthorityChecker .checkSystemPermission (userName , PrivilegeType .SYSTEM )
696- || AuthorityChecker .checkSystemPermission (userName , PrivilegeType .USE_TRIGGER ),
697- PrivilegeType .SYSTEM );
723+ return checkGlobalAuth (userName , PrivilegeType .USE_TRIGGER );
698724 }
699725
700726 // ============================== database related ===========================
@@ -744,11 +770,7 @@ public TSStatus visitDeleteStorageGroup(
744770 if (AuthorityChecker .SUPER_USER .equals (context .getUsername ())) {
745771 return SUCCEED ;
746772 }
747- return AuthorityChecker .getTSStatus (
748- AuthorityChecker .checkSystemPermission (context .getUsername (), PrivilegeType .SYSTEM )
749- || AuthorityChecker .checkSystemPermission (
750- context .getUsername (), PrivilegeType .MANAGE_DATABASE ),
751- PrivilegeType .SYSTEM );
773+ return checkGlobalAuth (context .getUsername (), PrivilegeType .MANAGE_DATABASE );
752774 }
753775
754776 private TSStatus checkCreateOrAlterDatabasePermission (String userName , PartialPath databaseName ) {
@@ -762,10 +784,7 @@ private TSStatus checkCreateOrAlterDatabasePermission(String userName, PartialPa
762784 return SUCCEED ;
763785 }
764786
765- return AuthorityChecker .getTSStatus (
766- AuthorityChecker .checkSystemPermission (userName , PrivilegeType .SYSTEM )
767- || AuthorityChecker .checkSystemPermission (userName , PrivilegeType .MANAGE_DATABASE ),
768- PrivilegeType .SYSTEM );
787+ return checkGlobalAuth (userName , PrivilegeType .MANAGE_DATABASE );
769788 }
770789
771790 private TSStatus checkShowOrCountDatabasePermission (
@@ -855,6 +874,11 @@ public TSStatus visitExplainAnalyze(
855874 return statement .getQueryStatement ().accept (this , context );
856875 }
857876
877+ @ Override
878+ public TSStatus visitExplain (ExplainStatement explainStatement , TreeAccessCheckContext context ) {
879+ return explainStatement .getQueryStatement ().accept (this , context );
880+ }
881+
858882 // ============================= timeseries related =================================
859883 public static TSStatus checkTimeSeriesPermission (
860884 String userName , List <? extends PartialPath > checkedPaths , PrivilegeType permission ) {
@@ -1280,6 +1304,12 @@ public TSStatus visitShowCurrentTimestamp(
12801304 return SUCCEED ;
12811305 }
12821306
1307+ @ Override
1308+ public TSStatus visitLoadConfiguration (
1309+ LoadConfigurationStatement loadConfigurationStatement , TreeAccessCheckContext context ) {
1310+ return checkOnlySuperUser (context .getUsername ());
1311+ }
1312+
12831313 // ======================== TTL related ===========================
12841314 @ Override
12851315 public TSStatus visitSetTTL (SetTTLStatement statement , TreeAccessCheckContext context ) {
@@ -1429,4 +1459,11 @@ protected void setCanSeeAuditDB(AuthorityInformationStatement statement, String
14291459 statement .setCanSeeAuditDB (false );
14301460 }
14311461 }
1462+
1463+ private TSStatus checkOnlySuperUser (String userName ) {
1464+ if (AuthorityChecker .SUPER_USER .equals (userName )) {
1465+ return SUCCEED ;
1466+ }
1467+ return AuthorityChecker .getTSStatus (false , "Only the admin user can perform this operation" );
1468+ }
14321469}
0 commit comments