From 274fac92cb118ae547c4b3c8306d39a0d05b2b96 Mon Sep 17 00:00:00 2001 From: yandongxiao Date: Fri, 2 Aug 2024 14:12:26 +0800 Subject: [PATCH] [Feature] Support setting session vars in user property (#48477) Signed-off-by: yandongxiao (cherry picked from commit 17d191495d431edde7c6d1b890f6af57b7638b1c) # 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 --- .../authentication/AuthenticationMgr.java | 72 ++- .../authentication/UserProperty.java | 301 +++++++++-- .../starrocks/mysql/nio/AcceptListener.java | 9 + .../com/starrocks/persist/AlterUserInfo.java | 13 + .../java/com/starrocks/persist/EditLog.java | 7 +- .../starrocks/privilege/AuthorizationMgr.java | 5 + .../java/com/starrocks/qe/ConnectContext.java | 41 ++ .../com/starrocks/qe/DDLStmtExecutor.java | 2 +- .../com/starrocks/qe/ExecuteAsExecutor.java | 11 +- .../java/com/starrocks/qe/SetExecutor.java | 2 +- .../java/com/starrocks/qe/VariableMgr.java | 14 + .../analyzer/ShowUserPropertyAnalyzer.java | 2 +- .../com/starrocks/sql/ast/AlterUserStmt.java | 10 + .../sql/ast/BaseCreateAlterUserStmt.java | 25 + .../com/starrocks/sql/ast/CreateUserStmt.java | 10 + .../sql/ast/ShowUserPropertyStmt.java | 11 +- .../com/starrocks/sql/parser/AstBuilder.java | 38 ++ .../com/starrocks/sql/parser/StarRocks.g4 | 6 +- .../AuthenticationManagerTest.java | 203 ++++++- .../authentication/UserPropertyTest.java | 499 ++++++++++++++++++ .../sql/analyzer/AnalyzeStmtTest.java | 2 +- .../starrocks/sql/ast/ExecuteAsStmtTest.java | 12 + 22 files changed, 1234 insertions(+), 61 deletions(-) create mode 100644 fe/fe-core/src/test/java/com/starrocks/authentication/UserPropertyTest.java 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 index ea640c074a26f..39693b651bf19 100644 --- 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 properties) throws DdlException { writeLock(); try { if (!userToAuthenticationInfo.containsKey(userIdentity)) { @@ -375,7 +389,11 @@ 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 { @@ -383,18 +401,23 @@ public void alterUser(UserIdentity userIdentity, UserAuthenticationInfo userAuth } } - private void updateUserPropertyNoLock(String user, List> properties) throws DdlException { + private void updateUserPropertyNoLock(String user, List> 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> 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> 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 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 userToAuthInfo = getUserToAuthenticationInfo(); + Map.Entry 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(); + } } diff --git a/fe/fe-core/src/main/java/com/starrocks/authentication/UserProperty.java b/fe/fe-core/src/main/java/com/starrocks/authentication/UserProperty.java index d0a2cf350e0aa..43baaeb5afc2b 100644 --- a/fe/fe-core/src/main/java/com/starrocks/authentication/UserProperty.java +++ b/fe/fe-core/src/main/java/com/starrocks/authentication/UserProperty.java @@ -15,66 +15,303 @@ package com.starrocks.authentication; +import com.google.common.collect.Lists; import com.google.gson.annotations.SerializedName; +import com.starrocks.analysis.StringLiteral; +import com.starrocks.catalog.Database; +import com.starrocks.catalog.InternalCatalog; import com.starrocks.common.Config; import com.starrocks.common.DdlException; +import com.starrocks.common.ErrorCode; +import com.starrocks.common.ErrorReport; import com.starrocks.common.Pair; -import com.starrocks.sql.ast.SetUserPropertyVar; +import com.starrocks.connector.exception.StarRocksConnectorException; +import com.starrocks.qe.SessionVariable; +import com.starrocks.qe.VariableMgr; +import com.starrocks.server.CatalogMgr; +import com.starrocks.server.GlobalStateMgr; +import com.starrocks.server.MetadataMgr; +import com.starrocks.sql.ast.SystemVariable; +import com.starrocks.sql.ast.UserIdentity; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import java.lang.reflect.Field; import java.util.List; +import java.util.Map; +import java.util.TreeMap; +// UserProperty is a class that represents the properties that are identified. public class UserProperty { + private static final Logger LOG = LogManager.getLogger(UserProperty.class); + + // Because session variables does not include these two properties, we define them here. + public static final String PROP_MAX_USER_CONNECTIONS = "max_user_connections"; + public static final String PROP_DATABASE = "database"; + // In order to keep consistent with database, we support user to set session.catalog = xxx or catalog = yyy + public static final String PROP_CATALOG = SessionVariable.CATALOG; + public static final String PROP_SESSION_PREFIX = "session."; + + public static final long MAX_CONN_DEFAULT_VALUE = 1024; + public static final String CATALOG_DEFAULT_VALUE = InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME; + public static final String DATABASE_DEFAULT_VALUE = ""; + + // If the values is empty, we remove the key from the session variables. + public static final String EMPTY_VALUE = ""; + @SerializedName(value = "m") - private long maxConn = 1024; + private long maxConn = MAX_CONN_DEFAULT_VALUE; - private static final String PROP_MAX_USER_CONNECTIONS = "max_user_connections"; + @SerializedName(value = "d") + private String database = DATABASE_DEFAULT_VALUE; - public long getMaxConn() { - return maxConn; - } + @SerializedName(value = "c") + private String catalog = CATALOG_DEFAULT_VALUE; + + @SerializedName(value = "s") + private Map sessionVariables = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); - public void setMaxConn(long maxConn) { - this.maxConn = maxConn; + public void update(String userName, List> properties) throws DdlException { + AuthenticationMgr authenticationMgr = GlobalStateMgr.getCurrentState().getAuthenticationMgr(); + UserIdentity user = authenticationMgr.getUserIdentityByName(userName); + update(user, properties); } - public void update(List> properties) throws DdlException { - // copy - long newMaxConn = maxConn; + // update the user properties + // we should check the properties and throw exceptions if the properties are invalid + public void update(UserIdentity user, List> properties) throws DdlException { + if (properties == null || properties.isEmpty()) { + return; + } - // update + String newDatabase = ""; for (Pair entry : properties) { String key = entry.first; String value = entry.second; - String[] keyArr = key.split("\\" + SetUserPropertyVar.DOT_SEPARATOR); - if (keyArr[0].equalsIgnoreCase(PROP_MAX_USER_CONNECTIONS)) { - if (keyArr.length != 1) { - throw new DdlException(PROP_MAX_USER_CONNECTIONS + " format error"); + if (key.equalsIgnoreCase(PROP_MAX_USER_CONNECTIONS)) { + long newMaxConn = checkMaxConn(value); + setMaxConn(newMaxConn); + } else if (key.equalsIgnoreCase(PROP_DATABASE)) { + // we do not check database existence here, because we should + // check catalog existence first. + newDatabase = value; + } else if (key.equalsIgnoreCase(PROP_CATALOG)) { + checkCatalog(value); + setCatalog(value); + } else if (key.startsWith(PROP_SESSION_PREFIX)) { + String sessionKey = key.substring(PROP_SESSION_PREFIX.length()); + if (sessionKey.equalsIgnoreCase(PROP_CATALOG)) { + checkCatalog(value); + setCatalog(value); + } else { + checkSessionVariable(sessionKey, value); + setSessionVariable(sessionKey, value); } + } else { + throw new DdlException("Unknown user property(" + key + ")"); + } + } + if (!newDatabase.isEmpty()) { + checkDatabase(newDatabase); + setDatabase(newDatabase); + } + } - try { - newMaxConn = Long.parseLong(value); - } catch (NumberFormatException e) { - throw new DdlException(PROP_MAX_USER_CONNECTIONS + " is not a number"); + // We do not check the variable default_session_database and default_session_catalog here, because we have checked them + // when set properties. And we never should throw exceptions, this may cause the system can be started normally. + public void updateForReplayJournal(List> properties) { + for (Pair entry : properties) { + try { + String key = entry.first; + String value = entry.second; + if (key.equalsIgnoreCase(PROP_MAX_USER_CONNECTIONS)) { + long maxConn = checkMaxConn(value); + setMaxConn(maxConn); + } else if (key.equalsIgnoreCase(PROP_DATABASE)) { + setDatabase(value); + } else if (key.equalsIgnoreCase(PROP_CATALOG)) { + setCatalog(value); + } else if (key.startsWith(PROP_SESSION_PREFIX)) { + String sessionKey = key.substring(PROP_SESSION_PREFIX.length()); + if (sessionKey.equalsIgnoreCase(PROP_CATALOG)) { + setCatalog(value); + } else { + setSessionVariable(sessionKey, value); + } } + } catch (Exception e) { + // we should never throw an exception when replaying journal + LOG.warn("update user property from journal failed: ", e); + } + } + } - if (newMaxConn <= 0 || newMaxConn > 10000) { - throw new DdlException(PROP_MAX_USER_CONNECTIONS + - " is not valid, the value must be between 1 and 10000"); - } - if (newMaxConn > Config.qe_max_connection) { - throw new DdlException( - PROP_MAX_USER_CONNECTIONS + - " is not valid, the value must be less than qe_max_connection(" - + Config.qe_max_connection + ")"); + public String getCatalogDbName() { + return getCatalog() + "." + getDatabase(); + } + + public long getMaxConn() { + return maxConn; + } + + public String getDatabase() { + return database; + } + + public void setDatabase(String sessionDatabase) { + if (sessionDatabase.equalsIgnoreCase(EMPTY_VALUE)) { + this.database = DATABASE_DEFAULT_VALUE; + } else { + this.database = sessionDatabase; + } + } + + public Map getSessionVariables() { + return sessionVariables; + } + + public void setSessionVariables(Map sessions) { + this.sessionVariables = sessions; + } + + // check the session variable + private void checkSessionVariable(String sessionKey, String value) throws DdlException { + if (value.equalsIgnoreCase(EMPTY_VALUE)) { + return; + } + // check whether the variable exists + SystemVariable variable = new SystemVariable(sessionKey, new StringLiteral(value)); + VariableMgr.checkSystemVariableExist(variable); + + // check whether the value is valid + Field field = VariableMgr.getField(sessionKey); + if (field == null || !canAssignValue(field, value)) { + ErrorReport.reportDdlException(ErrorCode.ERR_WRONG_TYPE_FOR_VAR, value); + } + + // check flags of the variable, e.g. whether the variable is read-only + VariableMgr.checkUpdate(variable); + } + + // check whether the catalog exist + private void checkCatalog(String catalogName) throws DdlException { + if (catalogName.equalsIgnoreCase(EMPTY_VALUE)) { + return; + } + + if (!CatalogMgr.isInternalCatalog(catalogName)) { + if (!GlobalStateMgr.getCurrentState().getCatalogMgr().catalogExists(catalogName)) { + ErrorReport.reportDdlException(ErrorCode.ERR_BAD_CATALOG_ERROR, catalogName); + } + } + } + + // check whether the database exist + // we need to reset the database if it checks failed + private void checkDatabase(String newDatabase) { + if (newDatabase.equalsIgnoreCase(DATABASE_DEFAULT_VALUE)) { + return; + } + + // check whether the database exists + MetadataMgr metadataMgr = GlobalStateMgr.getCurrentState().getMetadataMgr(); + Database db = metadataMgr.getDb(getCatalog(), newDatabase); + if (db == null) { + String catalogDbName = getCatalogDbName(); + throw new StarRocksConnectorException(catalogDbName + " not exists"); + } + } + + + public static List> changeToPairList(Map properties) { + List> list = Lists.newArrayList(); + if (properties == null || properties.size() == 0) { + return list; + } + + for (Map.Entry entry : properties.entrySet()) { + list.add(Pair.create(entry.getKey(), entry.getValue())); + } + return list; + } + + private boolean canAssignValue(Field field, String value) { + Class fieldType = field.getType(); + try { + if (fieldType == int.class || fieldType == Integer.class) { + Integer.parseInt(value); + } else if (fieldType == boolean.class || fieldType == Boolean.class) { + if (!value.equalsIgnoreCase("true") && !value.equalsIgnoreCase("false")) { + throw new IllegalArgumentException("Invalid boolean value"); } + } else if (fieldType == byte.class || fieldType == Byte.class) { + Byte.parseByte(value); + } else if (fieldType == short.class || fieldType == Short.class) { + Short.parseShort(value); + } else if (fieldType == long.class || fieldType == Long.class) { + Long.parseLong(value); + } else if (fieldType == float.class || fieldType == Float.class) { + Float.parseFloat(value); + } else if (fieldType == double.class || fieldType == Double.class) { + Double.parseDouble(value); + } else if (fieldType == String.class) { + return true; } else { - throw new DdlException("Unknown user property(" + key + ")"); + return false; } + return true; + } catch (Exception e) { + return false; } + } + + private void setSessionVariable(String sessionKey, String value) { + if (value.equalsIgnoreCase(EMPTY_VALUE)) { + sessionVariables.remove(sessionKey); + } else { + sessionVariables.put(sessionKey, value); + } + } + + public String getCatalog() { + return catalog; + } + + public void setCatalog(String catalog) { + if (catalog.equalsIgnoreCase(EMPTY_VALUE)) { + this.catalog = CATALOG_DEFAULT_VALUE; + } else { + this.catalog = catalog; + } + } + + private long checkMaxConn(String value) throws DdlException { + if (value.equalsIgnoreCase(EMPTY_VALUE)) { + return MAX_CONN_DEFAULT_VALUE; + } + + try { + long newMaxConn = Long.parseLong(value); + + if (newMaxConn <= 0 || newMaxConn > 10000) { + throw new DdlException(PROP_MAX_USER_CONNECTIONS + " is not valid, the value must be between 1 and 10000"); + } + + if (newMaxConn > Config.qe_max_connection) { + throw new DdlException( + PROP_MAX_USER_CONNECTIONS + " is not valid, the value must be less than qe_max_connection(" + + Config.qe_max_connection + ")"); + } + + return newMaxConn; + } catch (NumberFormatException e) { + throw new DdlException(PROP_MAX_USER_CONNECTIONS + " is not a number"); + } + } - // set - maxConn = newMaxConn; + private void setMaxConn(long value) { + maxConn = value; } } diff --git a/fe/fe-core/src/main/java/com/starrocks/mysql/nio/AcceptListener.java b/fe/fe-core/src/main/java/com/starrocks/mysql/nio/AcceptListener.java index 9c035d962ee83..7c70cbdbd7800 100644 --- a/fe/fe-core/src/main/java/com/starrocks/mysql/nio/AcceptListener.java +++ b/fe/fe-core/src/main/java/com/starrocks/mysql/nio/AcceptListener.java @@ -33,6 +33,7 @@ // under the License. package com.starrocks.mysql.nio; +import com.starrocks.authentication.UserProperty; import com.starrocks.common.Pair; import com.starrocks.common.util.LogUtil; import com.starrocks.mysql.MysqlProto; @@ -98,6 +99,14 @@ public void handleEvent(AcceptingChannel channel) { if (registerResult.first) { connection.setCloseListener( streamConnection -> connectScheduler.unregisterConnection(context)); + + // We place the set session environment code here, because we want to notify user if there + // are some errors when setting session environment. + // Unfortunately, the client cannot receive the message. + UserProperty userProperty = context.getGlobalStateMgr().getAuthenticationMgr() + .getUserProperty(context.getCurrentUserIdentity().getUser()); + context.updateByUserProperty(userProperty); + MysqlProto.sendResponsePacket(context); } else { context.getState().setError(registerResult.second); diff --git a/fe/fe-core/src/main/java/com/starrocks/persist/AlterUserInfo.java b/fe/fe-core/src/main/java/com/starrocks/persist/AlterUserInfo.java index 2588eb8dddce8..213882adfdcab 100644 --- a/fe/fe-core/src/main/java/com/starrocks/persist/AlterUserInfo.java +++ b/fe/fe-core/src/main/java/com/starrocks/persist/AlterUserInfo.java @@ -26,6 +26,7 @@ import java.io.DataInput; import java.io.DataOutput; import java.io.IOException; +import java.util.Map; public class AlterUserInfo implements Writable { @SerializedName(value = "u") @@ -33,11 +34,19 @@ public class AlterUserInfo implements Writable { @SerializedName(value = "a") UserAuthenticationInfo authenticationInfo; + @SerializedName(value = "p") + Map properties; + public AlterUserInfo(UserIdentity userIdentity, UserAuthenticationInfo authenticationInfo) { this.userIdentity = userIdentity; this.authenticationInfo = authenticationInfo; } + public AlterUserInfo(UserIdentity userIdentity, UserAuthenticationInfo authenticationInfo, Map properties) { + this(userIdentity, authenticationInfo); + this.properties = properties; + } + public UserIdentity getUserIdentity() { return userIdentity; } @@ -46,6 +55,10 @@ public UserAuthenticationInfo getAuthenticationInfo() { return authenticationInfo; } + public Map getProperties() { + return properties; + } + @Override public void write(DataOutput out) throws IOException { Text.writeString(out, GsonUtils.GSON.toJson(this)); diff --git a/fe/fe-core/src/main/java/com/starrocks/persist/EditLog.java b/fe/fe-core/src/main/java/com/starrocks/persist/EditLog.java index 46954e659b82f..e7a8f93a1fde4 100644 --- a/fe/fe-core/src/main/java/com/starrocks/persist/EditLog.java +++ b/fe/fe-core/src/main/java/com/starrocks/persist/EditLog.java @@ -1070,7 +1070,7 @@ public void loadJournal(GlobalStateMgr globalStateMgr, JournalEntity journal) case OperationType.OP_ALTER_USER_V2: { AlterUserInfo info = (AlterUserInfo) journal.getData(); globalStateMgr.getAuthenticationMgr().replayAlterUser( - info.getUserIdentity(), info.getAuthenticationInfo()); + info.getUserIdentity(), info.getAuthenticationInfo(), info.getProperties()); break; } case OperationType.OP_UPDATE_USER_PROP_V2: @@ -1872,8 +1872,9 @@ public void logCreateUser( logEdit(OperationType.OP_CREATE_USER_V2, info); } - public void logAlterUser(UserIdentity userIdentity, UserAuthenticationInfo authenticationInfo) { - AlterUserInfo info = new AlterUserInfo(userIdentity, authenticationInfo); + public void logAlterUser(UserIdentity userIdentity, UserAuthenticationInfo authenticationInfo, + Map properties) { + AlterUserInfo info = new AlterUserInfo(userIdentity, authenticationInfo, properties); logEdit(OperationType.OP_ALTER_USER_V2, info); } diff --git a/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java b/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java index beda9ac96031f..391f7aeaf6532 100644 --- a/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java +++ b/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java @@ -1766,4 +1766,9 @@ public void saveV2(DataOutputStream dos) throws IOException { throw new IOException("failed to save AuthenticationManager!", e); } } + + // get all role ids of the user, including the default roles and the inactivated roles + public Set getAllRoleIds(UserIdentity user) throws PrivilegeException { + return getRoleIdsByUser(user); + } } diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java b/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java index a1b1b26585119..c7ee2cda1cd73 100644 --- a/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java @@ -38,6 +38,8 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; +import com.starrocks.analysis.StringLiteral; +import com.starrocks.authentication.UserProperty; import com.starrocks.cluster.ClusterNamespace; import com.starrocks.common.DdlException; import com.starrocks.common.ErrorCode; @@ -1112,6 +1114,45 @@ public void cleanTemporaryTable() { } } + // We can not make sure the set variables are all valid. Even if some variables are invalid, we should let user continue + // to execute SQL. + public void updateByUserProperty(UserProperty userProperty) { + try { + // set session variables + Map sessionVariables = userProperty.getSessionVariables(); + for (Map.Entry entry : sessionVariables.entrySet()) { + SystemVariable variable = new SystemVariable(entry.getKey(), new StringLiteral(entry.getValue())); + modifySystemVariable(variable, true); + } + + // set catalog and database + boolean dbHasBeenSetByUser = !getCurrentCatalog().equals(VariableMgr.getDefaultValue(SessionVariable.CATALOG)) || + !getDatabase().isEmpty(); + if (!dbHasBeenSetByUser) { + String catalog = userProperty.getCatalog(); + String database = userProperty.getDatabase(); + if (catalog.equals(UserProperty.CATALOG_DEFAULT_VALUE)) { + if (!database.equals(UserProperty.DATABASE_DEFAULT_VALUE)) { + changeCatalogDb(userProperty.getCatalogDbName()); + } + } else { + if (database.equals(UserProperty.DATABASE_DEFAULT_VALUE)) { + changeCatalog(catalog); + } else { + changeCatalogDb(userProperty.getCatalogDbName()); + } + SystemVariable variable = new SystemVariable(SessionVariable.CATALOG, new StringLiteral(catalog)); + modifySystemVariable(variable, true); + } + } + } catch (Exception e) { + LOG.warn("set session env failed: ", e); + // In handshake, we will send error message to client. But it seems that client will ignore it. + getState().setOk(0L, 0, + String.format("set session variables from user property failed: %s", e.getMessage())); + } + } + /** * Set thread-local context for the scope, and remove it after leaving the scope */ diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/DDLStmtExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/DDLStmtExecutor.java index fd63b29da761a..0bc2e5c282f8c 100644 --- a/fe/fe-core/src/main/java/com/starrocks/qe/DDLStmtExecutor.java +++ b/fe/fe-core/src/main/java/com/starrocks/qe/DDLStmtExecutor.java @@ -506,7 +506,7 @@ public ShowResultSet visitCreateUserStatement(CreateUserStmt stmt, ConnectContex public ShowResultSet visitAlterUserStatement(AlterUserStmt stmt, ConnectContext context) { ErrorReport.wrapWithRuntimeException(() -> { context.getGlobalStateMgr().getAuthenticationMgr() - .alterUser(stmt.getUserIdentity(), stmt.getAuthenticationInfo()); + .alterUser(stmt.getUserIdentity(), stmt.getAuthenticationInfo(), stmt.getProperties()); }); return null; } diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteAsExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteAsExecutor.java index 030b323d7c199..ac656cc5d09c8 100644 --- a/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteAsExecutor.java +++ b/fe/fe-core/src/main/java/com/starrocks/qe/ExecuteAsExecutor.java @@ -16,7 +16,9 @@ package com.starrocks.qe; import com.google.common.base.Preconditions; +import com.starrocks.authentication.UserProperty; import com.starrocks.sql.ast.ExecuteAsStmt; +import com.starrocks.sql.ast.UserIdentity; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -38,7 +40,12 @@ public static void execute(ExecuteAsStmt stmt, ConnectContext ctx) { Preconditions.checkArgument(!stmt.isAllowRevert()); LOG.info("{} EXEC AS {} from now on", ctx.getCurrentUserIdentity(), stmt.getToUser()); - ctx.setCurrentUserIdentity(stmt.getToUser()); - ctx.setCurrentRoleIds(stmt.getToUser()); + UserIdentity user = stmt.getToUser(); + ctx.setCurrentUserIdentity(user); + ctx.setCurrentRoleIds(user); + + UserProperty userProperty = ctx.getGlobalStateMgr().getAuthenticationMgr() + .getUserProperty(user.getUser()); + ctx.updateByUserProperty(userProperty); } } diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/SetExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/SetExecutor.java index fe60cf26bb82b..07071eb9840a0 100644 --- a/fe/fe-core/src/main/java/com/starrocks/qe/SetExecutor.java +++ b/fe/fe-core/src/main/java/com/starrocks/qe/SetExecutor.java @@ -85,7 +85,7 @@ private void setVariablesOfAllType(SetListItem var) throws DdlException { } userAuthenticationInfo.setPassword(setPassVar.getPassword()); GlobalStateMgr.getCurrentState().getAuthenticationMgr() - .alterUser(setPassVar.getUserIdent(), userAuthenticationInfo); + .alterUser(setPassVar.getUserIdent(), userAuthenticationInfo, null); } } diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/VariableMgr.java b/fe/fe-core/src/main/java/com/starrocks/qe/VariableMgr.java index e83e018c8a56d..60b726f08ae6d 100644 --- a/fe/fe-core/src/main/java/com/starrocks/qe/VariableMgr.java +++ b/fe/fe-core/src/main/java/com/starrocks/qe/VariableMgr.java @@ -272,6 +272,12 @@ public static SessionVariable newSessionVariable() { return (SessionVariable) DEFAULT_SESSION_VARIABLE.clone(); } + // Check if this sessionVariable can be set correctly + public static void checkUpdate(SystemVariable sessionVariable) throws DdlException { + VarContext ctx = VariableMgr.getVarContext(sessionVariable.getVariable()); + checkUpdate(sessionVariable, ctx.getFlag()); + } + // Check if this setVar can be set correctly private static void checkUpdate(SystemVariable setVar, int flag) throws DdlException { if ((flag & READ_ONLY) != 0) { @@ -661,6 +667,14 @@ public static boolean shouldForwardToLeader(String name) { } } + public static Field getField(String name) { + VarContext ctx = getVarContext(name); + if (ctx == null) { + return null; + } + return ctx.getField(); + } + @Retention(RetentionPolicy.RUNTIME) public @interface VarAttr { // Name in show variables and set statement; diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ShowUserPropertyAnalyzer.java b/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ShowUserPropertyAnalyzer.java index 3b1240fe24a14..eb9d23531f6c9 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ShowUserPropertyAnalyzer.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ShowUserPropertyAnalyzer.java @@ -23,7 +23,7 @@ public class ShowUserPropertyAnalyzer { public static void analyze(ShowUserPropertyStmt statment, ConnectContext context) { String user = statment.getUser(); if (Strings.isNullOrEmpty(user)) { - statment.setUser(context.getQualifiedUser()); + statment.setUser(context.getCurrentUserIdentity().getUser()); } else { statment.setUser(user); } diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/ast/AlterUserStmt.java b/fe/fe-core/src/main/java/com/starrocks/sql/ast/AlterUserStmt.java index 36d9fdb72bdc1..fbb39047c2f50 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/ast/AlterUserStmt.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/ast/AlterUserStmt.java @@ -17,17 +17,27 @@ import com.starrocks.analysis.UserDesc; import com.starrocks.sql.parser.NodePosition; +<<<<<<< HEAD import java.util.Collections; +======= +import java.util.Map; +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) public class AlterUserStmt extends BaseCreateAlterUserStmt { private final boolean ifExists; +<<<<<<< HEAD public AlterUserStmt(UserDesc userDesc, boolean ifExists) { this(userDesc, ifExists, NodePosition.ZERO); } public AlterUserStmt(UserDesc userDesc, boolean ifExists, NodePosition pos) { super(userDesc, null, Collections.emptyList(), pos); +======= + public AlterUserStmt(UserIdentity userIdentity, boolean ifExists, UserAuthOption userAuthOption, + Map properties, NodePosition pos) { + super(userIdentity, userAuthOption, properties, pos); +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) this.ifExists = ifExists; } diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/ast/BaseCreateAlterUserStmt.java b/fe/fe-core/src/main/java/com/starrocks/sql/ast/BaseCreateAlterUserStmt.java index 6353dc4f45e23..bf7d45b834933 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/ast/BaseCreateAlterUserStmt.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/ast/BaseCreateAlterUserStmt.java @@ -18,7 +18,16 @@ import com.starrocks.authentication.UserAuthenticationInfo; import com.starrocks.sql.parser.NodePosition; +<<<<<<< HEAD import java.util.List; +======= +import java.util.Map; + +// CreateUserStmt and AlterUserStmt share the same parameter and check logic +public abstract class BaseCreateAlterUserStmt extends DdlStmt { + protected UserIdentity userIdentity; + protected UserAuthOption authOption; +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) // CreateUserStmt and AlterUserStmt share the same parameter and check logic public class BaseCreateAlterUserStmt extends DdlStmt { @@ -33,6 +42,7 @@ public class BaseCreateAlterUserStmt extends DdlStmt { // used in new RBAC privilege framework private UserAuthenticationInfo authenticationInfo = null; +<<<<<<< HEAD @Deprecated protected String userForAuthPlugin; @Deprecated @@ -53,6 +63,17 @@ public BaseCreateAlterUserStmt(UserDesc userDesc, SetRoleType setRoleType, List< this.setRoleType = setRoleType; this.defaultRoles = defaultRoles; +======= + private final Map properties; + + public BaseCreateAlterUserStmt(UserIdentity userIdentity, UserAuthOption authOption, + Map properties, NodePosition pos) { + super(pos); + + this.userIdentity = userIdentity; + this.authOption = authOption; + this.properties = properties; +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) } public UserIdentity getUserIdentity() { @@ -87,6 +108,10 @@ public void setAuthenticationInfo(UserAuthenticationInfo authenticationInfo) { this.authenticationInfo = authenticationInfo; } + public Map getProperties() { + return properties; + } + @Override public R accept(AstVisitor visitor, C context) { return visitor.visitBaseCreateAlterUserStmt(this, context); diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateUserStmt.java b/fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateUserStmt.java index f59a3c6fcdcf5..1f451f5ef1343 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateUserStmt.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/ast/CreateUserStmt.java @@ -18,6 +18,7 @@ import com.starrocks.sql.parser.NodePosition; import java.util.List; +import java.util.Map; /* * We support the following create user stmts: @@ -35,6 +36,7 @@ */ public class CreateUserStmt extends BaseCreateAlterUserStmt { +<<<<<<< HEAD private final boolean ifNotExists; public CreateUserStmt(boolean ifNotExists, UserDesc userDesc, List defaultRoles) { @@ -43,6 +45,14 @@ public CreateUserStmt(boolean ifNotExists, UserDesc userDesc, List defau public CreateUserStmt(boolean ifNotExists, UserDesc userDesc, List defaultRoles, NodePosition pos) { super(userDesc, SetRoleType.ROLE, defaultRoles, pos); +======= + public CreateUserStmt(UserIdentity userIdentity, boolean ifNotExists, + UserAuthOption authOption, + List defaultRoles, + Map properties, + NodePosition pos) { + super(userIdentity, authOption, properties, pos); +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) this.ifNotExists = ifNotExists; } diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/ast/ShowUserPropertyStmt.java b/fe/fe-core/src/main/java/com/starrocks/sql/ast/ShowUserPropertyStmt.java index 90fd2defb0155..8197f2c639835 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/ast/ShowUserPropertyStmt.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/ast/ShowUserPropertyStmt.java @@ -18,6 +18,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.starrocks.authentication.AuthenticationMgr; +import com.starrocks.authentication.UserProperty; import com.starrocks.catalog.Column; import com.starrocks.catalog.ScalarType; import com.starrocks.common.CaseSensibility; @@ -29,6 +30,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; // Show Property Stmt // syntax: @@ -70,9 +72,12 @@ public List> getRows(ConnectContext connectContext) { List> rows = new ArrayList<>(); AuthenticationMgr authenticationManager = GlobalStateMgr.getCurrentState().getAuthenticationMgr(); - // Currently only "max_user_connections" is supported - long maxConn = authenticationManager.getMaxConn(user); - rows.add(Lists.newArrayList("max_user_connections", String.valueOf(maxConn))); + UserProperty userProperty = authenticationManager.getUserProperty(user); + rows.add(Lists.newArrayList(UserProperty.PROP_MAX_USER_CONNECTIONS, String.valueOf(userProperty.getMaxConn()))); + rows.add(Lists.newArrayList(UserProperty.PROP_DATABASE, userProperty.getDatabase())); + for (Map.Entry entry : userProperty.getSessionVariables().entrySet()) { + rows.add(Lists.newArrayList(String.format("%s.%s", "session", entry.getKey()), entry.getValue())); + } if (pattern == null) { return rows; diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java b/fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java index d95c4cc272d93..6a3f46565492f 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java @@ -82,6 +82,7 @@ import com.starrocks.analysis.UserVariableExpr; import com.starrocks.analysis.VarBinaryLiteral; import com.starrocks.analysis.VariableExpr; +import com.starrocks.authentication.UserProperty; import com.starrocks.catalog.AggregateType; import com.starrocks.catalog.ArrayType; import com.starrocks.catalog.FunctionSet; @@ -3580,6 +3581,9 @@ public ParseNode visitSetUserPropertyStatement(StarRocksParser.SetUserPropertySt List propertyList = visit(context.userPropertyList().property(), Property.class); for (Property property : propertyList) { SetUserPropertyVar setVar = new SetUserPropertyVar(property.getKey(), property.getValue()); + if (!property.getKey().equalsIgnoreCase(UserProperty.PROP_MAX_USER_CONNECTIONS)) { + throw new ParsingException("Please use ALTER USER syntax to set user properties."); + } list.add(setVar); } } @@ -5255,8 +5259,20 @@ public ParseNode visitCreateUserStatement(StarRocksParser.CreateUserStatementCon roles.addAll(context.roleList().identifierOrString().stream().map(this::visit).map( s -> ((Identifier) s).getValue()).collect(toList())); } +<<<<<<< HEAD return new CreateUserStmt(ifNotExists, userDesc, roles, createPos(context)); +======= + + Map properties = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); + if (context.properties() != null) { + List propertyList = visit(context.properties().property(), Property.class); + for (Property property : propertyList) { + properties.put(property.getKey(), property.getValue()); + } + } + return new CreateUserStmt(user, ifNotExists, authOption, roles, properties, createPos(context)); +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) } @Override @@ -5290,6 +5306,7 @@ public ParseNode visitAlterUserStatement(StarRocksParser.AlterUserStatementConte return new SetDefaultRoleStmt(user, setRoleType, roles, createPos(context)); } +<<<<<<< HEAD stop = context.authOption().stop; UserAuthOption authOption = (UserAuthOption) visit(context.authOption()); if (authOption.getAuthPlugin() == null) { @@ -5300,6 +5317,27 @@ public ParseNode visitAlterUserStatement(StarRocksParser.AlterUserStatementConte authOption.isPasswordPlain(), createPos(start, stop)); } return new AlterUserStmt(userDesc, context.EXISTS() != null, createPos(context)); +======= + if (context.authOption() != null) { + UserAuthOption authOption = (UserAuthOption) visitIfPresent(context.authOption()); + Map properties = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); + if (context.properties() != null) { + List propertyList = visit(context.properties().property(), Property.class); + for (Property property : propertyList) { + properties.put(property.getKey(), property.getValue()); + } + } + return new AlterUserStmt(user, context.EXISTS() != null, authOption, properties, createPos(context)); + } + + // handle alter user xxx set properties + List list = new ArrayList<>(); + List propertyList = visit(context.properties().property(), Property.class); + for (Property property : propertyList) { + list.add(new SetUserPropertyVar(property.getKey(), property.getValue())); + } + return new SetUserPropertyStmt(user.getUser(), list, createPos(context)); +>>>>>>> 17d191495d ([Feature] Support setting session vars in user property (#48477)) } @Override diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4 b/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4 index a63f1c08ddedb..84597dda662c3 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4 +++ b/fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4 @@ -1458,6 +1458,7 @@ showTriggersStatement showUserPropertyStatement : SHOW PROPERTY (FOR string)? (LIKE string)? + | SHOW PROPERTIES (FOR string)? (LIKE string)? ; showVariablesStatement @@ -1475,7 +1476,7 @@ helpStatement // ------------------------------------------- Authz Statement ----------------------------------------------------- createUserStatement - : CREATE USER (IF NOT EXISTS)? user authOption? (DEFAULT ROLE roleList)? + : CREATE USER (IF NOT EXISTS)? user authOption? (DEFAULT ROLE roleList)? properties? ; dropUserStatement @@ -1485,6 +1486,7 @@ dropUserStatement alterUserStatement : ALTER USER (IF EXISTS)? user authOption | ALTER USER (IF EXISTS)? user DEFAULT ROLE (NONE| ALL | roleList) + | ALTER USER (IF EXISTS)? user SET properties ; showUserStatement @@ -2691,4 +2693,4 @@ nonReserved | DOTDOTDOT | NGRAMBF | FIELD | ARRAY_ELEMENT - ; \ No newline at end of file + ; diff --git a/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java b/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java index 2c808fbfff3ad..f6ff1d5847b50 100644 --- a/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java +++ b/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java @@ -25,8 +25,10 @@ import com.starrocks.qe.ConnectContext; import com.starrocks.qe.DDLStmtExecutor; import com.starrocks.qe.SetDefaultRoleExecutor; +import com.starrocks.server.CatalogMgr; import com.starrocks.sql.analyzer.SemanticException; import com.starrocks.sql.ast.AlterUserStmt; +import com.starrocks.sql.ast.CreateCatalogStmt; import com.starrocks.sql.ast.CreateRoleStmt; import com.starrocks.sql.ast.CreateUserStmt; import com.starrocks.sql.ast.DropUserStmt; @@ -255,6 +257,71 @@ public void testCreateUserWithDefaultRole() throws Exception { Assert.assertEquals(2, s.size()); } + @Test + public void testCreateUserPersistWithProperties() throws Exception { + AuthenticationMgr masterManager = ctx.getGlobalStateMgr().getAuthenticationMgr(); + String user = "user123"; + + // 1. create empty image + UtFrameUtils.PseudoJournalReplayer.resetFollowerJournalQueue(); + UtFrameUtils.PseudoImage emptyImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(emptyImage.getDataOutputStream()); + + // 2. create user with properties + String sql = "create user user123 properties (\"session.tx_visible_wait_timeout\" = \"100\", " + + "\"session.metadata_collect_query_timeout\" = \"200\")"; + CreateUserStmt stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.createUser(stmt); + UserProperty userProperty = masterManager.getUserProperty(user); + Assert.assertEquals(2, userProperty.getSessionVariables().size()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("tx_visible_wait_timeout")); + Assert.assertEquals("200", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + + // 2.1. create user with default catalog or database, we expect it will be failed + sql = "create user user2 properties (\"default_session_catalog\" = \"my_catalog\")"; + stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + try { + masterManager.createUser(stmt); + Assert.assertEquals(1, 2); + } catch (Exception e) { + } + + // 3. save final image + UtFrameUtils.PseudoImage finalImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(finalImage.getDataOutputStream()); + + // 4 verify replay... + + // 4.1 load empty image + AuthenticationMgr followerManager = new AuthenticationMgr(); + SRMetaBlockReader srMetaBlockReader = new SRMetaBlockReader(emptyImage.getDataInputStream()); + followerManager.loadV2(srMetaBlockReader); + + // 4.2 replay update user property + CreateUserInfo createUserInfo = (CreateUserInfo) + UtFrameUtils.PseudoJournalReplayer.replayNextJournal(OperationType.OP_CREATE_USER_V2); + followerManager.replayCreateUser( + createUserInfo.getUserIdentity(), + createUserInfo.getAuthenticationInfo(), + createUserInfo.getUserProperty(), + createUserInfo.getUserPrivilegeCollection(), + createUserInfo.getPluginId(), + createUserInfo.getPluginVersion()); + userProperty = followerManager.getUserProperty(user); + Assert.assertEquals(2, userProperty.getSessionVariables().size()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("tx_visible_wait_timeout")); + Assert.assertEquals("200", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + + // 4.3 verify final image + AuthenticationMgr finalManager = new AuthenticationMgr(); + srMetaBlockReader = new SRMetaBlockReader(finalImage.getDataInputStream()); + finalManager.loadV2(srMetaBlockReader); + userProperty = finalManager.getUserProperty(user); + Assert.assertEquals(2, userProperty.getSessionVariables().size()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("tx_visible_wait_timeout")); + Assert.assertEquals("200", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + } + @Test public void testDropAlterUser() throws Exception { UserIdentity testUser = UserIdentity.createAnalyzedUserIdentWithIp("test", "%"); @@ -310,6 +377,97 @@ public void testDropAlterUser() throws Exception { Assert.assertThrows(SemanticException.class, () -> manager.getMaxConn("test")); } + @Test + public void testAlterPersistWithProperties() throws Exception { + AuthenticationMgr masterManager = ctx.getGlobalStateMgr().getAuthenticationMgr(); + CatalogMgr catalogMgr = ctx.getGlobalStateMgr().getCatalogMgr(); + + // 1. create empty image + UtFrameUtils.PseudoJournalReplayer.resetFollowerJournalQueue(); + UtFrameUtils.PseudoImage emptyImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(emptyImage.getDataOutputStream()); + + // create two catalogs + String catalogName = "catalog"; + String createExternalCatalog = "CREATE EXTERNAL CATALOG catalog " + "PROPERTIES( " + " \"type\"=\"hive\", " + + " \"hive.metastore.uris\"=\"thrift://xx.xx.xx.xx:9083\" " + ");"; + CreateCatalogStmt createCatalogStmt = (CreateCatalogStmt) UtFrameUtils.parseStmtWithNewParser(createExternalCatalog, ctx); + catalogMgr.createCatalog(createCatalogStmt); + + String newCatalogName = "new_catalog"; + createExternalCatalog = "CREATE EXTERNAL CATALOG new_catalog " + "PROPERTIES( " + " \"type\"=\"hive\", " + + " \"hive.metastore.uris\"=\"thrift://xx.xx.xx.xx:9083\" " + ");"; + createCatalogStmt = (CreateCatalogStmt) UtFrameUtils.parseStmtWithNewParser(createExternalCatalog, ctx); + catalogMgr.createCatalog(createCatalogStmt); + + // 2. create user with properties + String sql = "create user user1 default role root properties (\"max_user_connections\" = \"100\", " + + "\"session.metadata_collect_query_timeout\" = \"100\", \"session.catalog\" = \"catalog\")"; + CreateUserStmt createUserStmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.createUser(createUserStmt); + UserProperty userProperty = masterManager.getUserProperty("user1"); + Assert.assertEquals(1, userProperty.getSessionVariables().size()); + Assert.assertEquals(100, userProperty.getMaxConn()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + Assert.assertEquals(catalogName, userProperty.getCatalog()); + + // 3. alter user with properties + sql = "alter user user1 set properties (\"max_user_connections\" = \"200\", \"catalog\" = \"new_catalog\")"; + SetUserPropertyStmt setUserPropertyStmt = (SetUserPropertyStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.updateUserProperty(setUserPropertyStmt.getUser(), setUserPropertyStmt.getPropertyPairList()); + Assert.assertEquals(1, userProperty.getSessionVariables().size()); + Assert.assertEquals(200, userProperty.getMaxConn()); + Assert.assertEquals(newCatalogName, userProperty.getCatalog()); + Assert.assertTrue(userProperty.getSessionVariables().get("metadata_collect_query_timeout").equals("100")); + + // 4. save final image + UtFrameUtils.PseudoImage finalImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(finalImage.getDataOutputStream()); + + // 5 verify replay... + + // 5.1 load empty image + AuthenticationMgr followerManager = new AuthenticationMgr(); + SRMetaBlockReader srMetaBlockReader = new SRMetaBlockReader(emptyImage.getDataInputStream()); + followerManager.loadV2(srMetaBlockReader); + + // 5.2 replay create user + CreateUserInfo createUserInfo = (CreateUserInfo) + UtFrameUtils.PseudoJournalReplayer.replayNextJournal(OperationType.OP_CREATE_USER_V2); + followerManager.replayCreateUser( + createUserInfo.getUserIdentity(), + createUserInfo.getAuthenticationInfo(), + createUserInfo.getUserProperty(), + createUserInfo.getUserPrivilegeCollection(), + createUserInfo.getPluginId(), + createUserInfo.getPluginVersion()); + userProperty = followerManager.getUserProperty("user1"); + Assert.assertEquals(1, userProperty.getSessionVariables().size()); + Assert.assertEquals(100, userProperty.getMaxConn()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + Assert.assertEquals(catalogName, userProperty.getCatalog()); + + // 5.2 replay alter user + UserPropertyInfo propertyInfo = + (UserPropertyInfo) UtFrameUtils.PseudoJournalReplayer.replayNextJournal(OperationType.OP_UPDATE_USER_PROP_V3); + followerManager.replayUpdateUserProperty(propertyInfo); + userProperty = followerManager.getUserProperty("user1"); + Assert.assertEquals(1, userProperty.getSessionVariables().size()); + Assert.assertEquals(200, userProperty.getMaxConn()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + Assert.assertEquals(newCatalogName, userProperty.getCatalog()); + + // 4.3 verify final image + AuthenticationMgr finalManager = new AuthenticationMgr(); + srMetaBlockReader = new SRMetaBlockReader(finalImage.getDataInputStream()); + finalManager.loadV2(srMetaBlockReader); + userProperty = finalManager.getUserProperty("user1"); + Assert.assertEquals(1, userProperty.getSessionVariables().size()); + Assert.assertEquals(200, userProperty.getMaxConn()); + Assert.assertEquals("100", userProperty.getSessionVariables().get("metadata_collect_query_timeout")); + Assert.assertEquals(newCatalogName, userProperty.getCatalog()); + } + @Test public void testDropAlterPersist() throws Exception { UserIdentity testUser = UserIdentity.createAnalyzedUserIdentWithIp("test", "%"); @@ -336,7 +494,7 @@ public void testDropAlterPersist() throws Exception { // 3. alter user sql = "alter user test identified by 'abc'"; AlterUserStmt alterUserStmt = (AlterUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); - masterManager.alterUser(alterUserStmt.getUserIdentity(), alterUserStmt.getAuthenticationInfo()); + masterManager.alterUser(alterUserStmt.getUserIdentity(), alterUserStmt.getAuthenticationInfo(), null); Assert.assertEquals(testUser, masterManager.checkPassword( testUser.getUser(), "10.1.1.1", scramble, seed)); @@ -378,7 +536,7 @@ public void testDropAlterPersist() throws Exception { // 7.2 replay alter user AlterUserInfo alterInfo = (AlterUserInfo) UtFrameUtils.PseudoJournalReplayer.replayNextJournal(OperationType.OP_ALTER_USER_V2); - followerManager.replayAlterUser(alterInfo.getUserIdentity(), alterInfo.getAuthenticationInfo()); + followerManager.replayAlterUser(alterInfo.getUserIdentity(), alterInfo.getAuthenticationInfo(), null); Assert.assertEquals(testUser, followerManager.checkPassword( testUser.getUser(), "10.1.1.1", scramble, seed)); // 7.2.1 replay update user property @@ -586,4 +744,45 @@ public void testIsRoleInSession() throws Exception { Assert.assertTrue(e.getMessage().contains("IS_ROLE_IN_SESSION currently only supports a single parameter")); } } + + @Test + public void testSetUserPropertyPersist() throws Exception { + AuthenticationMgr masterManager = ctx.getGlobalStateMgr().getAuthenticationMgr(); + Assert.assertTrue(masterManager.doesUserExist(UserIdentity.ROOT)); + + // 1. create empty image + UtFrameUtils.PseudoJournalReplayer.resetFollowerJournalQueue(); + UtFrameUtils.PseudoImage emptyImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(emptyImage.getDataOutputStream()); + + // 2. update user property + String sql = "set property for 'root' 'max_user_connections' = '555'"; + SetUserPropertyStmt setUserPropertyStmt = (SetUserPropertyStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.updateUserProperty("root", setUserPropertyStmt.getPropertyPairList()); + Assert.assertEquals(555, masterManager.getMaxConn("root")); + + // 3. save final image + UtFrameUtils.PseudoImage finalImage = new UtFrameUtils.PseudoImage(); + masterManager.saveV2(finalImage.getDataOutputStream()); + + // 4 verify replay... + + // 4.1 load empty image + AuthenticationMgr followerManager = new AuthenticationMgr(); + SRMetaBlockReader srMetaBlockReader = new SRMetaBlockReader(emptyImage.getDataInputStream()); + followerManager.loadV2(srMetaBlockReader); + + // 4.2 replay update user property + UserPropertyInfo userPropertyInfo = (UserPropertyInfo) + UtFrameUtils.PseudoJournalReplayer.replayNextJournal(OperationType.OP_UPDATE_USER_PROP_V3); + followerManager.replayUpdateUserProperty(userPropertyInfo); + Assert.assertEquals(555, followerManager.getMaxConn("root")); + + // 4.3 verify final image + AuthenticationMgr finalManager = new AuthenticationMgr(); + srMetaBlockReader = new SRMetaBlockReader(finalImage.getDataInputStream()); + finalManager.loadV2(srMetaBlockReader); + Assert.assertTrue(finalManager.doesUserExist(UserIdentity.ROOT)); + Assert.assertEquals(555, finalManager.getMaxConn("root")); + } } diff --git a/fe/fe-core/src/test/java/com/starrocks/authentication/UserPropertyTest.java b/fe/fe-core/src/test/java/com/starrocks/authentication/UserPropertyTest.java new file mode 100644 index 0000000000000..1fcdfb19a3430 --- /dev/null +++ b/fe/fe-core/src/test/java/com/starrocks/authentication/UserPropertyTest.java @@ -0,0 +1,499 @@ +// Copyright 2021-present StarRocks, Inc. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +package com.starrocks.authentication; + +import com.starrocks.catalog.InternalCatalog; +import com.starrocks.common.Pair; +import com.starrocks.privilege.AuthorizationMgr; +import com.starrocks.privilege.PrivilegeException; +import com.starrocks.qe.ConnectContext; +import com.starrocks.qe.SessionVariable; +import com.starrocks.qe.StmtExecutor; +import com.starrocks.qe.VariableMgr; +import com.starrocks.server.GlobalStateMgr; +import com.starrocks.sql.ast.CreateRoleStmt; +import com.starrocks.sql.ast.CreateUserStmt; +import com.starrocks.sql.ast.DropCatalogStmt; +import com.starrocks.sql.ast.GrantPrivilegeStmt; +import com.starrocks.sql.ast.GrantRoleStmt; +import com.starrocks.sql.ast.UserIdentity; +import com.starrocks.sql.plan.ConnectorPlanTestBase; +import com.starrocks.utframe.StarRocksAssert; +import com.starrocks.utframe.UtFrameUtils; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +public class UserPropertyTest { + + private static String databaseName = "myDB"; + + private static String catalogName = "myCatalog"; + + private static ConnectContext connectContext; + + private static StarRocksAssert starRocksAssert; + + private static AuthorizationMgr authorizationManager; + + @BeforeClass + public static void beforeClass() throws Exception { + UtFrameUtils.createMinStarRocksCluster(); + UtFrameUtils.addMockBackend(10002); + UtFrameUtils.addMockBackend(10003); + + connectContext = UtFrameUtils.initCtxForNewPrivilege(UserIdentity.ROOT); + ConnectorPlanTestBase.mockHiveCatalog(connectContext); + starRocksAssert = new StarRocksAssert(connectContext); + + authorizationManager = starRocksAssert.getCtx().getGlobalStateMgr().getAuthorizationMgr(); + starRocksAssert.getCtx().setRemoteIP("localhost"); + authorizationManager.initBuiltinRolesAndUsers(); + + authorizationManager = starRocksAssert.getCtx().getGlobalStateMgr().getAuthorizationMgr(); + starRocksAssert.getCtx().setRemoteIP("localhost"); + authorizationManager.initBuiltinRolesAndUsers(); + ctxToRoot(); + } + + private static void ctxToRoot() throws PrivilegeException { + starRocksAssert.getCtx().setCurrentUserIdentity(UserIdentity.ROOT); + starRocksAssert.getCtx().setCurrentRoleIds( + starRocksAssert.getCtx().getGlobalStateMgr().getAuthorizationMgr().getRoleIdsByUser(UserIdentity.ROOT)); + + starRocksAssert.getCtx().setQualifiedUser(UserIdentity.ROOT.getUser()); + } + + @Before + public void setUp() throws Exception { + GlobalStateMgr.getCurrentState().clear(); + connectContext = UtFrameUtils.createDefaultCtx(); + starRocksAssert = new StarRocksAssert(connectContext); + if (starRocksAssert.getCtx().getGlobalStateMgr().getCatalogMgr().catalogExists(catalogName)) { + DropCatalogStmt dropCatalogStmt = (DropCatalogStmt) UtFrameUtils.parseStmtWithNewParser( + String.format("DROP CATALOG IF EXISTS %s", catalogName), starRocksAssert.getCtx()); + starRocksAssert.getCtx().getGlobalStateMgr().getCatalogMgr().dropCatalog(dropCatalogStmt); + } + } + + @Test + public void testUpdate_WithMaxConn() throws Exception { + try { + // set max connections too large + List> properties = new ArrayList<>(); + UserProperty userProperty = new UserProperty(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "200000")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + } + + try { + // set max connections too small + List> properties = new ArrayList<>(); + UserProperty userProperty = new UserProperty(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "0")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + } + + try { + // set max connections to a invalid value + List> properties = new ArrayList<>(); + UserProperty userProperty = new UserProperty(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "xx")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + } + + try { + // set max connections to a valid value + List> properties = new ArrayList<>(); + UserProperty userProperty = new UserProperty(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "100")); + userProperty.update("root", properties); + Assert.assertEquals(100, userProperty.getMaxConn()); + } catch (Exception e) { + Assert.assertEquals(1, 2); + } + + try { + // set max connections to a default value + List> properties = new ArrayList<>(); + UserProperty userProperty = new UserProperty(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "")); + userProperty.update("root", properties); + Assert.assertEquals(UserProperty.MAX_CONN_DEFAULT_VALUE, userProperty.getMaxConn()); + } catch (Exception e) { + throw e; + } + } + + @Test + public void testUpdate_WithCatalog() throws Exception { + try { + // set the catalog property + String catalogName = "myCatalog"; + String createExternalCatalog = "CREATE EXTERNAL CATALOG myCatalog " + "PROPERTIES( " + " \"type\"=\"hive\", " + + " \"hive.metastore.uris\"=\"thrift://xx.xx.xx.xx:9083\" " + ");"; + starRocksAssert.withCatalog(createExternalCatalog); + + // set by catalog + List> properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "2000")); + properties.add(new Pair<>(UserProperty.PROP_CATALOG, catalogName)); + UserProperty userProperty = new UserProperty(); + userProperty.update("root", properties); + Assert.assertEquals(2000, userProperty.getMaxConn()); + Assert.assertEquals(catalogName, userProperty.getCatalog()); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + + // set by session.catalog + properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "3000")); + properties.add(new Pair<>("session.catalog", catalogName)); + userProperty = new UserProperty(); + userProperty.update("root", properties); + Assert.assertEquals(3000, userProperty.getMaxConn()); + Assert.assertEquals(catalogName, userProperty.getCatalog()); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + + // reset the catalog property + properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "")); + properties.add(new Pair<>(UserProperty.PROP_CATALOG, "")); + userProperty = new UserProperty(); + userProperty.update("root", properties); + Assert.assertEquals(UserProperty.MAX_CONN_DEFAULT_VALUE, userProperty.getMaxConn()); + Assert.assertEquals(VariableMgr.getDefaultValue(SessionVariable.CATALOG), userProperty.getCatalog()); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } catch (Exception e) { + throw e; + } + } + + @Test + public void testUpdate_WithUser() throws Exception { + try { + String createExternalCatalog = "CREATE EXTERNAL CATALOG myCatalog " + "PROPERTIES( " + " \"type\"=\"hive\", " + + " \"hive.metastore.uris\"=\"thrift://xx.xx.xx.xx:9083\" " + ");"; + starRocksAssert.withCatalog(createExternalCatalog); + + AuthenticationMgr authenticationManager = starRocksAssert.getCtx().getGlobalStateMgr().getAuthenticationMgr(); + String createUserSql = "CREATE USER 'test' IDENTIFIED BY ''"; + CreateUserStmt createUserStmt = + (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(createUserSql, starRocksAssert.getCtx()); + authenticationManager.createUser(createUserStmt); + + UserProperty userProperty = authenticationManager.getUserProperty("test"); + List> properties = UserProperty.changeToPairList( + userProperty.getSessionVariables()); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "2000")); + properties.add(new Pair<>(UserProperty.PROP_CATALOG, catalogName)); + authenticationManager.updateUserProperty("test", properties); + Assert.assertEquals(catalogName, userProperty.getCatalog()); + Assert.assertEquals(2000, userProperty.getMaxConn()); + + // we create a role 'r1' and grant it to user 'test' + AuthorizationMgr authorizationMgr = starRocksAssert.getCtx().getGlobalStateMgr().getAuthorizationMgr(); + String createRoleSql = "CREATE ROLE r1"; + CreateRoleStmt createRoleStmt = + (CreateRoleStmt) UtFrameUtils.parseStmtWithNewParser(createRoleSql, starRocksAssert.getCtx()); + authorizationMgr.createRole(createRoleStmt); + + String grantRoleSql = "GRANT r1 TO USER test"; + GrantRoleStmt grantRoleStmt = + (GrantRoleStmt) UtFrameUtils.parseStmtWithNewParser(grantRoleSql, starRocksAssert.getCtx()); + authorizationMgr.grantRole(grantRoleStmt); + + GrantPrivilegeStmt grantPrivilegeStmt = (GrantPrivilegeStmt) UtFrameUtils.parseStmtWithNewParser( + "grant CREATE DATABASE on CATALOG myCatalog to role r1", + starRocksAssert.getCtx()); + authorizationMgr.grant(grantPrivilegeStmt); + + // Set Default Role + UserIdentity testUser = authenticationManager.getUserIdentityByName("test"); + authorizationMgr.setUserDefaultRole(authorizationMgr.getAllRoleIds(testUser), testUser); + + // EXECUTE AS: the catalog property of root user is default_catalog, the catalog property of test user is myCatalog + Assert.assertEquals(InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME, starRocksAssert.getCtx().getCurrentCatalog()); + new StmtExecutor(starRocksAssert.getCtx(), UtFrameUtils.parseStmtWithNewParser( + String.format("EXECUTE AS test WITH NO REVERT;"), starRocksAssert.getCtx())).execute(); + Assert.assertEquals(catalogName, starRocksAssert.getCtx().getCurrentCatalog()); + } catch (Exception e) { + throw e; + } + } + + @Test + public void testUpdate_WithDatabase() throws Exception { + try { + UserProperty userProperty = new UserProperty(); + List> properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_DATABASE, "xxx")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + } + + try { + // set database for root user + starRocksAssert.withDatabase(databaseName); + UserProperty userProperty = new UserProperty(); + List> properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_DATABASE, databaseName)); + userProperty.update("root", properties); + Assert.assertEquals(databaseName, userProperty.getDatabase()); + + // reset database for root user + userProperty = new UserProperty(); + properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_DATABASE, UserProperty.DATABASE_DEFAULT_VALUE)); + userProperty.update("root", properties); + Assert.assertEquals(UserProperty.DATABASE_DEFAULT_VALUE, userProperty.getDatabase()); + } catch (Exception e) { + throw e; + } + } + + @Test + public void testUpdate_WithSessionVariables() throws Exception { + UserProperty userProperty = new UserProperty(); + try { + // session.aaa is not a valid session variable + List> properties = new ArrayList<>(); + properties.add(new Pair<>("session.aaa", "bbb")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + Assert.assertEquals(1, 1); + } + + userProperty = new UserProperty(); + try { + // the value type of session.wait_timeout is not correct + List> properties = new ArrayList<>(); + properties.add(new Pair<>("session.wait_timeout", "bbb")); + userProperty.update("root", properties); + Assert.assertEquals(1, 2); + } catch (Exception e) { + Assert.assertEquals(1, 1); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } + + userProperty = new UserProperty(); + try { + // init_connect is a global session variable, can't be set. + List> properties = new ArrayList<>(); + properties.add(new Pair<>("session.init_connect", "bbb")); + userProperty.update("root", properties); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } catch (Exception e) { + Assert.assertEquals(1, 1); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } + + userProperty = new UserProperty(); + try { + // system_time_zone is a read-only session variable, can't be set. + List> properties = new ArrayList<>(); + properties.add(new Pair<>("session.system_time_zone", "Asia/Shanghai")); + userProperty.update("root", properties); + } catch (Exception e) { + Assert.assertEquals(1, 1); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } + + userProperty = new UserProperty(); + try { + // session.wait_timeout is a valid session variable + List> properties = new ArrayList<>(); + properties.add(new Pair<>("session.wait_timeout", "1000")); + userProperty.update("root", properties); + Assert.assertEquals("1000", userProperty.getSessionVariables().get("wait_timeout")); + + // reset session.wait_timeout + properties = new ArrayList<>(); + properties.add(new Pair<>("session.wait_timeout", UserProperty.EMPTY_VALUE)); + userProperty.update("root", properties); + Assert.assertEquals(0, userProperty.getSessionVariables().size()); + } catch (Exception e) { + Assert.assertEquals(1, 2); + } + } + + @Test + public void testUpdateForReplayJournal() { + try { + // updateForReplayJournal must not throw any exception + List> properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "2000")); + properties.add(new Pair<>(UserProperty.PROP_DATABASE, "database")); + properties.add(new Pair<>(UserProperty.PROP_CATALOG, "catalog")); + properties.add(new Pair<>("session.aaa", "bbb")); + properties.add(new Pair<>("xxx", "yyy")); + + UserProperty userProperty = new UserProperty(); + userProperty.updateForReplayJournal(properties); + + Assert.assertEquals(2000, userProperty.getMaxConn()); + Assert.assertEquals("database", userProperty.getDatabase()); + Assert.assertEquals("catalog", userProperty.getCatalog()); + Map sessionVariables = userProperty.getSessionVariables(); + Assert.assertEquals(1, sessionVariables.size()); + Assert.assertEquals("bbb", sessionVariables.get("aaa")); + } catch (Exception e) { + throw e; + } + + try { + // set the user connection to a invalid value + List> properties = new ArrayList<>(); + properties.add(new Pair<>(UserProperty.PROP_MAX_USER_CONNECTIONS, "200d")); + properties.add(new Pair<>(UserProperty.PROP_DATABASE, "database")); + properties.add(new Pair<>(UserProperty.PROP_CATALOG, "catalog")); + properties.add(new Pair<>("session.aaa", "bbb")); + properties.add(new Pair<>("xxx", "yyy")); + + UserProperty userProperty = new UserProperty(); + userProperty.updateForReplayJournal(properties); + + Assert.assertEquals(UserProperty.MAX_CONN_DEFAULT_VALUE, userProperty.getMaxConn()); + Assert.assertEquals("database", userProperty.getDatabase()); + Assert.assertEquals("catalog", userProperty.getCatalog()); + Map sessionVariables = userProperty.getSessionVariables(); + Assert.assertEquals(1, sessionVariables.size()); + Assert.assertEquals("bbb", sessionVariables.get("aaa")); + } catch (Exception e) { + throw e; + } + } + + @Test + public void testUpdateSessionContext_WithSomeAbnormalCases() throws Exception { + ConnectContext context = new ConnectContext(null); + UserProperty userProperty = null; + try { + // Update By default UserProperty + userProperty = new UserProperty(); + context.updateByUserProperty(userProperty); + } catch (Exception e) { + throw e; + } + Assert.assertEquals(InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME, context.getCurrentCatalog()); + Assert.assertEquals(UserProperty.DATABASE_DEFAULT_VALUE, context.getDatabase()); + Assert.assertEquals(UserProperty.MAX_CONN_DEFAULT_VALUE, userProperty.getMaxConn()); + + try { + // database does not exist + userProperty = new UserProperty(); + userProperty.setDatabase("database"); + context.updateByUserProperty(userProperty); + } catch (Exception e) { + throw e; + } + Assert.assertEquals("", context.getDatabase()); + + try { + // session variable is not valid + userProperty = new UserProperty(); + Map sessionVariables = userProperty.getSessionVariables(); + sessionVariables.put("aaa", "bbb"); + userProperty.setSessionVariables(sessionVariables); + context.updateByUserProperty(userProperty); + } catch (Exception e) { + throw e; + } + + try { + // session variable is valid + userProperty = new UserProperty(); + Map sessionVariables = userProperty.getSessionVariables(); + sessionVariables.put("statistic_collect_parallel", "2"); + userProperty.setSessionVariables(sessionVariables); + context.updateByUserProperty(userProperty); + } catch (Exception e) { + throw e; + } + Assert.assertEquals(2, context.getSessionVariable().getStatisticCollectParallelism()); + + try { + // catalog is valid + String createExternalCatalog = "CREATE EXTERNAL CATALOG myCatalog " + "PROPERTIES( " + " \"type\"=\"hive\", " + + " \"hive.metastore.uris\"=\"thrift://xx.xx.xx.xx:9083\" " + ");"; + starRocksAssert.withCatalog(createExternalCatalog); + + userProperty = new UserProperty(); + Map sessionVariables = userProperty.getSessionVariables(); + sessionVariables.put("catalog", "myCatalog"); + userProperty.setSessionVariables(sessionVariables); + + context.setCurrentUserIdentity(UserIdentity.ROOT); + context.setCurrentRoleIds( + starRocksAssert.getCtx().getGlobalStateMgr().getAuthorizationMgr().getRoleIdsByUser(UserIdentity.ROOT)); + context.updateByUserProperty(userProperty); + } catch (Exception e) { + throw e; + } + Assert.assertEquals("myCatalog", context.getSessionVariable().getCatalog()); + + } + + @Test + public void testGetCatalogDbName() { + UserProperty userProperty = new UserProperty(); + userProperty.setDatabase("db"); + userProperty.setCatalog("catalog"); + String name = userProperty.getCatalogDbName(); + Assert.assertEquals("catalog.db", name); + + } + + @Test + public void testGetMaxConn() { + UserProperty userProperty = new UserProperty(); + long maxConnections = userProperty.getMaxConn(); + Assert.assertEquals(1024, maxConnections); + } + + @Test + public void testGetDefaultSessionDatabase() { + UserProperty userProperty = new UserProperty(); + String defaultSessionDatabase = userProperty.getDatabase(); + Assert.assertEquals("", defaultSessionDatabase); + } + + @Test + public void testGetDefaultSessionCatalog() { + UserProperty userProperty = new UserProperty(); + String defaultSessionCatalog = userProperty.getCatalog(); + Assert.assertEquals(InternalCatalog.DEFAULT_INTERNAL_CATALOG_NAME, defaultSessionCatalog); + } + + @Test + public void testGetSessionVariables() { + UserProperty userProperty = new UserProperty(); + Map sessionVariables = userProperty.getSessionVariables(); + Assert.assertEquals(0, sessionVariables.size()); + } +} diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/analyzer/AnalyzeStmtTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/analyzer/AnalyzeStmtTest.java index 6c4b5e0b5e6fa..c462783e5ad08 100644 --- a/fe/fe-core/src/test/java/com/starrocks/sql/analyzer/AnalyzeStmtTest.java +++ b/fe/fe-core/src/test/java/com/starrocks/sql/analyzer/AnalyzeStmtTest.java @@ -122,7 +122,7 @@ public void testShowUserProperty() { @Test public void testSetUserProperty() { - String sql = "SET PROPERTY FOR 'tom' 'max_user_connections' = 'value', 'test' = 'true'"; + String sql = "SET PROPERTY FOR 'tom' 'max_user_connections' = '100'"; SetUserPropertyStmt setUserPropertyStmt = (SetUserPropertyStmt) analyzeSuccess(sql); Assert.assertEquals("tom", setUserPropertyStmt.getUser()); } diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/ast/ExecuteAsStmtTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/ast/ExecuteAsStmtTest.java index 0490b205b76aa..23c80c47ddb12 100644 --- a/fe/fe-core/src/test/java/com/starrocks/sql/ast/ExecuteAsStmtTest.java +++ b/fe/fe-core/src/test/java/com/starrocks/sql/ast/ExecuteAsStmtTest.java @@ -16,6 +16,7 @@ package com.starrocks.sql.ast; import com.starrocks.authentication.AuthenticationMgr; +import com.starrocks.authentication.UserProperty; import com.starrocks.privilege.AuthorizationMgr; import com.starrocks.privilege.PrivilegeException; import com.starrocks.qe.ConnectContext; @@ -91,6 +92,17 @@ public void testWithNoRevert() throws Exception { auth.doesUserExist((UserIdentity) any); minTimes = 0; result = true; + + auth.getUserProperty(anyString); + minTimes = 0; + result = new UserProperty(); + } + }; + + new Expectations(ctx) { + { + ctx.updateByUserProperty((UserProperty) any); + minTimes = 0; } };