[BugFix] Fix resource group in audit log and persist planCostRange (S…

…tarRocks#30726) Signed-off-by: zihe.liu <ziheliu1024@gmail.com>
wyb · Sep 12, 2023 · 57161cf · 57161cf
1 parent db10d50
commit 57161cf
Show file tree

Hide file tree

Showing 13 changed files with 320 additions and 158 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroup.java b/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroup.java
@@ -15,23 +15,17 @@
 package com.starrocks.catalog;
 
 import com.google.gson.annotations.SerializedName;
-import com.starrocks.common.io.Text;
-import com.starrocks.common.io.Writable;
-import com.starrocks.persist.gson.GsonUtils;
 import com.starrocks.qe.ShowResultSetMetaData;
 import com.starrocks.thrift.TWorkGroup;
 import com.starrocks.thrift.TWorkGroupType;
 
-import java.io.DataInput;
-import java.io.DataOutput;
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
 
-public class ResourceGroup implements Writable {
+public class ResourceGroup {
     public static final String GROUP_TYPE = "type";
     public static final String USER = "user";
     public static final String ROLE = "role";
@@ -110,11 +104,6 @@ public class ResourceGroup implements Writable {
     public ResourceGroup() {
     }
 
-    public static ResourceGroup read(DataInput in) throws IOException {
-        String json = Text.readString(in);
-        return GsonUtils.GSON.fromJson(json, ResourceGroup.class);
-    }
-
     private List<String> showClassifier(ResourceGroupClassifier classifier) {
         List<String> row = new ArrayList<>();
         row.add(this.name);
@@ -156,11 +145,6 @@ public void setVersion(long version) {
         this.version = version;
     }
 
-    @Override
-    public void write(DataOutput out) throws IOException {
-        Text.writeString(out, GsonUtils.GSON.toJson(this));
-    }
-
     public List<List<String>> show() {
         if (classifiers.isEmpty()) {
             return Collections.singletonList(showClassifier(new ResourceGroupClassifier()));

diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroupClassifier.java b/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroupClassifier.java
@@ -16,21 +16,11 @@
 
 import com.google.common.collect.ImmutableSet;
 import com.google.gson.annotations.SerializedName;
-import com.starrocks.common.io.Text;
-import com.starrocks.common.io.Writable;
-import com.starrocks.persist.gson.GsonUtils;
 import com.starrocks.server.GlobalStateMgr;
-import com.starrocks.sql.ast.DmlStmt;
-import com.starrocks.sql.ast.LoadStmt;
-import com.starrocks.sql.ast.QueryStatement;
-import com.starrocks.sql.ast.StatementBase;
 import com.starrocks.thrift.TQueryType;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.net.util.SubnetUtils;
 
-import java.io.DataInput;
-import java.io.DataOutput;
-import java.io.IOException;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
@@ -39,7 +29,7 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-public class ResourceGroupClassifier implements Writable {
+public class ResourceGroupClassifier {
     public static final Pattern USER_PATTERN = Pattern.compile("^[a-zA-Z][a-zA-Z0-9_]{1,63}/?[.a-zA-Z0-9_-]{0,63}$");
     public static final Pattern USE_ROLE_PATTERN = Pattern.compile("^\\w+$");
     public static final ImmutableSet<String> SUPPORTED_QUERY_TYPES =
@@ -66,11 +56,6 @@ public class ResourceGroupClassifier implements Writable {
     @SerializedName(value = "planMemCostRange")
     private CostRange planMemCostRange;
 
-    public static ResourceGroupClassifier read(DataInput in) throws IOException {
-        String json = Text.readString(in);
-        return GsonUtils.GSON.fromJson(json, ResourceGroupClassifier.class);
-    }
-
     public long getResourceGroupId() {
         return resourceGroupId;
     }
@@ -131,14 +116,16 @@ public void setPlanCpuCostRange(CostRange planCpuCostRange) {
         this.planCpuCostRange = planCpuCostRange;
     }
 
+    public CostRange getPlanCpuCostRange() {
+        return planCpuCostRange;
+    }
+
     public void setPlanMemCostRange(CostRange planMemCostRange) {
         this.planMemCostRange = planMemCostRange;
     }
 
-    @Override
-    public void write(DataOutput out) throws IOException {
-        String json = GsonUtils.GSON.toJson(this);
-        Text.writeString(out, json);
+    public CostRange getPlanMemCostRange() {
+        return planMemCostRange;
     }
 
     public boolean isSatisfied(String user, List<String> activeRoles, QueryType queryType, String sourceIp,
@@ -259,20 +246,6 @@ public enum QueryType {
         public static QueryType fromTQueryType(TQueryType type) {
             return type == TQueryType.LOAD ? INSERT : SELECT;
         }
-
-        public static QueryType fromStatement(StatementBase stmt) {
-            if (stmt instanceof QueryStatement) {
-                return SELECT;
-            }
-            if (stmt instanceof DmlStmt) {
-                return INSERT;
-            }
-            if (stmt instanceof LoadStmt) {
-                return INSERT;
-            }
-
-            return SELECT;
-        }
     }
 
     /**
@@ -286,13 +259,15 @@ public static class CostRange {
         private static final Pattern STR_RANGE_PATTERN = Pattern.compile(STR_RANGE_REGEX, Pattern.CASE_INSENSITIVE);
 
         public static final String FORMAT_STR_RANGE_MESSAGE = "the format must be '[min, max)' " +
-                "where min and max are double (including infinity and -infinity) " +
+                "where min and max are finite double " +
                 "and min must be less than max";
 
+        @SerializedName(value = "min")
         private final double min;
+        @SerializedName(value = "max")
         private final double max;
 
-        public CostRange(double min, double max) {
+        private CostRange(double min, double max) {
             this.min = min;
             this.max = max;
         }
@@ -304,8 +279,12 @@ public static CostRange fromString(String rangeStr) {
             }
 
             try {
-                double min = parseDoubleWithInfinity(matcher.group(1));
-                double max = parseDoubleWithInfinity(matcher.group(2));
+                double min = Double.parseDouble(matcher.group(1));
+                double max = Double.parseDouble(matcher.group(2));
+
+                if (!Double.isFinite(min) || !Double.isFinite(max)) {
+                    return null;
+                }
 
                 if (min >= max) {
                     return null;
@@ -325,15 +304,5 @@ public boolean contains(double value) {
         public String toString() {
             return "[" + min + ", " + max + ")";
         }
-
-        private static double parseDoubleWithInfinity(String input) {
-            if ("infinity".equalsIgnoreCase(input)) {
-                return Double.POSITIVE_INFINITY;
-            } else if ("-infinity".equalsIgnoreCase(input)) {
-                return Double.NEGATIVE_INFINITY;
-            } else {
-                return Double.parseDouble(input);
-            }
-        }
     }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroupMgr.java b/fe/fe-core/src/main/java/com/starrocks/catalog/ResourceGroupMgr.java
@@ -246,7 +246,7 @@ public Set<String> getAllResourceGroupNames() {
 
     @Override
     public void write(DataOutput out) throws IOException {
-        List<ResourceGroup> resourceGroups = resourceGroupMap.values().stream().collect(Collectors.toList());
+        List<ResourceGroup> resourceGroups = new ArrayList<>(resourceGroupMap.values());
         SerializeData data = new SerializeData();
         data.resourceGroups = resourceGroups;
 

diff --git a/fe/fe-core/src/main/java/com/starrocks/persist/ResourceGroupOpEntry.java b/fe/fe-core/src/main/java/com/starrocks/persist/ResourceGroupOpEntry.java
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
 package com.starrocks.persist;
 
 import com.google.gson.annotations.SerializedName;

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ConnectProcessor.java b/fe/fe-core/src/main/java/com/starrocks/qe/ConnectProcessor.java
@@ -74,6 +74,7 @@
 import com.starrocks.thrift.TQueryOptions;
 import com.starrocks.thrift.TWorkGroup;
 import org.apache.commons.codec.binary.Hex;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -716,6 +717,11 @@ public TMasterOpResult proxyExecute(TMasterOpRequest request) {
             } else if (executor.getProxyResultBuffer() != null) {  // query statement
                 result.setChannelBufferList(executor.getProxyResultBuffer());
             }
+
+            String resourceGroupName = ctx.getAuditEventBuilder().build().resourceGroup;
+            if (StringUtils.isNotEmpty(resourceGroupName)) {
+                result.setResource_group_name(resourceGroupName);
+            }
         }
         return result;
     }

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/LeaderOpExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/LeaderOpExecutor.java
@@ -110,6 +110,10 @@ public void execute() throws Exception {
                 }
             }
         }
+
+        if (result.isSetResource_group_name()) {
+            ctx.getAuditEventBuilder().setResourceGroup(result.getResource_group_name());
+        }
     }
 
     private void afterForward() throws DdlException {

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java
@@ -200,7 +200,6 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
-import static com.starrocks.qe.CoordinatorPreprocessor.prepareResourceGroup;
 import static com.starrocks.sql.common.UnsupportedException.unsupportedException;
 
 // Do one COM_QUERY process.
@@ -477,9 +476,6 @@ public void execute() throws Exception {
                 return;
             }
             if (isForwardToLeader()) {
-                // Write the resource group information to audit log.
-                prepareResourceGroup(context, ResourceGroupClassifier.QueryType.fromStatement(parsedStmt));
-
                 forwardToLeader();
                 return;
             } else {

diff --git a/fe/fe-core/src/main/java/com/starrocks/service/FrontendServiceImpl.java b/fe/fe-core/src/main/java/com/starrocks/service/FrontendServiceImpl.java
@@ -1223,7 +1223,8 @@ public TMasterOpResult forward(TMasterOpRequest params) throws TException {
         }
 
         // add this log so that we can track this stmt
-        LOG.info("receive forwarded stmt {} from FE: {}", params.getStmt_id(), clientAddr.getHostname());
+        LOG.info("receive forwarded stmt {} from FE: {}",
+                params.getStmt_id(), clientAddr != null ? clientAddr.getHostname() : "unknown");
         ConnectContext context = new ConnectContext(null);
         ConnectProcessor processor = new ConnectProcessor(context);
         TMasterOpResult result = processor.proxyExecute(params);
@@ -1963,7 +1964,7 @@ public TImmutablePartitionResult updateImmutablePartition(TImmutablePartitionReq
         return result;
     }
 
-    public synchronized TImmutablePartitionResult updateImmutablePartitionInternal(TImmutablePartitionRequest request) 
+    public synchronized TImmutablePartitionResult updateImmutablePartitionInternal(TImmutablePartitionRequest request)
             throws UserException {
         long dbId = request.getDb_id();
         long tableId = request.getTable_id();
@@ -2100,7 +2101,7 @@ private static void buildTablets(PhysicalPartition physicalPartition, List<TTabl
                     if (bePathsMap.keySet().size() < quorum) {
                         throw new UserException(
                                 "Tablet lost replicas. Check if any backend is down or not. tablet_id: "
-                                + tablet.getId() + ", backends: " + Joiner.on(",").join(localTablet.getBackends()));
+                                        + tablet.getId() + ", backends: " + Joiner.on(",").join(localTablet.getBackends()));
                     }
                     // replicas[0] will be the primary replica
                     // getNormalReplicaBackendPathMap returns a linkedHashMap, it's keysets is stable
@@ -2614,7 +2615,6 @@ public TReleaseSlotResponse releaseSlot(TReleaseSlotRequest request) throws TExc
         return res;
     }
 
-
     @Override
     public TGetDictQueryParamResponse getDictQueryParam(TGetDictQueryParamRequest request) throws TException {
         Database db = GlobalStateMgr.getCurrentState().getDb(request.getDb_name());
@@ -2645,7 +2645,7 @@ public TGetDictQueryParamResponse getDictQueryParam(TGetDictQueryParamRequest re
             List<Long> allPartitions = dictTable.getAllPartitionIds();
             response.setPartition(
                     OlapTableSink.createPartition(
-                        db.getId(), dictTable, dictTable.supportedAutomaticPartition(), allPartitions));
+                            db.getId(), dictTable, dictTable.supportedAutomaticPartition(), allPartitions));
             response.setLocation(OlapTableSink.createLocation(
                     dictTable, dictTable.getClusterId(), allPartitions, dictTable.enableReplicatedStorage()));
             response.setNodes_info(GlobalStateMgr.getCurrentState().createNodesInfo(dictTable.getClusterId()));

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ResourceGroupAnalyzer.java b/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ResourceGroupAnalyzer.java
@@ -142,8 +142,11 @@ public static ResourceGroupClassifier convertPredicateToClassifier(List<Predicat
                 classifier.getRole() == null &&
                 (classifier.getQueryTypes() == null || classifier.getQueryTypes().isEmpty()) &&
                 classifier.getSourceIp() == null &&
-                classifier.getDatabases() == null) {
-            throw new SemanticException("At least one of ('user', 'role', 'query_type', 'source_ip') should be given");
+                classifier.getDatabases() == null &&
+                classifier.getPlanCpuCostRange() == null &&
+                classifier.getPlanMemCostRange() == null) {
+            throw new SemanticException("At least one of ('user', 'role', 'query_type', 'db', 'source_ip', " +
+                    "'plan_cpu_cost_range', 'plan_mem_cost_range') should be given");
         }
         return classifier;
     }

diff --git a/fe/fe-core/src/test/java/com/starrocks/analysis/ResourceGroupMgrTest.java b/fe/fe-core/src/test/java/com/starrocks/analysis/ResourceGroupMgrTest.java
-Original file line number
+Diff line change
@@ Expand Up / @@ -110,6 +110,10 @@ public void execute() throws Exception { @@
                     }
                 }
             }
+            if (result.isSetResource_group_name()) {
+                ctx.getAuditEventBuilder().setResourceGroup(result.getResource_group_name());
+            }
         }
         private void afterForward() throws DdlException {
@@ Expand Down @@