Fix NPE in PostgreSQL plugin when executing statements prepared with generated keys

CHANGES.md

@@ -5,6 +5,8 @@ Release Notes.
 9.9.0
 ------------------
 
+* Fix the `NullPointerException` thrown by the `postgresql-8.x` plugin when executing statements prepared with
+  `Statement.RETURN_GENERATED_KEYS` or column names (apache/skywalking#14138).
 
 All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/267?closed=1)
 

apm-sniffer/apm-sdk-plugin/postgresql-8.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/jdbc/postgresql/PreparedStatementExecuteMethodsInterceptor.java

@@ -43,6 +43,13 @@ public class PreparedStatementExecuteMethodsInterceptor implements InstanceMetho
     public final void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
                                    Class<?>[] argumentsTypes, MethodInterceptResult result) {
         StatementEnhanceInfos cacheObject = (StatementEnhanceInfos) objInst.getSkyWalkingDynamicField();
+        /*
+         * The statement created by PgConnection#prepareStatement(String, String[]) has no dynamic field. It is
+         * wrapped and traced by SWPreparedStatement, see JDBCPrepareStatementWithStringArrayInterceptor.
+         */
+        if (cacheObject == null || cacheObject.getConnectionInfo() == null) {
+            return;
+        }
         ConnectionInfo connectInfo = cacheObject.getConnectionInfo();
         AbstractSpan span = ContextManager.createExitSpan(
             buildOperationName(connectInfo, method.getName(), cacheObject.getStatementName()), connectInfo
@@ -68,7 +75,7 @@ public final void beforeMethod(EnhancedInstance objInst, Method method, Object[]
     public final Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
                                     Class<?>[] argumentsTypes, Object ret) {
         StatementEnhanceInfos cacheObject = (StatementEnhanceInfos) objInst.getSkyWalkingDynamicField();
-        if (cacheObject.getConnectionInfo() != null) {
+        if (cacheObject != null && cacheObject.getConnectionInfo() != null) {
             ContextManager.stopSpan();
         }
         return ret;
@@ -78,7 +85,7 @@ public final Object afterMethod(EnhancedInstance objInst, Method method, Object[
     public final void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
                                             Class<?>[] argumentsTypes, Throwable t) {
         StatementEnhanceInfos cacheObject = (StatementEnhanceInfos) objInst.getSkyWalkingDynamicField();
-        if (cacheObject.getConnectionInfo() != null) {
+        if (cacheObject != null && cacheObject.getConnectionInfo() != null) {
             ContextManager.activeSpan().log(t);
         }
     }

apm-sniffer/apm-sdk-plugin/postgresql-8.x-plugin/src/test/java/org/apache/skywalking/apm/plugin/jdbc/postgresql/PreparedStatementExecuteMethodsInterceptorTest.java

@@ -0,0 +1,115 @@
+/*
+ * 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.skywalking.apm.plugin.jdbc.postgresql;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.when;
+import java.lang.reflect.Method;
+import java.sql.SQLException;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
+import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
+import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
+import org.apache.skywalking.apm.agent.test.helper.SegmentHelper;
+import org.apache.skywalking.apm.agent.test.tools.AgentServiceRule;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStorage;
+import org.apache.skywalking.apm.agent.test.tools.SegmentStoragePoint;
+import org.apache.skywalking.apm.agent.test.tools.SpanAssert;
+import org.apache.skywalking.apm.agent.test.tools.TracingSegmentRunner;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
+import org.apache.skywalking.apm.plugin.jdbc.JDBCPluginConfig;
+import org.apache.skywalking.apm.plugin.jdbc.define.StatementEnhanceInfos;
+import org.apache.skywalking.apm.plugin.jdbc.trace.ConnectionInfo;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+@RunWith(TracingSegmentRunner.class)
+public class PreparedStatementExecuteMethodsInterceptorTest {
+
+    private static final String SQL = "INSERT INTO test (name) VALUES (?)";
+
+    @SegmentStoragePoint
+    private SegmentStorage segmentStorage;
+
+    @Rule
+    public AgentServiceRule serviceRule = new AgentServiceRule();
+    @Rule
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    private PreparedStatementExecuteMethodsInterceptor interceptor;
+
+    @Mock
+    private ConnectionInfo connectionInfo;
+    @Mock
+    private EnhancedInstance objectInstance;
+    @Mock
+    private Method method;
+
+    @Before
+    public void setUp() {
+        JDBCPluginConfig.Plugin.JDBC.SQL_BODY_MAX_LENGTH = 2048;
+        interceptor = new PreparedStatementExecuteMethodsInterceptor();
+    }
+
+    @Test
+    public void testCreateDatabaseSpan() throws Throwable {
+        when(objectInstance.getSkyWalkingDynamicField()).thenReturn(
+            new StatementEnhanceInfos(connectionInfo, SQL, "PreparedStatement"));
+        when(method.getName()).thenReturn("executeUpdate");
+        when(connectionInfo.getComponent()).thenReturn(ComponentsDefine.POSTGRESQL_DRIVER);
+        when(connectionInfo.getDBType()).thenReturn("PostgreSQL");
+        when(connectionInfo.getDatabaseName()).thenReturn("test");
+        when(connectionInfo.getDatabasePeer()).thenReturn("localhost:5432");
+
+        interceptor.beforeMethod(objectInstance, method, new Object[0], null, null);
+        interceptor.afterMethod(objectInstance, method, new Object[0], null, null);
+
+        assertThat(segmentStorage.getTraceSegments().size(), is(1));
+        TraceSegment segment = segmentStorage.getTraceSegments().get(0);
+        assertThat(SegmentHelper.getSpans(segment).size(), is(1));
+        AbstractTracingSpan span = SegmentHelper.getSpans(segment).get(0);
+        SpanAssert.assertLayer(span, SpanLayer.DB);
+        assertThat(span.getOperationName(), is("PostgreSQL/JDBC/PreparedStatement/executeUpdate"));
+        SpanAssert.assertTag(span, 0, "PostgreSQL");
+        SpanAssert.assertTag(span, 1, "test");
+        SpanAssert.assertTag(span, 2, SQL);
+    }
+
+    /**
+     * A statement created by {@code PgConnection#prepareStatement(String, String[])}, which is also reached through
+     * {@code prepareStatement(String, int)} with {@code RETURN_GENERATED_KEYS}, is wrapped by {@code SWPreparedStatement}
+     * without setting the dynamic field, so the interceptor must skip it instead of throwing NPE.
+     */
+    @Test
+    public void testSkipStatementWithoutEnhanceInfos() throws Throwable {
+        when(objectInstance.getSkyWalkingDynamicField()).thenReturn(null);
+
+        interceptor.beforeMethod(objectInstance, method, new Object[0], null, null);
+        interceptor.handleMethodException(objectInstance, method, new Object[0], null, new SQLException());
+        interceptor.afterMethod(objectInstance, method, new Object[0], null, null);
+
+        assertThat(segmentStorage.getTraceSegments().size(), is(0));
+    }
+}