feat: Add JDBC log S3 upload with IAM role support

kentoyoshida · kentoyoshida · commit 9239ca740611 · 2025-06-27T06:45:31.000Z
- Add JdbcLogUploader class for automatic JDBC log upload to S3
- Support both IAM role and explicit AWS credentials authentication
- Add configuration options for S3 upload (bucket, prefix, region, credentials)
- Enable JDBC tracing when log upload is enabled
- Upload logs only when communication errors occur
- Implement proper resource management with AutoCloseable
- Add cross-platform support using system temp directory
- Update README with configuration examples and usage instructions
- Add AWS SDK v2 dependencies with conflict exclusions

This feature helps with debugging Snowflake connection issues by
automatically uploading JDBC driver logs to S3 when errors occur.
diff --git a/README.md b/README.md
@@ -30,6 +30,12 @@ Snowflake output plugin for Embulk loads records to Snowflake.
 - **merge_rule**: list of column assignments for updating existing records used in merge mode, for example `"foo" = T."foo" + S."foo"` (`T` means target table and `S` means source table). (string array, default: always overwrites with new values)
 - **batch_size**: size of a single batch insert (integer, default: 16777216)
 - **match_by_column_name**: specify whether to load semi-structured data into columns in the target table that match corresponding columns represented in the data. ("case_sensitive", "case_insensitive", "none", default: "none")
+- **upload_jdbc_log_to_s3**: enable automatic upload of JDBC driver logs to S3 when communication errors occur (boolean, default: false)
+- **s3_bucket**: S3 bucket name for JDBC log upload (string, required when upload_jdbc_log_to_s3 is true)
+- **s3_prefix**: S3 key prefix for JDBC log upload (string, required when upload_jdbc_log_to_s3 is true)
+- **s3_region**: AWS region for S3 bucket (string, required when upload_jdbc_log_to_s3 is true)
+- **s3_access_key_id**: AWS access key ID for S3 access (string, optional - uses IAM role if not specified)
+- **s3_secret_access_key**: AWS secret access key for S3 access (string, optional - uses IAM role if not specified)
 - **default_timezone**: If input column type (embulk type) is timestamp, this plugin needs to format the timestamp into a SQL string. This default_timezone option is used to control the timezone. You can overwrite timezone for each columns using column_options option. (string, default: `UTC`)
 - **column_options**: advanced: a key-value pairs where key is a column name and value is options for the column.
   - **type**: type of a column when this plugin creates new tables (e.g. `VARCHAR(255)`, `INTEGER NOT NULL UNIQUE`). This used when this plugin creates intermediate tables (insert, truncate_insert and merge modes), when it creates the target table (insert_direct and replace modes), and when it creates nonexistent target table automatically. (string, default: depends on input column type. `BIGINT` if input column type is long, `BOOLEAN` if boolean, `DOUBLE PRECISION` if double, `CLOB` if string, `TIMESTAMP` if timestamp)
@@ -62,6 +68,49 @@ Snowflake output plugin for Embulk loads records to Snowflake.
   * Transactional: Yes.
   * Resumable: No.
 
+## JDBC Log Upload to S3
+
+This plugin supports automatic upload of JDBC driver logs to S3 when communication errors occur. This feature is useful for debugging connection issues with Snowflake.
+
+### Configuration Example
+
+```yaml
+out:
+  type: snowflake
+  host: your-account.snowflakecomputing.com
+  user: your_user
+  password: your_password
+  warehouse: your_warehouse
+  database: your_database
+  schema: your_schema
+  table: your_table
+  mode: insert
+  
+  # JDBC log upload configuration
+  upload_jdbc_log_to_s3: true
+  s3_bucket: your-log-bucket
+  s3_prefix: snowflake-jdbc-logs
+  s3_region: us-east-1
+  
+  # Optional: Explicit AWS credentials (uses IAM role if not specified)
+  s3_access_key_id: YOUR_ACCESS_KEY_ID
+  s3_secret_access_key: YOUR_SECRET_ACCESS_KEY
+```
+
+### Authentication Methods
+
+1. **IAM Role (Recommended)**: Leave `s3_access_key_id` and `s3_secret_access_key` unspecified. The plugin will use the default AWS credentials provider chain (IAM role, environment variables, etc.).
+
+2. **Explicit Credentials**: Specify both `s3_access_key_id` and `s3_secret_access_key` for explicit authentication.
+
+### Behavior
+
+- JDBC logs are only uploaded when communication errors occur during Snowflake operations
+- The plugin automatically finds the latest `snowflake_jdbc*.log` file in the system temp directory
+- Logs are uploaded to `s3://{bucket}/{prefix}/{filename}`
+- If S3 upload fails, a warning is logged but the original error is still thrown
+- If required S3 configuration is missing, a warning is logged and log upload is skipped
+
 ## Build
 
 ```
diff --git a/build.gradle b/build.gradle
@@ -52,6 +52,24 @@ dependencies {
 
     compile "org.embulk:embulk-output-jdbc:0.10.2"
     compile "net.snowflake:snowflake-jdbc:3.13.26"
+
+    implementation platform("software.amazon.awssdk:bom:2.25.20")
+
+    implementation("software.amazon.awssdk:s3") {
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-annotations"
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-core"
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-databind"
+        exclude group: "joda-time", module: "joda-time"
+        exclude group: "commons-logging", module: "commons-logging"
+    }
+
+    implementation("software.amazon.awssdk:regions") {
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-annotations"
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-core"
+        exclude group: "com.fasterxml.jackson.core", module: "jackson-databind"
+        exclude group: "joda-time", module: "joda-time"
+        exclude group: "commons-logging", module: "commons-logging"
+    }
 }
 embulkPlugin {
     mainClass = "org.embulk.output.SnowflakeOutputPlugin"
diff --git a/src/main/java/org/embulk/output/SnowflakeOutputPlugin.java b/src/main/java/org/embulk/output/SnowflakeOutputPlugin.java
@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
+import java.io.File;
 import java.io.IOException;
 import java.sql.SQLException;
 import java.sql.Types;
@@ -15,6 +16,7 @@
 import org.embulk.config.ConfigSource;
 import org.embulk.config.TaskSource;
 import org.embulk.output.jdbc.*;
+import org.embulk.output.s3.JdbcLogUploader;
 import org.embulk.output.snowflake.PrivateKeyReader;
 import org.embulk.output.snowflake.SnowflakeCopyBatchInsert;
 import org.embulk.output.snowflake.SnowflakeOutputConnection;
@@ -91,6 +93,30 @@ public interface SnowflakePluginTask extends PluginTask {
     @ConfigDefault("\"none\"")
     public MatchByColumnName getMatchByColumnName();
 
+    @Config("upload_jdbc_log_to_s3")
+    @ConfigDefault("false")
+    public boolean getUploadJdbcLogToS3();
+
+    @Config("s3_bucket")
+    @ConfigDefault("null")
+    public String getS3Bucket();
+
+    @Config("s3_prefix")
+    @ConfigDefault("null")
+    public String getS3Prefix();
+
+    @Config("s3_region")
+    @ConfigDefault("null")
+    public String getS3Region();
+
+    @Config("s3_access_key_id")
+    @ConfigDefault("null")
+    public String getS3AccessKeyId();
+
+    @Config("s3_secret_access_key")
+    @ConfigDefault("null")
+    public String getS3SecretAccessKey();
+
     public void setCopyIntoTableColumnNames(String[] columnNames);
 
     public String[] getCopyIntoTableColumnNames();
@@ -139,6 +165,9 @@ public static MatchByColumnName fromString(String value) {
   private static final int MASTER_TOKEN_INVALID_GS_CODE = 390115;
   private static final int ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE = 390195;
 
+  private static final String ENCOUNTERED_COMMUNICATION_ERROR_MESSAGE =
+      "JDBC driver encountered communication error";
+
   @Override
   protected Class<? extends PluginTask> getTaskClass() {
     return SnowflakePluginTask.class;
@@ -198,6 +227,10 @@ protected JdbcOutputConnector getConnector(PluginTask task, boolean retryableMet
     props.setProperty("CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX", "true");
     props.setProperty("MULTI_STATEMENT_COUNT", "0");
 
+    if (t.getUploadJdbcLogToS3()) {
+      props.setProperty("tracing", "ALL");
+    }
+
     props.putAll(t.getOptions());
 
     logConnectionProperties(url, props);
@@ -222,6 +255,42 @@ public ConfigDiff transaction(
         runDropStageWithRecovery(snowflakeCon, stageIdentifier, task);
       }
     } catch (Exception e) {
+      if (e instanceof SQLException) {
+        String message = e.getMessage();
+        if (message != null
+            && message.contains(ENCOUNTERED_COMMUNICATION_ERROR_MESSAGE)
+            && t.getUploadJdbcLogToS3() == true) {
+          final String s3Bucket = t.getS3Bucket();
+          final String s3Prefix = t.getS3Prefix();
+          final String s3Region = t.getS3Region();
+          final String s3AccessKeyId = t.getS3AccessKeyId();
+          final String s3SecretAccessKey = t.getS3SecretAccessKey();
+          if (s3Bucket == null || s3Prefix == null || s3Region == null) {
+            logger.warn("s3_bucket, s3_prefix, and s3_region must be set when upload_jdbc_log_to_s3 is true");
+          } else {
+            try (JdbcLogUploader jdbcLogUploader = new JdbcLogUploader(s3Bucket, s3Prefix, s3Region, s3AccessKeyId, s3SecretAccessKey)) {
+              // snowflake_jdbc*.log で最新のファイルを探してアップロード
+              String tmpDir = System.getProperty("java.io.tmpdir", "/tmp");
+              File logDir = new File(tmpDir);
+              File[] logFiles =
+                  logDir.listFiles(
+                      (dir, name) -> name.startsWith("snowflake_jdbc") && name.endsWith(".log"));
+              if (logFiles != null && logFiles.length > 0) {
+                // 最終更新日時が新しいファイルを選択
+                Optional<File> latest =
+                    Arrays.stream(logFiles).max(Comparator.comparingLong(File::lastModified));
+                if (latest.isPresent()) {
+                  jdbcLogUploader.uploadIfExists(latest.get());
+                }
+              } else {
+                logger.warn("No snowflake_jdbc*.log file found in {} for upload", tmpDir);
+              }
+            } catch (Exception uploadException) {
+              logger.warn("Failed to upload JDBC log to S3: {}", uploadException.getMessage());
+            }
+          }
+        }
+      }
       if (t.getDeleteStage() && t.getDeleteStageOnError()) {
         try {
           runDropStageWithRecovery(snowflakeCon, stageIdentifier, task);
diff --git a/src/main/java/org/embulk/output/s3/JdbcLogUploader.java b/src/main/java/org/embulk/output/s3/JdbcLogUploader.java
@@ -0,0 +1,67 @@
+package org.embulk.output.s3;
+
+import java.io.File;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+
+public class JdbcLogUploader implements AutoCloseable {
+  private final Logger logger = LoggerFactory.getLogger(JdbcLogUploader.class);
+
+  private final S3Client s3Client;
+  private final String bucket;
+  private final String prefix;
+  private final Region region;
+
+  public JdbcLogUploader(String bucket, String prefix, String region, String accessKeyId, String secretAccessKey) {
+    this.bucket = bucket;
+    this.prefix = prefix;
+    this.region = Region.of(region);
+    
+    if (accessKeyId != null && secretAccessKey != null) {
+      // Use explicit credentials
+      AwsBasicCredentials awsCreds = AwsBasicCredentials.create(accessKeyId, secretAccessKey);
+      this.s3Client = S3Client.builder()
+        .region(this.region)
+        .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
+        .build();
+    } else {
+      // Use default credentials provider (IAM role, environment variables, etc.)
+      this.s3Client = S3Client.builder()
+        .region(this.region)
+        .credentialsProvider(DefaultCredentialsProvider.create())
+        .build();
+    }
+  }
+
+  public void uploadIfExists(File file) {
+    if (!file.exists()) {
+      logger.warn("File not found: {}", file.getAbsolutePath());
+      return;
+    }
+
+    String key = prefix + "/" + file.getName();
+    try {
+      PutObjectRequest putObjectRequest =
+          PutObjectRequest.builder().bucket(bucket).key(key).build();
+
+      s3Client.putObject(putObjectRequest, RequestBody.fromFile(file));
+      logger.info("Uploaded {} to s3://{}/{}", file.getAbsolutePath(), bucket, key);
+    } catch (Exception e) {
+      logger.error("Failed to upload {} to S3: {}", file.getAbsolutePath(), e.getMessage(), e);
+    }
+  }
+
+  @Override
+  public void close() {
+    if (s3Client != null) {
+      s3Client.close();
+    }
+  }
+}
