SNOW-180297 Added JDBC_TRACE env var and JDBC_LOG_DIR env var to allo…

…w for getting JDBC trace logs (#202) * Added JDBC_TRACE env var and JDBC_LOG_DIR env var to allow for getting JDBC trace logs * SNOW-180297 minor midifications to move logic into functions * Added test Co-authored-by: Chris Conner <[email protected]>
snowflakedb · Aug 4, 2020 · e989829 · e989829
1 parent 9c0206a
commit e989829
Show file tree

Hide file tree

Showing 8 changed files with 112 additions and 0 deletions.
diff --git a/pom.xml b/pom.xml
@@ -425,6 +425,12 @@
             <version>2.13.2</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.github.stefanbirkner</groupId>
+            <artifactId>system-rules</artifactId>
+            <version>1.19.0</version>
+            <scope>test</scope>
+        </dependency>
 
         <!--Kafka JSON converter for SMT unit test-->
         <dependency>

diff --git a/pom_confluent.xml b/pom_confluent.xml
@@ -483,6 +483,12 @@
             <version>2.13.2</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.github.stefanbirkner</groupId>
+            <artifactId>system-rules</artifactId>
+            <version>1.19.0</version>
+            <scope>test</scope>
+        </dependency>
 
         <!--Kafka JSON converter for SMT unit test-->
         <dependency>

diff --git a/src/main/java/com/snowflake/kafka/connector/SnowflakeSinkConnectorConfig.java b/src/main/java/com/snowflake/kafka/connector/SnowflakeSinkConnectorConfig.java
@@ -70,6 +70,12 @@ public class SnowflakeSinkConnectorConfig
   static final String JVM_PROXY_HOST = "jvm.proxy.host";
   static final String JVM_PROXY_PORT = "jvm.proxy.port";
 
+  // JDBC logging directory Info (environment variable)
+  static final String SNOWFLAKE_JDBC_LOG_DIR = "JDBC_LOG_DIR";
+
+  // JDBC trace Info (environment variable)
+  public static final String SNOWFLAKE_JDBC_TRACE = "JDBC_TRACE";
+
   // Snowflake Metadata Flags
   private static final String SNOWFLAKE_METADATA_FLAGS = "Snowflake Metadata Flags";
   public static final String SNOWFLAKE_METADATA_CREATETIME = "snowflake.metadata.createtime";

diff --git a/src/main/java/com/snowflake/kafka/connector/Utils.java b/src/main/java/com/snowflake/kafka/connector/Utils.java
@@ -25,6 +25,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.BufferedReader;
+import java.io.File;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.URL;
@@ -68,6 +69,9 @@ public class Utils
   private static final String HTTP_PROXY_HOST = "http.proxyHost";
   private static final String HTTP_PROXY_PORT = "http.proxyPort";
 
+  //jdbc log dir
+  public static final String JAVA_IO_TMPDIR = "java.io.tmpdir";
+
   private static final Random random = new Random();
 
   //mvn repo
@@ -177,6 +181,26 @@ public static String pipeName(String appName, String table, int partition)
     return pipeName;
   }
 
+  /**
+   * Read JDBC logging directory from environment variable JDBC_LOG_DIR and set that in System property
+   */
+  public static void setJDBCLoggingDirectory()
+  {
+    String jdbcTmpDir = System.getenv(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_LOG_DIR);
+    if (jdbcTmpDir != null)
+    {
+      File jdbcTmpDirObj = new File(jdbcTmpDir);
+      if (jdbcTmpDirObj.isDirectory())
+      {
+        LOGGER.info(Logging.logMessage("jdbc tracing directory = {}", jdbcTmpDir));
+        System.setProperty(JAVA_IO_TMPDIR, jdbcTmpDir);
+      } else {
+        LOGGER.info(Logging.logMessage("invalid JDBC_LOG_DIR {} defaulting to {}", jdbcTmpDir,
+          System.getProperty(JAVA_IO_TMPDIR)));
+      }
+    }
+  }
+
   /**
    * Enable JVM proxy
    *
@@ -397,6 +421,9 @@ static String validateConfig(Map<String, String> config)
     // jvm proxy settings
     configIsValid = Utils.enableJVMProxy(config) && configIsValid;
 
+    // set jdbc logging directory
+    Utils.setJDBCLoggingDirectory();
+
     if (!configIsValid)
     {
       throw SnowflakeErrors.ERROR_0001.getException();

diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeConnectionServiceV1.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeConnectionServiceV1.java
@@ -7,6 +7,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
+import java.io.File;
 import java.nio.charset.StandardCharsets;
 import java.security.PrivateKey;
 import java.sql.Connection;

diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeURL.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeURL.java
@@ -17,6 +17,9 @@
 
 package com.snowflake.kafka.connector.internal;
 
+import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
+import net.snowflake.client.jdbc.internal.apache.arrow.util.VisibleForTesting;
+
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -69,11 +72,26 @@ else if (ssl)
     }
 
     jdbcUrl = "jdbc:snowflake://" + url + ":" + port;
+    if (enableJDBCTrace())
+    {
+      logInfo("enabling JDBC tracing");
+      jdbcUrl = jdbcUrl + "/?tracing=ALL";
+    }
 
     logDebug("parsed Snowflake URL: {}", urlStr);
 
   }
 
+  /**
+   * Read environment variable JDBC_TRACE to check whether trace is enabled
+   * @return whether to enable JDBC trace
+   */
+  boolean enableJDBCTrace()
+  {
+    String enableJDBCTrace = System.getenv(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_TRACE);
+    return enableJDBCTrace != null && enableJDBCTrace.toLowerCase().contains("true");
+  }
+
   String getJdbcUrl()
   {
     return jdbcUrl;

diff --git a/src/test/java/com/snowflake/kafka/connector/UtilsTest.java b/src/test/java/com/snowflake/kafka/connector/UtilsTest.java
@@ -2,13 +2,19 @@
 
 import com.snowflake.kafka.connector.internal.SnowflakeErrors;
 import com.snowflake.kafka.connector.internal.TestUtils;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.contrib.java.lang.system.EnvironmentVariables;
 
 import java.util.HashMap;
 import java.util.Map;
 
 public class UtilsTest
 {
+  @Rule
+  public final EnvironmentVariables environmentVariables
+    = new EnvironmentVariables();
+
   @Test
   public void testGetTopicToTableMap()
   {
@@ -119,4 +125,25 @@ public void testIsValidSnowflakeApplicationName()
     assert !Utils.isValidSnowflakeApplicationName("(1.f$-_");
   }
 
+  @Test
+  public void testSetJDBCLoggingDir()
+  {
+    String defaultTmpDir = System.getProperty(Utils.JAVA_IO_TMPDIR);
+
+    Utils.setJDBCLoggingDirectory();
+    assert !System.getProperty(Utils.JAVA_IO_TMPDIR).isEmpty();
+
+    environmentVariables.set(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_LOG_DIR, "/dummy_dir_not_exist");
+    Utils.setJDBCLoggingDirectory();
+    assert !System.getProperty(Utils.JAVA_IO_TMPDIR).equals("/dummy_dir_not_exist");
+
+    environmentVariables.set(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_LOG_DIR, "/usr");
+    Utils.setJDBCLoggingDirectory();
+    assert System.getProperty(Utils.JAVA_IO_TMPDIR).equals("/usr");
+
+    environmentVariables.set(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_LOG_DIR, defaultTmpDir);
+    Utils.setJDBCLoggingDirectory();
+    assert System.getProperty(Utils.JAVA_IO_TMPDIR).equals(defaultTmpDir);
+  }
+
 }
diff --git a/src/test/java/com/snowflake/kafka/connector/internal/SnowflakeURLTest.java b/src/test/java/com/snowflake/kafka/connector/internal/SnowflakeURLTest.java
@@ -16,10 +16,18 @@
  */
 package com.snowflake.kafka.connector.internal;
 
+import com.snowflake.kafka.connector.SnowflakeSinkConnectorConfig;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.contrib.java.lang.system.EnvironmentVariables;
+
+import java.io.IOException;
 
 public class SnowflakeURLTest
 {
+  @Rule
+  public final EnvironmentVariables environmentVariables
+    = new EnvironmentVariables();
 
   @Test
   public void createFromValidURL()
@@ -97,6 +105,19 @@ public void createFromValidURL()
     assert sfurl.getUrlWithoutPort().equals("account.region.aws.privatelink.snowflake.com");
   }
 
+  @Test
+  public void testJDBCTracing()
+  {
+    environmentVariables.set(SnowflakeSinkConnectorConfig.SNOWFLAKE_JDBC_TRACE, "true");
+
+    String url = "https://account.region.aws.privatelink.snowflake.com:443";
+
+    SnowflakeURL sfurl = new SnowflakeURL(url);
+
+    assert sfurl.getJdbcUrl().equals("jdbc:snowflake://account.region.aws.privatelink.snowflake.com:443/?tracing=ALL");
+
+  }
+
   @Test(expected = SnowflakeKafkaConnectorException.class)
   public void createFromInvalidURL()
   {