SNOW-909180: Pass query id in exceptions from file agent (#1581)

snowflakedb · Jan 2, 2024 · 55e5b15 · 55e5b15
1 parent 301005e
commit 55e5b15
Show file tree

Hide file tree

Showing 16 changed files with 787 additions and 180 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeSQLException.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeSQLException.java
@@ -51,31 +51,43 @@ public SnowflakeSQLException(String queryId, String reason, String sqlState, int
         queryId);
   }
 
-  public SnowflakeSQLException(String reason, String SQLState) {
-    super(reason, SQLState);
+  public SnowflakeSQLException(String reason, String sqlState) {
+    super(reason, sqlState);
     // log user error from GS at fine level
-    logger.debug("Snowflake exception: {}, sqlState:{}", reason, SQLState);
+    logger.debug("Snowflake exception: {}, sqlState:{}", reason, sqlState);
   }
 
+  /** use {@link SnowflakeSQLException#SnowflakeSQLException(String, String, int)} */
+  @Deprecated
   public SnowflakeSQLException(String sqlState, int vendorCode) {
+    this((String) null, sqlState, vendorCode);
+  }
+
+  public SnowflakeSQLException(String queryId, String sqlState, int vendorCode) {
     super(
         errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode)),
         sqlState,
         vendorCode);
-
+    this.queryId = queryId;
     logger.debug(
         "Snowflake exception: {}, sqlState:{}, vendorCode:{}",
         errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode)),
         sqlState,
         vendorCode);
   }
 
+  /** use {@link SnowflakeSQLException#SnowflakeSQLException(String, String, int, Object...)} */
+  @Deprecated
   public SnowflakeSQLException(String sqlState, int vendorCode, Object... params) {
+    this((String) null, sqlState, vendorCode, params);
+  }
+
+  public SnowflakeSQLException(String queryId, String sqlState, int vendorCode, Object... params) {
     super(
         errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode), params),
         sqlState,
         vendorCode);
-
+    this.queryId = queryId;
     logger.debug(
         "Snowflake exception: {}, sqlState:{}, vendorCode:{}",
         errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode), params),
@@ -100,12 +112,23 @@ public SnowflakeSQLException(Throwable ex, ErrorCode errorCode, Object... params
     this(ex, errorCode.getSqlState(), errorCode.getMessageCode(), params);
   }
 
+  /**
+   * @deprecated use {@link SnowflakeSQLException#SnowflakeSQLException(String, Throwable, String,
+   *     int, Object...)}
+   */
+  @Deprecated
   public SnowflakeSQLException(Throwable ex, String sqlState, int vendorCode, Object... params) {
+    this(null, ex, sqlState, vendorCode, params);
+  }
+
+  public SnowflakeSQLException(
+      String queryId, Throwable ex, String sqlState, int vendorCode, Object... params) {
     super(
         errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode), params),
         sqlState,
         vendorCode,
         ex);
+    this.queryId = queryId;
 
     logger.debug(
         "Snowflake exception: "
@@ -121,6 +144,15 @@ public SnowflakeSQLException(ErrorCode errorCode, Object... params) {
         errorCode.getMessageCode());
   }
 
+  public SnowflakeSQLException(String queryId, ErrorCode errorCode, Object... params) {
+    super(
+        errorResourceBundleManager.getLocalizedMessage(
+            String.valueOf(errorCode.getMessageCode()), params),
+        errorCode.getSqlState(),
+        errorCode.getMessageCode());
+    this.queryId = queryId;
+  }
+
   public SnowflakeSQLException(
       ErrorCode errorCode, int retryCount, boolean issocketTimeoutNoBackoff, long elapsedSeconds) {
     super(

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeSQLLoggedException.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeSQLLoggedException.java
@@ -232,47 +232,58 @@ public SnowflakeSQLLoggedException(SFBaseSession session, int vendorCode, String
     sendTelemetryData(null, SQLState, vendorCode, session, this);
   }
 
+  public SnowflakeSQLLoggedException(
+      String queryId, SFBaseSession session, int vendorCode, String SQLState) {
+    super(queryId, SQLState, vendorCode);
+    sendTelemetryData(queryId, SQLState, vendorCode, session, this);
+  }
+
   public SnowflakeSQLLoggedException(SFBaseSession session, String SQLState, String reason) {
     super(reason, SQLState);
     sendTelemetryData(null, SQLState, -1, session, this);
   }
 
   public SnowflakeSQLLoggedException(
       SFBaseSession session, int vendorCode, String SQLState, Object... params) {
-    super(SQLState, vendorCode, params);
-    String reason =
-        errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode), params);
-    sendTelemetryData(null, SQLState, vendorCode, session, this);
+    this(null, session, vendorCode, SQLState, params);
+  }
+
+  public SnowflakeSQLLoggedException(
+      String queryId, SFBaseSession session, int vendorCode, String SQLState, Object... params) {
+    super(queryId, SQLState, vendorCode, params);
+    sendTelemetryData(queryId, SQLState, vendorCode, session, this);
   }
 
   public SnowflakeSQLLoggedException(
       SFBaseSession session, ErrorCode errorCode, Throwable ex, Object... params) {
     super(ex, errorCode, params);
-    // add telemetry
     sendTelemetryData(null, errorCode.getSqlState(), errorCode.getMessageCode(), session, this);
   }
 
   public SnowflakeSQLLoggedException(
       SFBaseSession session, String SQLState, int vendorCode, Throwable ex, Object... params) {
     super(ex, SQLState, vendorCode, params);
-    // add telemetry
-    String reason =
-        errorResourceBundleManager.getLocalizedMessage(String.valueOf(vendorCode), params);
     sendTelemetryData(null, SQLState, vendorCode, session, this);
   }
 
+  public SnowflakeSQLLoggedException(
+      String queryId,
+      SFBaseSession session,
+      String SQLState,
+      int vendorCode,
+      Throwable ex,
+      Object... params) {
+    super(queryId, ex, SQLState, vendorCode, params);
+    sendTelemetryData(queryId, SQLState, vendorCode, session, this);
+  }
+
   public SnowflakeSQLLoggedException(SFBaseSession session, ErrorCode errorCode, Object... params) {
     super(errorCode, params);
-    // add telemetry
-    String reason =
-        errorResourceBundleManager.getLocalizedMessage(
-            String.valueOf(errorCode.getMessageCode()), params);
     sendTelemetryData(null, null, -1, session, this);
   }
 
   public SnowflakeSQLLoggedException(SFBaseSession session, SFException e) {
     super(e);
-    // add telemetry
     sendTelemetryData(null, null, -1, session, this);
   }
 

diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/QueryIdHelper.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/QueryIdHelper.java
@@ -0,0 +1,12 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+package net.snowflake.client.jdbc.cloud.storage;
+
+import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial;
+
+class QueryIdHelper {
+  static String queryIdFromEncMatOr(RemoteStoreFileEncryptionMaterial encMat, String queryId) {
+    return encMat != null && encMat.getQueryId() != null ? encMat.getQueryId() : queryId;
+  }
+}