[Feature] Support setting session vars in user property (#48477)

Signed-off-by: yandongxiao <[email protected]> (cherry picked from commit 17d1914) # Conflicts: # fe/fe-core/src/main/java/com/starrocks/sql/ast/AlterUserStmt.java # fe/fe-core/src/main/java/com/starrocks/sql/ast/BaseCreateAlterUserStmt.java # fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateUserStmt.java # fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java
StarRocks · Aug 15, 2024 · 274fac9 · 274fac9
1 parent 363b693
commit 274fac9
Show file tree

Hide file tree

Showing 22 changed files with 1,234 additions and 61 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/authentication/AuthenticationMgr.java b/fe/fe-core/src/main/java/com/starrocks/authentication/AuthenticationMgr.java
@@ -338,13 +338,25 @@ public void createUser(CreateUserStmt stmt) throws DdlException {
                         + " : user " + stmt.getUserIdentity() + " already exists");
                 return;
             }
-            userToAuthenticationInfo.put(userIdentity, info);
 
             UserProperty userProperty = null;
-            if (!userNameToProperty.containsKey(userIdentity.getUser())) {
+            String userName = userIdentity.getUser();
+            if (userNameToProperty.containsKey(userName)) {
+                userProperty = userNameToProperty.get(userName);
+            } else {
                 userProperty = new UserProperty();
-                userNameToProperty.put(userIdentity.getUser(), userProperty);
             }
+
+            if (stmt.getProperties() != null) {
+                // If we create the user with properties, we need to call userProperty.update to check and update userProperty.
+                // If there are failures, update method will throw an exception
+                userProperty.update(userIdentity, UserProperty.changeToPairList(stmt.getProperties()));
+            }
+
+            // If all checks are passed, we can add the user to the userToAuthenticationInfo and userNameToProperty
+            userToAuthenticationInfo.put(userIdentity, info);
+            userNameToProperty.put(userName, userProperty);
+
             GlobalStateMgr globalStateMgr = GlobalStateMgr.getCurrentState();
             AuthorizationMgr authorizationManager = globalStateMgr.getAuthorizationMgr();
             // init user privilege
@@ -363,8 +375,10 @@ public void createUser(CreateUserStmt stmt) throws DdlException {
         }
     }
 
-    public void alterUser(UserIdentity userIdentity, UserAuthenticationInfo userAuthenticationInfo)
-            throws DdlException {
+    // This method is used to update user information, including authentication information and user properties
+    // Note: if properties is null, we should keep the original properties
+    public void alterUser(UserIdentity userIdentity, UserAuthenticationInfo userAuthenticationInfo,
+                          Map<String, String> properties) throws DdlException {
         writeLock();
         try {
             if (!userToAuthenticationInfo.containsKey(userIdentity)) {
@@ -375,26 +389,35 @@ public void alterUser(UserIdentity userIdentity, UserAuthenticationInfo userAuth
             }
 
             updateUserNoLock(userIdentity, userAuthenticationInfo, true);
-            GlobalStateMgr.getCurrentState().getEditLog().logAlterUser(userIdentity, userAuthenticationInfo);
+            if (properties != null && properties.size() > 0) {
+                UserProperty userProperty = userNameToProperty.get(userIdentity.getUser());
+                userProperty.update(userIdentity, UserProperty.changeToPairList(properties));
+            }
+            GlobalStateMgr.getCurrentState().getEditLog().logAlterUser(userIdentity, userAuthenticationInfo, properties);
         } catch (AuthenticationException e) {
             throw new DdlException("failed to alter user " + userIdentity, e);
         } finally {
             writeUnlock();
         }
     }
 
-    private void updateUserPropertyNoLock(String user, List<Pair<String, String>> properties) throws DdlException {
+    private void updateUserPropertyNoLock(String user, List<Pair<String, String>> properties, boolean isReplay)
+            throws DdlException {
         UserProperty userProperty = userNameToProperty.getOrDefault(user, null);
         if (userProperty == null) {
             throw new DdlException("user '" + user + "' doesn't exist");
         }
-        userProperty.update(properties);
+        if (isReplay) {
+            userProperty.updateForReplayJournal(properties);
+        } else {
+            userProperty.update(user, properties);
+        }
     }
 
     public void updateUserProperty(String user, List<Pair<String, String>> properties) throws DdlException {
         try {
             writeLock();
-            updateUserPropertyNoLock(user, properties);
+            updateUserPropertyNoLock(user, properties, false);
             UserPropertyInfo propertyInfo = new UserPropertyInfo(user, properties);
             GlobalStateMgr.getCurrentState().getEditLog().logUpdateUserPropertyV2(propertyInfo);
             LOG.info("finished to update user '{}' with properties: {}", user, properties);
@@ -406,16 +429,20 @@ public void updateUserProperty(String user, List<Pair<String, String>> propertie
     public void replayUpdateUserProperty(UserPropertyInfo info) throws DdlException {
         try {
             writeLock();
-            updateUserPropertyNoLock(info.getUser(), info.getProperties());
+            updateUserPropertyNoLock(info.getUser(), info.getProperties(), true);
         } finally {
             writeUnlock();
         }
     }
 
-    public void replayAlterUser(UserIdentity userIdentity, UserAuthenticationInfo info) throws AuthenticationException {
+    public void replayAlterUser(UserIdentity userIdentity, UserAuthenticationInfo info,
+                                Map<String, String> properties) throws AuthenticationException {
         writeLock();
         try {
             updateUserNoLock(userIdentity, info, true);
+            // updateForReplayJournal will catch all exceptions when replaying user properties
+            UserProperty userProperty = userNameToProperty.get(userIdentity.getUser());
+            userProperty.updateForReplayJournal(UserProperty.changeToPairList(properties));
         } finally {
             writeUnlock();
         }
@@ -485,8 +512,7 @@ public void replayCreateUser(
         }
     }
 
-    private void updateUserNoLock(
-            UserIdentity userIdentity, UserAuthenticationInfo info, boolean shouldExists)
+    private void updateUserNoLock(UserIdentity userIdentity, UserAuthenticationInfo info, boolean shouldExists)
             throws AuthenticationException {
         if (userToAuthenticationInfo.containsKey(userIdentity)) {
             if (!shouldExists) {
@@ -650,4 +676,24 @@ public void loadV2(SRMetaBlockReader reader) throws IOException, SRMetaBlockExce
         this.nameToSecurityIntegrationMap = ret.nameToSecurityIntegrationMap;
         this.userToAuthenticationInfo = ret.userToAuthenticationInfo;
     }
+
+    public UserProperty getUserProperty(String userName) {
+        UserProperty userProperty = userNameToProperty.get(userName);
+        if (userProperty == null) {
+            throw new SemanticException("Unknown user: " + userName);
+        }
+        return userProperty;
+    }
+
+    public UserIdentity getUserIdentityByName(String userName) {
+        Map<UserIdentity, UserAuthenticationInfo> userToAuthInfo = getUserToAuthenticationInfo();
+        Map.Entry<UserIdentity, UserAuthenticationInfo> matchedUserIdentity = userToAuthInfo.entrySet().stream()
+                .filter(entry -> (entry.getKey().getUser().equals(userName)))
+                .findFirst().orElse(null);
+        if (matchedUserIdentity == null) {
+            throw new SemanticException("Unknown user: " + userName);
+        }
+
+        return matchedUserIdentity.getKey();
+    }
 }