mysql数据库禁写插件测试问题修复

Signed-off-by: daizhenyu <[email protected]>

sermant-plugins/sermant-database-write-prohibition/mysql-mariadb-3.x-plugin/src/main/java/com/huaweicloud/sermant/mariadbv3/declarers/StandardClientDeclarer.java

@@ -35,8 +35,9 @@ public ClassMatcher getClassMatcher() {
 
     @Override
     public InterceptDeclarer[] getInterceptDeclarers(ClassLoader classLoader) {
-        return new InterceptDeclarer[] {
-                MariadbV3EnhancementHelper.getSendQueryInterceptDeclarer()
+        return new InterceptDeclarer[]{
+                MariadbV3EnhancementHelper.getSendQueryInterceptDeclarer(),
+                MariadbV3EnhancementHelper.getExecutePipelineInterceptDeclarer()
         };
     }
 }

sermant-plugins/sermant-database-write-prohibition/mysql-mariadb-3.x-plugin/src/main/java/com/huaweicloud/sermant/mariadbv3/interceptors/AbstractMariadbV3Interceptor.java

@@ -0,0 +1,46 @@
+/*
+ *  Copyright (C) 2024-2024 Huawei Technologies Co., Ltd. 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
+ *
+ *       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 com.huaweicloud.sermant.mariadbv3.interceptors;
+
+import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
+import com.huaweicloud.sermant.database.constant.DatabaseType;
+import com.huaweicloud.sermant.database.entity.DatabaseInfo;
+import com.huaweicloud.sermant.database.interceptor.AbstractDatabaseInterceptor;
+
+import org.mariadb.jdbc.HostAddress;
+import org.mariadb.jdbc.client.impl.StandardClient;
+
+/**
+ * sendQuery方法拦截器
+ *
+ * @author daizhenyu
+ * @since 2024-01-30
+ **/
+public abstract class AbstractMariadbV3Interceptor extends AbstractDatabaseInterceptor {
+    @Override
+    protected void createAndCacheDatabaseInfo(ExecuteContext context) {
+        DatabaseInfo info = new DatabaseInfo(DatabaseType.MYSQL);
+        context.setLocalFieldValue(DATABASE_INFO, info);
+        StandardClient client = (StandardClient) context.getObject();
+        if (client.getContext() != null) {
+            info.setDatabaseName(client.getContext().getDatabase());
+        }
+        HostAddress hostAddress = client.getHostAddress();
+        info.setHostAddress(hostAddress.host);
+        info.setPort(hostAddress.port);
+    }
+}

sermant-plugins/sermant-database-write-prohibition/mysql-mariadb-3.x-plugin/src/main/java/com/huaweicloud/sermant/mariadbv3/interceptors/ExecutePipelineInterceptor.java

@@ -0,0 +1,64 @@
+/*
+ *  Copyright (C) 2024-2024 Huawei Technologies Co., Ltd. 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
+ *
+ *       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 com.huaweicloud.sermant.mariadbv3.interceptors;
+
+import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
+import com.huaweicloud.sermant.database.config.DatabaseWriteProhibitionManager;
+import com.huaweicloud.sermant.database.controller.DatabaseController;
+import com.huaweicloud.sermant.database.handler.DatabaseHandler;
+import com.huaweicloud.sermant.database.utils.SqlParserUtils;
+
+import org.mariadb.jdbc.message.ClientMessage;
+
+/**
+ * executePipeline方法拦截器
+ *
+ * @author daizhenyu
+ * @since 2024-01-30
+ **/
+public class ExecutePipelineInterceptor extends AbstractMariadbV3Interceptor {
+    /**
+     * 无参构造方法
+     */
+    public ExecutePipelineInterceptor() {
+    }
+
+    /**
+     * 有参构造方法
+     *
+     * @param handler 写操作处理器
+     */
+    public ExecutePipelineInterceptor(DatabaseHandler handler) {
+        this.handler = handler;
+    }
+
+    @Override
+    protected ExecuteContext doBefore(ExecuteContext context) {
+        String database = getDataBaseInfo(context).getDatabaseName();
+        if (!DatabaseWriteProhibitionManager.getMySqlProhibitionDatabases().contains(database)) {
+            return context;
+        }
+        ClientMessage[] clientMessages = (ClientMessage[]) context.getArguments()[0];
+        for (ClientMessage clientMessage : clientMessages) {
+            if (SqlParserUtils.isWriteOperation(clientMessage.description())) {
+                DatabaseController.disableDatabaseWriteOperation(database, context);
+                return context;
+            }
+        }
+        return context;
+    }
+}

sermant-plugins/sermant-database-write-prohibition/mysql-mariadb-3.x-plugin/src/main/java/com/huaweicloud/sermant/mariadbv3/interceptors/SendQueryInterceptor.java

@@ -18,15 +18,10 @@
 
 import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
 import com.huaweicloud.sermant.database.config.DatabaseWriteProhibitionManager;
-import com.huaweicloud.sermant.database.constant.DatabaseType;
 import com.huaweicloud.sermant.database.controller.DatabaseController;
-import com.huaweicloud.sermant.database.entity.DatabaseInfo;
 import com.huaweicloud.sermant.database.handler.DatabaseHandler;
-import com.huaweicloud.sermant.database.interceptor.AbstractDatabaseInterceptor;
 import com.huaweicloud.sermant.database.utils.SqlParserUtils;
 
-import org.mariadb.jdbc.HostAddress;
-import org.mariadb.jdbc.client.impl.StandardClient;
 import org.mariadb.jdbc.message.ClientMessage;
 
 /**
@@ -35,7 +30,7 @@
  * @author daizhenyu
  * @since 2024-01-30
  **/
-public class SendQueryInterceptor extends AbstractDatabaseInterceptor {
+public class SendQueryInterceptor extends AbstractMariadbV3Interceptor {
     /**
      * 无参构造方法
      */
@@ -61,17 +56,4 @@ protected ExecuteContext doBefore(ExecuteContext context) {
         }
         return context;
     }
-
-    @Override
-    protected void createAndCacheDatabaseInfo(ExecuteContext context) {
-        DatabaseInfo info = new DatabaseInfo(DatabaseType.MYSQL);
-        context.setLocalFieldValue(DATABASE_INFO, info);
-        StandardClient client = (StandardClient) context.getObject();
-        if (client.getContext() != null) {
-            info.setDatabaseName(client.getContext().getDatabase());
-        }
-        HostAddress hostAddress = client.getHostAddress();
-        info.setHostAddress(hostAddress.host);
-        info.setPort(hostAddress.port);
-    }
 }

sermant-plugins/sermant-database-write-prohibition/mysql-mariadb-3.x-plugin/src/main/java/com/huaweicloud/sermant/mariadbv3/utils/MariadbV3EnhancementHelper.java

@@ -20,6 +20,7 @@
 import com.huaweicloud.sermant.core.plugin.agent.matcher.ClassMatcher;
 import com.huaweicloud.sermant.core.plugin.agent.matcher.MethodMatcher;
 import com.huaweicloud.sermant.database.handler.DatabaseHandler;
+import com.huaweicloud.sermant.mariadbv3.interceptors.ExecutePipelineInterceptor;
 import com.huaweicloud.sermant.mariadbv3.interceptors.SendQueryInterceptor;
 
 /**
@@ -35,6 +36,8 @@ public class MariadbV3EnhancementHelper {
 
     private static final String SEND_QUERY_METHOD_NAME = "sendQuery";
 
+    private static final String EXECUTE_PIPELINE_METHOD_NAME = "executePipeline";
+
     private MariadbV3EnhancementHelper() {
     }
 
@@ -77,7 +80,32 @@ public static InterceptDeclarer getSendQueryInterceptDeclarer(DatabaseHandler ha
                 new SendQueryInterceptor(handler));
     }
 
+    /**
+     * 获取executePipeline方法无参拦截器
+     *
+     * @return InterceptDeclarer executePipeline方法无参拦截器
+     */
+    public static InterceptDeclarer getExecutePipelineInterceptDeclarer() {
+        return InterceptDeclarer.build(getExecutePipelineMethodMatcher(),
+                new ExecutePipelineInterceptor());
+    }
+
+    /**
+     * 获取executePipeline方法有参拦截器
+     *
+     * @param handler 数据库自定义处理器
+     * @return InterceptDeclarer executePipeline方法有参拦截器
+     */
+    public static InterceptDeclarer getExecutePipelineInterceptDeclarer(DatabaseHandler handler) {
+        return InterceptDeclarer.build(getExecutePipelineMethodMatcher(),
+                new ExecutePipelineInterceptor(handler));
+    }
+
     private static MethodMatcher getSendQueryMethodMatcher() {
         return MethodMatcher.nameEquals(SEND_QUERY_METHOD_NAME);
     }
+
+    private static MethodMatcher getExecutePipelineMethodMatcher() {
+        return MethodMatcher.nameEquals(EXECUTE_PIPELINE_METHOD_NAME);
+    }
 }