embulk-input-athena 0.1.0 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: c3fb336d83e38770353f87f111859738619bb3e6
-  data.tar.gz: 7fe07d1b0fd72c8e08ba963e21b590c26242da79
+  metadata.gz: df05f5540bab30f61316b799b7cb80b849da32d6
+  data.tar.gz: 0a854b7fc362bdff91f8423bd788a119426a4ff0
 SHA512:
-  metadata.gz: 4483e0985f1741775a6abbf8a9183231be12f52b5f28dc278c850fbf926fba5a2090e9fd660fd6d5b93782dec551b644828e6b767c096254ccc269deb755728f
-  data.tar.gz: bb989515c90ebdd73c651099fcfe6b136d1bfcdc1a5c30c979bbb5f0b3e2d782b8212110cb282f3129888bd347cc6e9878ae5ab2deb4118d108d58201a765f2c
+  metadata.gz: b84f7b31683d9d239bd41bde61d35ee9d003fe5d0537b96b9f06310ac9123a752288197b671d7d11d335a215bb9334a4589013d9ac65d7cddef3cd033d53371f
+  data.tar.gz: c6ea2231c7e1dc107f89b4d0fdf3ba1ee7e495b1be14de86e3c75782fb2219b0818fa094e3d14ac23a492748b7acc5a5781e4ed798ed2a41115b8a783c3f1ca6

data/.gitignore

@@ -10,4 +10,5 @@ build/
 /.metadata/
 .classpath
 .project
-bin
+bin/
+.vscode/

data/README.md

@@ -1,5 +1,8 @@
 # Athena input plugin for Embulk
 
+[![Gem Version](https://badge.fury.io/rb/embulk-input-athena.svg)](https://badge.fury.io/rb/embulk-input-athena)
+[![MIT License](http://img.shields.io/badge/license-MIT-blue.svg?style=flat)](LICENSE)
+
 Athena input plugin for Embulk loads records from Athena(AWS).
 
 ## Overview
@@ -11,14 +14,16 @@ Athena input plugin for Embulk loads records from Athena(AWS).
 
 ## Configuration
 
-* **database**: description (string, required)
-* **athena_url**: description (string, required)
-* **s3_staging_dir**: description (string, required)
-* **access_key**: description (string, required)
-* **secret_key**: description (string, required)
-* **query**: description (string, required)
-* **columns**: description (string, required)
-* **options**: description (string, default: {})
+* **driver_path**: path to the jar file of the Athena JDBC driver. If not set, the bundled JDBC driver(AthenaJDBC41-1.1.0.jar) will be used. (string)
+* **database**: database name (string, required)
+* **athena_url**: Athena url (string, required)
+* **s3_staging_dir**: The S3 location to which your query output is written, for example s3://query-results-bucket/folder/. (string, required)
+* **access_key**: AWS access key (string, required)
+* **secret_key**: AWS secret key (string, required)
+* **query**: SQL to run (string, required)
+* **columns**: columns (string, required)
+* **options**: extra JDBC properties (string, default: {})
+* **null_to_zero**: if true, convert long, double and boolean value from null to zero (boolean, default: false)
 
 ## Example
 
@@ -35,12 +40,19 @@ in:
   columns:
     - {name: uid, type: string}
     - {name: created_at, type: timestamp}
+  null_to_zero: true
 ```
 
 ## Build
 
-```
+```bash
 $ docker-compose up -d
 $ docker-compose exec embulk bash
 embulk>$ ./gradlew gem  # -t to watch change of files and rebuild continuously
+
+embulk>$ embulk preview -I lib sample.yml
+embulk>$ embulk preview -L . sample.yml
+
+curl -u shinji19 https://rubygems.org/api/v1/api_key.yaml > ~/.gem/credentials
+embulk>$ ./gradlew gemPush
 ```

data/build.gradle

@@ -1,21 +1,22 @@
 plugins {
     id "com.jfrog.bintray" version "1.1"
-    id "com.github.jruby-gradle.base" version "0.1.5"
+    id "com.github.jruby-gradle.base" version "1.5.0"
     id "java"
     id "checkstyle"
+    // for task download
+    id "de.undercouch.download" version "3.4.2"
 }
 import com.github.jrubygradle.JRubyExec
 repositories {
     mavenCentral()
     jcenter()
-    // for athena jdbc
-    maven { url "https://maven.atlassian.com/repository/public" }
+    maven { url "https://dl.bintray.com/embulk-input-jdbc/maven" }
 }
 configurations {
     provided
 }
 
-version = "0.1.0"
+version = "0.1.6"
 
 sourceCompatibility = 1.8
 targetCompatibility = 1.8
@@ -23,16 +24,24 @@ targetCompatibility = 1.8
 dependencies {
     compile  "org.embulk:embulk-core:0.8.39"
     provided "org.embulk:embulk-core:0.8.39"
-    // https://mvnrepository.com/artifact/com.amazonaws.athena.jdbc/AthenaJDBC41
-    // compile group: 'com.amazonaws.athena.jdbc', name: 'AthenaJDBC41', version: '1.0.1-atlassian-hosted'
+    // TODO: maven...
     compile files ('build/AthenaJDBC41-1.1.0.jar')
     compile group: 'com.amazonaws', name: 'aws-java-sdk', version: '1.11.301'
+    // compile group: 'com.amazonaws', name: 'aws-java-sdk-athena', version: '1.11.301'
+    compile 'org.embulk.input.jdbc:embulk-input-jdbc:0.9.1'
     testCompile "junit:junit:4.+"
 }
 
-task classpath(type: Copy, dependsOn: ["jar"]) {
+task downloadFile(type: Download) {
+    src 'https://s3.amazonaws.com/athena-downloads/drivers/AthenaJDBC41-1.1.0.jar'
+    dest buildDir
+    onlyIfModified true
+}
+
+task classpath(type: Copy, dependsOn: ["downloadFile", "jar"]) {
     doFirst { file("classpath").deleteDir() }
     from (configurations.runtime - configurations.provided + files(jar.archivePath))
+    from ("build/AthenaJDBC41-1.1.0.jar'")
     into "classpath"
 }
 clean { delete "classpath" }
@@ -55,14 +64,16 @@ task checkstyle(type: Checkstyle) {
 }
 
 task gem(type: JRubyExec, dependsOn: ["gemspec", "classpath"]) {
-    jrubyArgs "-rrubygems/gem_runner", "-eGem::GemRunner.new.run(ARGV)", "build"
-    script "${project.name}.gemspec"
+    jrubyArgs "-S"
+    script "gem"
+    scriptArgs "build", "${project.name}.gemspec"
     doLast { ant.move(file: "${project.name}-${project.version}.gem", todir: "pkg") }
 }
 
 task gemPush(type: JRubyExec, dependsOn: ["gem"]) {
-    jrubyArgs "-rrubygems/gem_runner", "-eGem::GemRunner.new.run(ARGV)", "push"
-    script "pkg/${project.name}-${project.version}.gem"
+    jrubyArgs "-S"
+    script "gem"
+    scriptArgs "push", "pkg/${project.name}-${project.version}.gem"
 }
 
 task "package"(dependsOn: ["gemspec", "classpath"]) {

data/docker-compose.yml

@@ -7,4 +7,10 @@ services:
     working_dir: /root/embulk-input-athena
     volumes:
       - ./:/root/embulk-input-athena
-    command: tail -f /dev/null
+    tty: true
+  postgres:
+    image: postgres
+    ports:
+      - 5432:5432
+    environment:
+      - POSTGRES_PASSWORD=postgres

data/src/main/java/org/embulk/input/athena/AthenaInputConnection.java

@@ -1,3 +1,4 @@
+/*
 package org.embulk.input.athena;
 
 import java.util.List;
@@ -17,6 +18,7 @@ public class AthenaInputConnection
             throws SQLException
     {
         super(connection, null);
+        connection.setAutoCommit(true);
     }
 
     @Override
@@ -47,3 +49,4 @@ public class AthenaInputConnection
         return new SingleSelect(stmt);
     }
 }
+*/

data/src/main/java/org/embulk/input/athena/AthenaInputPlugin.java

@@ -1,9 +1,17 @@
 package org.embulk.input.athena;
 
+import com.google.common.base.Optional;
+
+import java.io.File;
+import java.io.FileFilter;
+import java.net.MalformedURLException;
+import java.net.URISyntaxException;
+import java.net.URL;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.List;
@@ -12,12 +20,14 @@ import java.util.Properties;
 import org.embulk.config.Config;
 import org.embulk.config.ConfigDefault;
 import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigException;
 import org.embulk.config.ConfigInject;
 import org.embulk.config.ConfigSource;
 import org.embulk.config.Task;
 import org.embulk.config.TaskReport;
 import org.embulk.config.TaskSource;
 import org.embulk.input.jdbc.ToStringMap;
+import org.embulk.plugin.PluginClassLoader;
 import org.embulk.spi.BufferAllocator;
 import org.embulk.spi.Column;
 import org.embulk.spi.ColumnVisitor;
@@ -28,9 +38,18 @@ import org.embulk.spi.PageOutput;
 import org.embulk.spi.Schema;
 import org.embulk.spi.SchemaConfig;
 import org.embulk.spi.time.Timestamp;
+import org.slf4j.Logger;
+
+public class AthenaInputPlugin implements InputPlugin
+{
+    protected final Logger logger = Exec.getLogger(getClass());
+
+    public interface PluginTask extends Task
+    {
+        @Config("driver_path")
+        @ConfigDefault("null")
+        public Optional<String> getDriverPath();
 
-public class AthenaInputPlugin implements InputPlugin {
-    public interface PluginTask extends Task {
         // database (required string)
         @Config("database")
         public String getDatabase();
@@ -63,12 +82,17 @@ public class AthenaInputPlugin implements InputPlugin {
         @ConfigDefault("{}")
         public ToStringMap getOptions();
 
+        @Config("null_to_zero")
+        @ConfigDefault("false")
+        public boolean getNullToZero();
+
         @ConfigInject
         BufferAllocator getBufferAllocator();
     }
 
     @Override
-    public ConfigDiff transaction(ConfigSource config, InputPlugin.Control control) {
+    public ConfigDiff transaction(ConfigSource config, InputPlugin.Control control)
+    {
         PluginTask task = config.loadConfig(PluginTask.class);
 
         Schema schema = task.getColumns().toSchema();
@@ -78,17 +102,20 @@ public class AthenaInputPlugin implements InputPlugin {
     }
 
     @Override
-    public ConfigDiff resume(TaskSource taskSource, Schema schema, int taskCount, InputPlugin.Control control) {
+    public ConfigDiff resume(TaskSource taskSource, Schema schema, int taskCount, InputPlugin.Control control)
+    {
         control.run(taskSource, schema, taskCount);
         return Exec.newConfigDiff();
     }
 
     @Override
-    public void cleanup(TaskSource taskSource, Schema schema, int taskCount, List<TaskReport> successTaskReports) {
+    public void cleanup(TaskSource taskSource, Schema schema, int taskCount, List<TaskReport> successTaskReports)
+    {
     }
 
     @Override
-    public TaskReport run(TaskSource taskSource, Schema schema, int taskIndex, PageOutput output) {
+    public TaskReport run(TaskSource taskSource, Schema schema, int taskIndex, PageOutput output)
+    {
         PluginTask task = taskSource.loadTask(PluginTask.class);
         BufferAllocator allocator = task.getBufferAllocator();
         PageBuilder pageBuilder = new PageBuilder(allocator, schema, output);
@@ -101,58 +128,93 @@ public class AthenaInputPlugin implements InputPlugin {
             connection = getAthenaConnection(task);
             statement = connection.createStatement();
             ResultSet resultSet = statement.executeQuery(task.getQuery());
+            boolean nullToZero = task.getNullToZero();
 
-            ResultSetMetaData m = resultSet.getMetaData();
             while (resultSet.next()) {
-                schema.visitColumns(new ColumnVisitor() {
+                schema.visitColumns(new ColumnVisitor()
+                {
                     @Override
-                    public void timestampColumn(Column column) {
+                    public void timestampColumn(Column column)
+                    {
                         try {
                             java.sql.Timestamp t = resultSet.getTimestamp(column.getName());
                             pageBuilder.setTimestamp(column, Timestamp.ofEpochMilli(t.getTime()));
-                        } catch (SQLException e) {
+                        }
+                        catch (SQLException e) {
                             e.printStackTrace();
+                            throw new RuntimeException(e);
                         }
                     }
 
                     @Override
-                    public void stringColumn(Column column) {
+                    public void stringColumn(Column column)
+                    {
                         try {
                             pageBuilder.setString(column, resultSet.getString(column.getName()));
-                        } catch (SQLException e) {
+                        }
+                        catch (SQLException e) {
                             e.printStackTrace();
+                            throw new RuntimeException(e);
                         }
                     }
 
                     @Override
-                    public void longColumn(Column column) {
+                    public void longColumn(Column column)
+                    {
                         try {
-                            pageBuilder.setLong(column, resultSet.getLong(column.getName()));
-                        } catch (SQLException e) {
+                            long ret = resultSet.getLong(column.getName());
+                            if (resultSet.wasNull() && !nullToZero){
+                                pageBuilder.setNull(column);
+                            }
+                            else {
+                                pageBuilder.setLong(column, ret);
+                            }
+                        }
+                        catch (SQLException e) {
                             e.printStackTrace();
+                            throw new RuntimeException(e);
                         }
                     }
 
                     @Override
-                    public void doubleColumn(Column column) {
+                    public void doubleColumn(Column column)
+                    {
                         try {
-                            pageBuilder.setDouble(column, resultSet.getDouble(column.getName()));
-                        } catch (SQLException e) {
+                            double ret = resultSet.getDouble(column.getName());
+                            if (resultSet.wasNull() && !nullToZero){
+                                pageBuilder.setNull(column);
+                            }
+                            else {
+                                pageBuilder.setDouble(column, ret);
+                            }
+                        }
+                        catch (SQLException e) {
                             e.printStackTrace();
+                            throw new RuntimeException(e);
                         }
                     }
 
                     @Override
-                    public void booleanColumn(Column column) {
+                    public void booleanColumn(Column column)
+                    {
                         try {
-                            pageBuilder.setBoolean(column, resultSet.getBoolean(column.getName()));
-                        } catch (SQLException e) {
+                            boolean ret = resultSet.getBoolean(column.getName());
+                            if (resultSet.wasNull() && !nullToZero){
+                                pageBuilder.setNull(column);
+                            }
+                            else {
+                                pageBuilder.setBoolean(column, ret);
+                            }
+                        }
+                        catch (SQLException e) {
                             e.printStackTrace();
+                            throw new RuntimeException(e);
                         }
                     }
 
                     @Override
-                    public void jsonColumn(Column column) {
+                    public void jsonColumn(Column column)
+                    {
                         // TODO:
                     }
                 });
@@ -164,19 +226,24 @@ public class AthenaInputPlugin implements InputPlugin {
             pageBuilder.close();
             resultSet.close();
             connection.close();
-        } catch (Exception e) {
+        }
+        catch (Exception e) {
             e.printStackTrace();
-        } finally {
+            throw new RuntimeException(e);
+        }
+        finally {
             try {
-                if (statement != null)
+                if (statement != null) {
                     statement.close();
-            } catch (Exception ex) {
-
+                }
             }
+            catch (Exception ex) { }
             try {
-                if (connection != null)
+                if (connection != null) {
                     connection.close();
-            } catch (Exception ex) {
+                }
+            }
+            catch (Exception ex) {
                 ex.printStackTrace();
             }
         }
@@ -185,12 +252,14 @@ public class AthenaInputPlugin implements InputPlugin {
     }
 
     @Override
-    public ConfigDiff guess(ConfigSource config) {
+    public ConfigDiff guess(ConfigSource config)
+    {
         return Exec.newConfigDiff();
     }
 
-    protected Connection getAthenaConnection(PluginTask task) throws ClassNotFoundException, SQLException {
-        Class.forName("com.amazonaws.athena.jdbc.AthenaDriver");
+    protected Connection getAthenaConnection(PluginTask task) throws ClassNotFoundException, SQLException
+    {
+        loadDriver("com.amazonaws.athena.jdbc.AthenaDriver", task.getDriverPath());
         Properties properties = new Properties();
         properties.put("s3_staging_dir", task.getS3StagingDir());
         properties.put("user", task.getAccessKey());
@@ -199,4 +268,84 @@ public class AthenaInputPlugin implements InputPlugin {
 
         return DriverManager.getConnection(task.getAthenaUrl(), properties);
     }
+
+    //
+    // copy from embulk-input-jdbc
+    //
+
+    protected void loadDriver(String className, Optional<String> driverPath)
+    {
+        if (driverPath.isPresent()) {
+            addDriverJarToClasspath(driverPath.get());
+        }
+        else {
+            try {
+                // Gradle test task will add JDBC driver to classpath
+                Class.forName(className);
+            }
+            catch (ClassNotFoundException ex) {
+                File root = findPluginRoot();
+                File driverLib = new File(root, "default_jdbc_driver");
+                File[] files = driverLib.listFiles(new FileFilter() {
+                    @Override
+                    public boolean accept(File file)
+                    {
+                        return file.isFile() && file.getName().endsWith(".jar");
+                    }
+                });
+                if (files == null || files.length == 0) {
+                    throw new RuntimeException("Cannot find JDBC driver in '" + root.getAbsolutePath() + "'.");
+                }
+                else {
+                    for (File file : files) {
+                        logger.info("JDBC Driver = " + file.getAbsolutePath());
+                        addDriverJarToClasspath(file.getAbsolutePath());
+                    }
+                }
+            }
+        }
+
+        // Load JDBC Driver
+        try {
+            Class.forName(className);
+        }
+        catch (ClassNotFoundException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+
+    protected void addDriverJarToClasspath(String glob)
+    {
+        // TODO match glob
+        PluginClassLoader loader = (PluginClassLoader) getClass().getClassLoader();
+        Path path = Paths.get(glob);
+        if (!path.toFile().exists()) {
+             throw new ConfigException("The specified driver jar doesn't exist: " + glob);
+        }
+        loader.addPath(Paths.get(glob));
+    }
+
+    protected File findPluginRoot()
+    {
+        try {
+            URL url = getClass().getResource("/" + getClass().getName().replace('.', '/') + ".class");
+            if (url.toString().startsWith("jar:")) {
+                url = new URL(url.toString().replaceAll("^jar:", "").replaceAll("![^!]*$", ""));
+            }
+
+            File folder = new File(url.toURI()).getParentFile();
+            for (;; folder = folder.getParentFile()) {
+                if (folder == null) {
+                    throw new RuntimeException("Cannot find 'embulk-input-xxx' folder.");
+                }
+
+                if (folder.getName().startsWith("embulk-input-")) {
+                    return folder;
+                }
+            }
+        }
+        catch (MalformedURLException | URISyntaxException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }