feat: include SQL statement in error message (googleapis#355)

thiagotnunes · Jul 17, 2020 · cc5ac48 · cc5ac48
1 parent cea51ad
commit cc5ac48
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 4 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java
@@ -607,7 +607,7 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(Iterable<Stateme
   }
 
   ResultSet executeQueryInternalWithOptions(
-      Statement statement,
+      final Statement statement,
       com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode,
       Options options,
       ByteString partitionToken) {
@@ -622,7 +622,7 @@ ResultSet executeQueryInternalWithOptions(
         new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {
           @Override
           CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
-            GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
+            GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
             if (resumeToken != null) {
               request.setResumeToken(resumeToken);
             }

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java
@@ -763,16 +763,24 @@ interface CloseableIterator<T> extends Iterator<T> {
   @VisibleForTesting
   static class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
       implements CloseableIterator<PartialResultSet> {
+    private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
     private static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
 
     private final ConsumerImpl consumer = new ConsumerImpl();
     private final BlockingQueue<PartialResultSet> stream;
+    private final Statement statement;
 
     private SpannerRpc.StreamingCall call;
     private SpannerException error;
 
-    // Visible for testing.
+    @VisibleForTesting
     GrpcStreamIterator(int prefetchChunks) {
+      this(null, prefetchChunks);
+    }
+
+    @VisibleForTesting
+    GrpcStreamIterator(Statement statement, int prefetchChunks) {
+      this.statement = statement;
       // One extra to allow for END_OF_STREAM message.
       this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
     }
@@ -839,6 +847,23 @@ public void onCompleted() {
 
       @Override
       public void onError(SpannerException e) {
+        if (statement != null) {
+          if (logger.isLoggable(Level.FINEST)) {
+            // Include parameter values if logging level is set to FINEST or higher.
+            e =
+                SpannerExceptionFactory.newSpannerExceptionPreformatted(
+                    e.getErrorCode(),
+                    String.format("%s - Statement: '%s'", e.getMessage(), statement.toString()),
+                    e);
+            logger.log(Level.FINEST, "Error executing statement", e);
+          } else {
+            e =
+                SpannerExceptionFactory.newSpannerExceptionPreformatted(
+                    e.getErrorCode(),
+                    String.format("%s - Statement: '%s'", e.getMessage(), statement.getSql()),
+                    e);
+          }
+        }
         error = e;
         addToStream(END_OF_STREAM);
       }

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java
@@ -188,7 +188,7 @@ private static ResourceInfo extractResourceInfo(Throwable cause) {
     return null;
   }
 
-  private static SpannerException newSpannerExceptionPreformatted(
+  static SpannerException newSpannerExceptionPreformatted(
       ErrorCode code, @Nullable String message, @Nullable Throwable cause) {
     // This is the one place in the codebase that is allowed to call constructors directly.
     DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED;

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -25,6 +25,7 @@
 import com.google.api.gax.grpc.testing.LocalChannelProvider;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.AbstractResultSet.GrpcStreamIterator;
 import com.google.cloud.spanner.AsyncResultSet.CallbackResponse;
 import com.google.cloud.spanner.AsyncResultSet.ReadyCallback;
 import com.google.cloud.spanner.AsyncRunner.AsyncWork;
@@ -33,6 +34,7 @@
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
 import com.google.common.base.Stopwatch;
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.protobuf.AbstractMessage;
 import com.google.spanner.v1.ExecuteSqlRequest;
@@ -52,6 +54,8 @@
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -1483,4 +1487,65 @@ public void testBatchCreateSessionsPermissionDenied() {
       }
     }
   }
+
+  @Test
+  public void testExceptionIncludesStatement() {
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(
+            Status.INVALID_ARGUMENT.withDescription("Invalid query").asRuntimeException()));
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    try (ResultSet rs =
+        client
+            .singleUse()
+            .executeQuery(
+                Statement.newBuilder("SELECT * FROM FOO WHERE ID=@id").bind("id").to(1L).build())) {
+      rs.next();
+      fail("missing expected exception");
+    } catch (SpannerException e) {
+      assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+      assertThat(e.getMessage()).contains("Statement: 'SELECT * FROM FOO WHERE ID=@id'");
+      // The error message should normally not include the parameter values to prevent sensitive
+      // information from accidentally being logged.
+      assertThat(e.getMessage()).doesNotContain("id: 1");
+    }
+
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(
+            Status.INVALID_ARGUMENT.withDescription("Invalid query").asRuntimeException()));
+    Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
+    Level currentLevel = logger.getLevel();
+    try (ResultSet rs =
+        client
+            .singleUse()
+            .executeQuery(
+                Statement.newBuilder("SELECT * FROM FOO WHERE ID=@id").bind("id").to(1L).build())) {
+      logger.setLevel(Level.FINEST);
+      rs.next();
+      fail("missing expected exception");
+    } catch (SpannerException e) {
+      // With log level set to FINEST the error should also include the parameter values.
+      assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+      assertThat(e.getMessage()).contains("Statement: 'SELECT * FROM FOO WHERE ID=@id {id: 1}'");
+    } finally {
+      logger.setLevel(currentLevel);
+    }
+  }
+
+  @Test
+  public void testReadDoesNotIncludeStatement() {
+    mockSpanner.setStreamingReadExecutionTime(
+        SimulatedExecutionTime.ofException(
+            Status.INVALID_ARGUMENT.withDescription("Invalid read").asRuntimeException()));
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    try (ResultSet rs =
+        client.singleUse().read("FOO", KeySet.singleKey(Key.of(1L)), ImmutableList.of("BAR"))) {
+      rs.next();
+      fail("missing expected exception");
+    } catch (SpannerException e) {
+      assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+      assertThat(e.getMessage()).doesNotContain("Statement:");
+    }
+  }
 }