Fix the issue of database write prohibition failure in some versions …

…of Postgresql

Signed-off-by: hanbingleixue <hanbingleixue@hotmail.com>

sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/java/com/huaweicloud/sermant/postgresqlv42/declarers/PgStatementDeclarer.java

@@ -0,0 +1,41 @@
+/*
+ *  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.postgresqlv42.declarers;
+
+import com.huaweicloud.sermant.core.plugin.agent.declarer.AbstractPluginDeclarer;
+import com.huaweicloud.sermant.core.plugin.agent.declarer.InterceptDeclarer;
+import com.huaweicloud.sermant.core.plugin.agent.matcher.ClassMatcher;
+import com.huaweicloud.sermant.postgresqlv42.utils.PostgresqlEnhancementHelper;
+
+/**
+ * PgStatement declarer
+ *
+ * @author zhp
+ * @since 2024-03-19
+ **/
+public class PgStatementDeclarer extends AbstractPluginDeclarer {
+    @Override
+    public ClassMatcher getClassMatcher() {
+        return PostgresqlEnhancementHelper.getPgStatementClassMatcher();
+    }
+
+    @Override
+    public InterceptDeclarer[] getInterceptDeclarers(ClassLoader classLoader) {
+        return new InterceptDeclarer[]{PostgresqlEnhancementHelper.getPgStatementExecuteInterceptDeclarer(),
+                PostgresqlEnhancementHelper.getPgStatementExecuteBatchInterceptDeclarer()};
+    }
+}

sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/java/com/huaweicloud/sermant/postgresqlv42/declarers/QueryExecutorImplDeclarer.java

@@ -19,7 +19,7 @@
 import com.huaweicloud.sermant.core.plugin.agent.declarer.AbstractPluginDeclarer;
 import com.huaweicloud.sermant.core.plugin.agent.declarer.InterceptDeclarer;
 import com.huaweicloud.sermant.core.plugin.agent.matcher.ClassMatcher;
-import com.huaweicloud.sermant.postgresqlv42.utils.QueryExecutorImplEnhancementHelper;
+import com.huaweicloud.sermant.postgresqlv42.utils.PostgresqlEnhancementHelper;
 
 /**
  * QueryExecutorImpl declarer
@@ -30,11 +30,11 @@
 public class QueryExecutorImplDeclarer extends AbstractPluginDeclarer {
     @Override
     public ClassMatcher getClassMatcher() {
-        return QueryExecutorImplEnhancementHelper.getQueryExecutorImplClassMatcher();
+        return PostgresqlEnhancementHelper.getQueryExecutorImplClassMatcher();
     }
 
     @Override
     public InterceptDeclarer[] getInterceptDeclarers(ClassLoader classLoader) {
-        return new InterceptDeclarer[]{QueryExecutorImplEnhancementHelper.getSendOneQueryInterceptDeclarer()};
+        return new InterceptDeclarer[]{PostgresqlEnhancementHelper.getSendOneQueryInterceptDeclarer()};
     }
 }

sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/java/com/huaweicloud/sermant/postgresqlv42/interceptors/PgStatementInterceptor.java

@@ -0,0 +1,87 @@
+/*
+ *  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.postgresqlv42.interceptors;
+
+import com.huaweicloud.sermant.core.common.LoggerFactory;
+import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
+import com.huaweicloud.sermant.database.handler.DatabaseHandler;
+import com.huaweicloud.sermant.database.interceptor.AbstractDatabaseInterceptor;
+import com.huaweicloud.sermant.database.utils.ThreadDatabaseUrlUtil;
+
+import org.postgresql.jdbc.PgStatement;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * Interceptor for PgStatement execute and executeBatch method
+ *
+ * @author zhp
+ * @since 2024-03-19
+ **/
+public class PgStatementInterceptor extends AbstractDatabaseInterceptor {
+    private static final Logger LOGGER = LoggerFactory.getLogger();
+
+    /**
+     * Non-parametric construction method
+     */
+    public PgStatementInterceptor() {
+    }
+
+    /**
+     * Parameterized construction method
+     *
+     * @param handler Database write operation handler
+     */
+    public PgStatementInterceptor(DatabaseHandler handler) {
+        this.handler = handler;
+    }
+
+    @Override
+    public ExecuteContext doBefore(ExecuteContext context) {
+        PgStatement statement = (PgStatement) context.getObject();
+
+        // Store link information to ensure that QueryExecutorImplInterceptor can retrieve database information
+        try {
+            Connection connection = statement.getConnection();
+            if (connection != null && connection.getMetaData() != null) {
+                ThreadDatabaseUrlUtil.setDatabaseUrl(connection.getMetaData().getURL());
+            }
+        } catch (SQLException e) {
+            LOGGER.log(Level.SEVERE, "can not obtain the database information.", e);
+        }
+        return context;
+    }
+
+    @Override
+    protected void createAndCacheDatabaseInfo(ExecuteContext context) {
+        // This interception point is mainly used to store link information into thread variables, without the need to
+        // cache basic database information, so the method is implemented as empty
+    }
+
+    @Override
+    public ExecuteContext after(ExecuteContext context) throws Exception {
+        if (handler != null) {
+            handler.doAfter(context);
+            return context;
+        }
+        ThreadDatabaseUrlUtil.removeDatabaseUrl();
+        return context;
+    }
+}

sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/java/com/huaweicloud/sermant/postgresqlv42/utils/QueryExecutorImplEnhancementHelper.java → sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/java/com/huaweicloud/sermant/postgresqlv42/utils/PostgresqlEnhancementHelper.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.postgresqlv42.interceptors.PgStatementInterceptor;
 import com.huaweicloud.sermant.postgresqlv42.interceptors.QueryExecutorImplInterceptor;
 
 /**
@@ -28,7 +29,7 @@
  * @author zhp
  * @since 2024-02-04
  **/
-public class QueryExecutorImplEnhancementHelper {
+public class PostgresqlEnhancementHelper {
     private static final String SEND_ONE_QUERY_METHOD_NAME = "sendOneQuery";
 
     private static final String ENHANCE_CLASS_NAME = "org.postgresql.core.v3.QueryExecutorImpl";
@@ -39,14 +40,30 @@ public class QueryExecutorImplEnhancementHelper {
 
     private static final String SIMPLE_PARAMETER_LIST_CLASS_NAME = "org.postgresql.core.v3.SimpleParameterList";
 
+    private static final String EXECUTE_METHOD_NAME = "execute";
+
+    private static final String EXECUTE_BATCH_METHOD_NAME = "executeBatch";
+
+    private static final String PG_STATEMENT_CLASS_NAME = "org.postgresql.jdbc.PgStatement";
+
+    private static final String QUERY_CLASS_NAME = "org.postgresql.core.Query";
+
+    private static final String PARAMETER_LIST_CLASS_NAME = "org.postgresql.core.ParameterList";
+
+    private static final String[] STATEMENT_EXECUTE_METHOD_PARAMS_TYPE = {
+            QUERY_CLASS_NAME,
+            PARAMETER_LIST_CLASS_NAME,
+            INT_CLASS_NAME
+    };
+
     private static final String[] EXECUTE_METHOD_PARAMS_TYPE = {
             SIMPLE_QUERY_CLASS_NAME,
             SIMPLE_PARAMETER_LIST_CLASS_NAME,
             INT_CLASS_NAME, INT_CLASS_NAME,
             INT_CLASS_NAME
     };
 
-    private QueryExecutorImplEnhancementHelper() {
+    private PostgresqlEnhancementHelper() {
     }
 
     private static MethodMatcher getSendOneQueryMethodMatcher() {
@@ -81,4 +98,60 @@ public static InterceptDeclarer getSendOneQueryInterceptDeclarer() {
     public static ClassMatcher getQueryExecutorImplClassMatcher() {
         return ClassMatcher.nameEquals(ENHANCE_CLASS_NAME);
     }
+
+    private static MethodMatcher getExecuteMethodMatcher() {
+        return MethodMatcher.nameEquals(EXECUTE_METHOD_NAME)
+                .and(MethodMatcher.paramTypesEqual(STATEMENT_EXECUTE_METHOD_PARAMS_TYPE));
+    }
+
+    private static MethodMatcher getExecuteBatchMethodMatcher() {
+        return MethodMatcher.nameEquals(EXECUTE_BATCH_METHOD_NAME);
+    }
+
+    /**
+     * Get ClassMatcher for PgStatement class
+     *
+     * @return ClassMatcher Database write operation handler
+     */
+    public static ClassMatcher getPgStatementClassMatcher() {
+        return ClassMatcher.nameEquals(PG_STATEMENT_CLASS_NAME);
+    }
+
+    /**
+     * Get the non-parameter interceptor declarer for PgStatement execute method
+     *
+     * @return InterceptDeclarer The non-parameter interceptor declarer for PgStatement execute method
+     */
+    public static InterceptDeclarer getPgStatementExecuteInterceptDeclarer() {
+        return InterceptDeclarer.build(getExecuteMethodMatcher(), new PgStatementInterceptor());
+    }
+
+    /**
+     * Get the parameterized interceptor declarer for the PgStatement execute method
+     *
+     * @param handler Database write operation handler
+     * @return InterceptDeclarer The parameterized interceptor declarer for the PgStatement execute method
+     */
+    public static InterceptDeclarer getPgStatementExecuteInterceptDeclarer(DatabaseHandler handler) {
+        return InterceptDeclarer.build(getExecuteMethodMatcher(), new PgStatementInterceptor(handler));
+    }
+
+    /**
+     * Get the non-parameter interceptor declarer for PgStatement executeBatch method
+     *
+     * @return InterceptDeclarer The non-parameter interceptor declarer for PgStatement executeBatch method
+     */
+    public static InterceptDeclarer getPgStatementExecuteBatchInterceptDeclarer() {
+        return InterceptDeclarer.build(getExecuteBatchMethodMatcher(), new PgStatementInterceptor());
+    }
+
+    /**
+     * Get the parameterized interceptor declarer for PgStatement executeBatch method
+     *
+     * @param handler Database write operation handler
+     * @return InterceptDeclarer The parameterized interceptor declarer for PgStatement executeBatch method
+     */
+    public static InterceptDeclarer getPgStatementExecuteBatchInterceptDeclarer(DatabaseHandler handler) {
+        return InterceptDeclarer.build(getExecuteBatchMethodMatcher(), new PgStatementInterceptor(handler));
+    }
 }

sermant-plugins/sermant-database-write-prohibition/postgresql-42.x-plugin/src/main/resources/META-INF/services/com.huaweicloud.sermant.core.plugin.agent.declarer.PluginDeclarer

@@ -13,4 +13,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-com.huaweicloud.sermant.postgresqlv42.declarers.QueryExecutorImplDeclarer
+com.huaweicloud.sermant.postgresqlv42.declarers.QueryExecutorImplDeclarer
+com.huaweicloud.sermant.postgresqlv42.declarers.PgStatementDeclarer