[shardingsphere] branch master updated: add `count mask rule` DistSQL resultSet (#23207)

[ji...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jianglongtao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git


The following commit(s) were added to refs/heads/master by this push:
     new e82d9235c9c add `count mask rule` DistSQL resultSet (#23207)
e82d9235c9c is described below

commit e82d9235c9c933d16c00e04e082140d1634e63c2
Author: Zichao <57...@users.noreply.github.com>
AuthorDate: Sat Dec 31 02:05:35 2022 +1300

    add `count mask rule` DistSQL resultSet (#23207)
    
    * add `count mask rule` DistSQL resultSet
    
    * add `count mask rule` DistSQL resultSet
---
 .../handler/query/CountMaskRuleResultSet.java      | 91 ++++++++++++++++++++++
 ...here.distsql.handler.resultset.DistSQLResultSet |  2 +-
 .../parser/core/MaskDistSQLStatementVisitor.java   | 12 ++-
 .../parser/statement/CountMaskRuleStatement.java   | 31 ++++++++
 4 files changed, 133 insertions(+), 3 deletions(-)

diff --git a/features/mask/distsql/handler/src/main/java/org/apache/shardingsphere/mask/distsql/handler/query/CountMaskRuleResultSet.java b/features/mask/distsql/handler/src/main/java/org/apache/shardingsphere/mask/distsql/handler/query/CountMaskRuleResultSet.java
new file mode 100644
index 00000000000..8322027fe78
--- /dev/null
+++ b/features/mask/distsql/handler/src/main/java/org/apache/shardingsphere/mask/distsql/handler/query/CountMaskRuleResultSet.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ *     http://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 org.apache.shardingsphere.mask.distsql.handler.query;
+
+import org.apache.shardingsphere.distsql.handler.resultset.DatabaseDistSQLResultSet;
+import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.mask.distsql.parser.statement.CountMaskRuleStatement;
+import org.apache.shardingsphere.mask.rule.MaskRule;
+import org.apache.shardingsphere.sql.parser.sql.common.statement.SQLStatement;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+/**
+ * Result set for count mask rule.
+ */
+public final class CountMaskRuleResultSet implements DatabaseDistSQLResultSet {
+    
+    private static final String MASK = "mask";
+    
+    private Iterator<Entry<String, LinkedList<Object>>> data = Collections.emptyIterator();
+    
+    @Override
+    public Collection<String> getColumnNames() {
+        return Arrays.asList("rule_name", "database", "count");
+    }
+    
+    @Override
+    public void init(final ShardingSphereDatabase database, final SQLStatement sqlStatement) {
+        Optional<MaskRule> rule = database.getRuleMetaData().findSingleRule(MaskRule.class);
+        Map<String, LinkedList<Object>> result = new LinkedHashMap<>();
+        rule.ifPresent(optional -> addMaskData(result, database.getName(), rule.get()));
+        data = result.entrySet().iterator();
+    }
+    
+    private void addMaskData(final Map<String, LinkedList<Object>> rowMap, final String databaseName, final MaskRule rule) {
+        addData(rowMap, MASK, databaseName, () -> rule.getTables().size());
+    }
+    
+    private void addData(final Map<String, LinkedList<Object>> rowMap, final String dataKey, final String databaseName, final Supplier<Integer> apply) {
+        rowMap.compute(dataKey, (key, value) -> buildRow(value, databaseName, apply.get()));
+    }
+    
+    private LinkedList<Object> buildRow(final LinkedList<Object> value, final String databaseName, final int count) {
+        if (null == value) {
+            return new LinkedList<>(Arrays.asList(databaseName, count));
+        }
+        value.set(1, (Integer) value.get(1) + count);
+        return value;
+    }
+    
+    @Override
+    public boolean next() {
+        return data.hasNext();
+    }
+    
+    @Override
+    public Collection<Object> getRowData() {
+        Entry<String, LinkedList<Object>> entry = data.next();
+        entry.getValue().addFirst(entry.getKey());
+        return entry.getValue();
+    }
+    
+    @Override
+    public String getType() {
+        return CountMaskRuleStatement.class.getName();
+    }
+}
diff --git a/features/mask/distsql/handler/src/main/resources/META-INF/services/org.apache.shardingsphere.distsql.handler.resultset.DistSQLResultSet b/features/mask/distsql/handler/src/main/resources/META-INF/services/org.apache.shardingsphere.distsql.handler.resultset.DistSQLResultSet
index 8bb10a389bb..847d8dbb05d 100644
--- a/features/mask/distsql/handler/src/main/resources/META-INF/services/org.apache.shardingsphere.distsql.handler.resultset.DistSQLResultSet
+++ b/features/mask/distsql/handler/src/main/resources/META-INF/services/org.apache.shardingsphere.distsql.handler.resultset.DistSQLResultSet
@@ -16,4 +16,4 @@
 #
 
 org.apache.shardingsphere.mask.distsql.handler.query.MaskRuleResultSet
-#org.apache.shardingsphere.mask.distsql.handler.query.CountMaskRuleResultSet
+org.apache.shardingsphere.mask.distsql.handler.query.CountMaskRuleResultSet
diff --git a/features/mask/distsql/parser/src/main/java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java b/features/mask/distsql/parser/src/main/java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java
index db6a776574e..f972b450747 100644
--- a/features/mask/distsql/parser/src/main/java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java
+++ b/features/mask/distsql/parser/src/main/java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java
@@ -19,7 +19,6 @@ package org.apache.shardingsphere.mask.distsql.parser.core;
 
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementBaseVisitor;
-import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.AlgorithmDefinitionContext;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.AlterMaskRuleContext;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.ColumnDefinitionContext;
@@ -28,10 +27,13 @@ import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementPars
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.MaskRuleDefinitionContext;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.PropertiesDefinitionContext;
 import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.PropertyContext;
+import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.ShowMaskRulesContext;
+import org.apache.shardingsphere.distsql.parser.autogen.MaskDistSQLStatementParser.CountMaskRuleContext;
 import org.apache.shardingsphere.distsql.parser.segment.AlgorithmSegment;
 import org.apache.shardingsphere.mask.distsql.parser.segment.MaskColumnSegment;
 import org.apache.shardingsphere.mask.distsql.parser.segment.MaskRuleSegment;
 import org.apache.shardingsphere.mask.distsql.parser.statement.AlterMaskRuleStatement;
+import org.apache.shardingsphere.mask.distsql.parser.statement.CountMaskRuleStatement;
 import org.apache.shardingsphere.mask.distsql.parser.statement.CreateMaskRuleStatement;
 import org.apache.shardingsphere.mask.distsql.parser.statement.DropMaskRuleStatement;
 import org.apache.shardingsphere.mask.distsql.parser.statement.ShowMaskRulesStatement;
@@ -40,6 +42,7 @@ import org.apache.shardingsphere.sql.parser.api.visitor.SQLVisitor;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.DatabaseSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.value.identifier.IdentifierValue;
 
+import java.util.Objects;
 import java.util.Properties;
 import java.util.stream.Collectors;
 
@@ -64,11 +67,16 @@ public final class MaskDistSQLStatementVisitor extends MaskDistSQLStatementBaseV
     }
     
     @Override
-    public ASTNode visitShowMaskRules(final MaskDistSQLStatementParser.ShowMaskRulesContext ctx) {
+    public ASTNode visitShowMaskRules(final ShowMaskRulesContext ctx) {
         return new ShowMaskRulesStatement(null == ctx.RULE() ? null : getIdentifierValue(ctx.ruleName()),
                 null == ctx.databaseName() ? null : (DatabaseSegment) visit(ctx.databaseName()));
     }
     
+    @Override
+    public ASTNode visitCountMaskRule(final CountMaskRuleContext ctx) {
+        return new CountMaskRuleStatement(Objects.nonNull(ctx.databaseName()) ? (DatabaseSegment) visit(ctx.databaseName()) : null);
+    }
+    
     @Override
     public ASTNode visitMaskRuleDefinition(final MaskRuleDefinitionContext ctx) {
         return new MaskRuleSegment(getIdentifierValue(ctx.ruleName()), ctx.columnDefinition().stream().map(each -> (MaskColumnSegment) visit(each)).collect(Collectors.toList()));
diff --git a/features/mask/distsql/statement/src/main/java/org/apache/shardingsphere/mask/distsql/parser/statement/CountMaskRuleStatement.java b/features/mask/distsql/statement/src/main/java/org/apache/shardingsphere/mask/distsql/parser/statement/CountMaskRuleStatement.java
new file mode 100644
index 00000000000..960b54783fb
--- /dev/null
+++ b/features/mask/distsql/statement/src/main/java/org/apache/shardingsphere/mask/distsql/parser/statement/CountMaskRuleStatement.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ *     http://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 org.apache.shardingsphere.mask.distsql.parser.statement;
+
+import org.apache.shardingsphere.distsql.parser.statement.rql.show.ShowRulesStatement;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.DatabaseSegment;
+
+/**
+ * Count mask rule statement.
+ */
+public final class CountMaskRuleStatement extends ShowRulesStatement {
+    
+    public CountMaskRuleStatement(final DatabaseSegment database) {
+        super(database);
+    }
+}