fluent-plugin-redshift 0.0.5 → 0.0.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: dd0f2cbfbaef5ac00d76a43d13bcdb2bcef77d83
-  data.tar.gz: 6071d3bb18d17fbf424fdced42febfc07adfe114
+  metadata.gz: ae21e38c3e70d5c36c0c52bcf3e38183756e4534
+  data.tar.gz: a2ebf5bc56f51b9c5f4dbd55fcf07499d9bc9bd8
 SHA512:
-  metadata.gz: 52dbbd84055f7e2eeba6fc33199544d247684179fb3a7ccf29400005e8b814f4146b1a075515384defdf011615de1c8f3140c1c4e456f6af7a7779a757056fdf
-  data.tar.gz: 34174f832c9ddcac43b96bd667cf872ea4403b453576f26921bb7fecf14bfb385c913ed34e74b115dca2184a306e14ae612e5d822a91cf0a3fdf1d1de020eb3c
+  metadata.gz: 064a44a6ed9086b1aef44eded321e984a9ccb8afc530e46b03ad99a77b2a2b384c24900ceea220f92c96871c1db18c9503910aff50ccc2b36d042e3d6077d80d
+  data.tar.gz: c85eb3832be2f660f8974b535e95430a51c2a13556a97d3f417eec816a05eb78c015ba079080847a46f48b3f6350cbddd3a06bae12464bcdcba4a26776e28390

data/VERSION

@@ -1 +1 @@
-0.0.5
+0.0.6

data/lib/fluent/plugin/out_redshift.rb

@@ -59,7 +59,8 @@ class RedshiftOutput < BufferedOutput
     }
     @delimiter = determine_delimiter(@file_type) if @delimiter.nil? or @delimiter.empty?
     $log.debug format_log("redshift file_type:#{@file_type} delimiter:'#{@delimiter}'")
-    @copy_sql_template = "copy #{table_name_with_schema} from '%s' CREDENTIALS 'aws_access_key_id=#{@aws_key_id};aws_secret_access_key=%s' delimiter '#{@delimiter}' GZIP ESCAPE #{@redshift_copy_base_options} #{@redshift_copy_options};"
+    @table_name_with_schema = [@redshift_schemaname, @redshift_tablename].compact.join('.')
+    @copy_sql_template = "copy #{@table_name_with_schema} from '%s' CREDENTIALS 'aws_access_key_id=#{@aws_key_id};aws_secret_access_key=%s' delimiter '#{@delimiter}' GZIP ESCAPE #{@redshift_copy_base_options} #{@redshift_copy_options};"
   end
 
   def start
@@ -72,6 +73,7 @@ class RedshiftOutput < BufferedOutput
     options[:s3_endpoint] = @s3_endpoint if @s3_endpoint
     @s3 = AWS::S3.new(options)
     @bucket = @s3.buckets[@s3_bucket]
+    @redshift_connection = RedshiftConnection.new(@db_conf)
   end
 
   def format(tag, time, record)
@@ -118,25 +120,26 @@ class RedshiftOutput < BufferedOutput
     $log.debug format_log("start copying. s3_uri=#{s3_uri}")
 
     begin
-      conn = PG.connect(@db_conf)
-      conn.exec(sql)
+      @redshift_connection.exec(sql)
       $log.info format_log("completed copying to redshift. s3_uri=#{s3_uri}")
-    rescue PG::Error => e
-      $log.error format_log("failed to copy data into redshift. s3_uri=#{s3_uri}"), :error=>e.to_s
-      raise e unless e.to_s =~ IGNORE_REDSHIFT_ERROR_REGEXP
-      return false # for debug
-    ensure
-      conn.close rescue nil if conn
+    rescue RedshiftError => e
+      if e.to_s =~ IGNORE_REDSHIFT_ERROR_REGEXP
+        $log.error format_log("failed to copy data into redshift due to load error. s3_uri=#{s3_uri}"), :error=>e.to_s
+        return false # for debug
+      end
+      raise e
     end
     true # for debug
   end
 
   protected
+
   def format_log(message)
     (@log_suffix and not @log_suffix.empty?) ? "#{message} #{@log_suffix}" : message
   end
 
   private
+
   def json?
     @file_type == 'json'
   end
@@ -158,11 +161,11 @@ class RedshiftOutput < BufferedOutput
 
   def create_gz_file_from_structured_data(dst_file, chunk, delimiter)
     # fetch the table definition from redshift
-    redshift_table_columns = fetch_table_columns
+    redshift_table_columns = @redshift_connection.fetch_table_columns(@redshift_tablename, @redshift_schemaname)
     if redshift_table_columns == nil
       raise "failed to fetch the redshift table definition."
     elsif redshift_table_columns.empty?
-      $log.warn format_log("no table on redshift. table_name=#{table_name_with_schema}")
+      $log.warn format_log("no table on redshift. table_name=#{@table_name_with_schema}")
       return nil
     end
 
@@ -204,27 +207,6 @@ class RedshiftOutput < BufferedOutput
     end
   end
 
-  def fetch_table_columns
-    begin
-      columns = nil
-      conn = PG.connect(@db_conf)
-      conn.exec(fetch_columns_sql_with_schema) do |result|
-        columns = result.collect{|row| row['column_name']}
-      end
-      columns
-    ensure
-      conn.close rescue nil if conn
-    end
-  end
-
-  def fetch_columns_sql_with_schema
-    @fetch_columns_sql ||= if @redshift_schemaname
-                             "select column_name from INFORMATION_SCHEMA.COLUMNS where table_schema = '#{@redshift_schemaname}' and table_name = '#{@redshift_tablename}' order by ordinal_position;"
-                           else
-                             "select column_name from INFORMATION_SCHEMA.COLUMNS where table_name = '#{@redshift_tablename}' order by ordinal_position;"
-                           end
-  end
-
   def json_to_hash(json_text)
     return nil if json_text.to_s.empty?
 
@@ -279,14 +261,108 @@ class RedshiftOutput < BufferedOutput
     s3path
   end
 
-  def table_name_with_schema
-    @table_name_with_schema ||= if @redshift_schemaname
-                                  "#{@redshift_schemaname}.#{@redshift_tablename}"
-                                else
-                                  @redshift_tablename
-                                end
+  class RedshiftError < StandardError
+    def initialize(msg)
+      case msg
+      when PG::Error
+        @pg_error = msg
+        super(msg.to_s)
+        set_backtrace(msg.backtrace)
+      else
+        super
+      end
+    end
+
+    attr_accessor :pg_error
+  end
+
+  class RedshiftConnection
+    REDSHIFT_CONNECT_TIMEOUT = 10.0  # 10sec
+
+    def initialize(db_conf)
+      @db_conf = db_conf
+      @connection = nil
+    end
+
+    attr_reader :db_conf
+
+    def fetch_table_columns(table_name, schema_name)
+      columns = nil
+      exec(fetch_columns_sql(table_name, schema_name)) do |result|
+        columns = result.collect{|row| row['column_name']}
+      end
+      columns
+    end
+
+    def exec(sql, &block)
+      conn = @connection
+      conn = create_redshift_connection if conn.nil?
+      if block
+        conn.exec(sql) {|result| block.call(result)}
+      else
+        conn.exec(sql)
+      end
+    rescue PG::Error => e
+      raise RedshiftError.new(e)
+    ensure
+      conn.close if conn && @connection.nil?
+    end
+
+    def connect_start
+      @connection = create_redshift_connection
+    end
+
+    def close
+      @connection.close rescue nil if @connection
+      @connection = nil
+    end
+
+    private
+
+    def create_redshift_connection
+      hostaddr = IPSocket.getaddress(db_conf[:host])
+      db_conf[:hostaddr] = hostaddr
+
+      conn = PG::Connection.connect_start(db_conf)
+      raise RedshiftError.new("Unable to create a new connection.") unless conn
+      if conn.status == PG::CONNECTION_BAD
+        raise RedshiftError.new("Connection failed: %s" % [ conn.error_message ])
+      end
+
+      socket = conn.socket_io
+      poll_status = PG::PGRES_POLLING_WRITING
+      until poll_status == PG::PGRES_POLLING_OK || poll_status == PG::PGRES_POLLING_FAILED
+        case poll_status
+        when PG::PGRES_POLLING_READING
+          IO.select([socket], nil, nil, REDSHIFT_CONNECT_TIMEOUT) or
+            raise RedshiftError.new("Asynchronous connection timed out!(READING)")
+        when PG::PGRES_POLLING_WRITING
+          IO.select(nil, [socket], nil, REDSHIFT_CONNECT_TIMEOUT) or
+            raise RedshiftError.new("Asynchronous connection timed out!(WRITING)")
+        end
+        poll_status = conn.connect_poll
+      end
+
+      unless conn.status == PG::CONNECTION_OK
+        raise RedshiftError, ("Connect failed: %s" % [conn.error_message.to_s.lines.uniq.join(" ")])
+      end
+
+      conn
+    rescue => e
+      conn.close rescue nil if conn
+      raise RedshiftError.new(e) if e.kind_of?(PG::Error)
+      raise e
+    end
+
+    def fetch_columns_sql(table_name, schema_name = nil)
+      sql = "select column_name from INFORMATION_SCHEMA.COLUMNS where table_name = '#{table_name}'"
+      sql << " and table_schema = '#{schema_name}'" if schema_name
+      sql << " order by ordinal_position;"
+      sql
+    end
   end
 end
 
 
+
 end

data/test/plugin/test_out_redshift.rb

@@ -12,6 +12,7 @@ class RedshiftOutputTest < Test::Unit::TestCase
     require 'pg'
     require 'csv'
     Fluent::Test.setup
+    PG::Error.module_eval { attr_accessor :result}
   end
 
   CONFIG_BASE= %[
@@ -211,54 +212,45 @@ class RedshiftOutputTest < Test::Unit::TestCase
     d_msgpack.run
   end
 
-  class PGConnectionMock
-    def initialize(options = {})
-      @return_keys = options[:return_keys] || ['key_a', 'key_b', 'key_c', 'key_d', 'key_e', 'key_f', 'key_g', 'key_h']
-      @target_schema = options[:schemaname] || nil
-      @target_table = options[:tablename] || 'test_table'
-    end
+  def setup_redshift_connection_mock(options = {})
+    options ||= {}
+    column_names = options[:column_names] || ['key_a', 'key_b', 'key_c', 'key_d', 'key_e', 'key_f', 'key_g', 'key_h']
+    schema_name = options[:schema_name]
+    table_name = options[:table_name] || 'test_table'
+    exec_sql_proc = options[:exec_sql_proc]
 
-    def expected_column_list_query
-      if @target_schema
-        /\Aselect column_name from INFORMATION_SCHEMA.COLUMNS where table_schema = '#{@target_schema}' and table_name = '#{@target_table}'/
+    column_list_query_regex =
+      if schema_name
+        /\Aselect column_name from INFORMATION_SCHEMA.COLUMNS where table_name = '#{table_name}' and table_schema = '#{schema_name}'/
       else
-        /\Aselect column_name from INFORMATION_SCHEMA.COLUMNS where table_name = '#{@target_table}'/
+        /\Aselect column_name from INFORMATION_SCHEMA.COLUMNS where table_name = '#{table_name}'/
       end
-    end
-
-    def expected_copy_query
-      if @target_schema
-        /\Acopy #{@target_schema}.#{@target_table} from/
+    copy_query_regex =
+      if schema_name
+        /\Acopy #{schema_name}.#{table_name} from/
       else
-        /\Acopy #{@target_table} from/
+        /\Acopy #{table_name} from/
       end
-    end
 
-    def exec(sql, &block)
-      if block_given?
-        if sql =~ expected_column_list_query
-          yield @return_keys.collect{|key| {'column_name' => key}}
+    flexmock(Fluent::RedshiftOutput::RedshiftConnection).new_instances do |conn|
+      conn.should_receive(:exec).and_return do |sql, block|
+        if exec_sql_proc
+          exec_sql_proc.call(sql, block)
+        elsif block
+          if sql =~ column_list_query_regex
+            block.call column_names.collect{|key| {'column_name' => key}}
+          else
+            block.call []
+          end
         else
-          yield []
-        end
-      else
-        unless sql =~ expected_copy_query
-          error = PG::Error.new("ERROR:  Load into table '#{@target_table}' failed.  Check 'stl_load_errors' system table for details.")
-          error.result = "ERROR:  Load into table '#{@target_table}' failed.  Check 'stl_load_errors' system table for details."
-          raise error
+          unless sql =~ copy_query_regex
+            error = PG::Error.new("ERROR:  Load into table '#{@target_table}' failed.  Check 'stl_load_errors' system table for details.")
+            error.result = "ERROR:  Load into table '#{@target_table}' failed.  Check 'stl_load_errors' system table for details."
+            raise Fluent::RedshiftOutput::RedshiftError.new(error)
+          end
         end
       end
     end
-
-    def close
-    end
-  end
-
-  def setup_pg_mock
-    # create mock of PG
-    def PG.connect(dbinfo)
-      return PGConnectionMock.new
-    end
   end
 
   def setup_s3_mock(expected_data)
@@ -305,9 +297,10 @@ class RedshiftOutputTest < Test::Unit::TestCase
     flexmock(Tempfile).new_instances.should_receive(:close!).at_least.once
   end
 
-  def setup_mocks(expected_data)
-    setup_pg_mock
-    setup_s3_mock(expected_data) end
+  def setup_mocks(expected_data, options = {})
+    setup_redshift_connection_mock(options)
+    setup_s3_mock(expected_data)
+  end
 
   def test_write_with_csv
     setup_mocks(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n])
@@ -421,38 +414,21 @@ class RedshiftOutputTest < Test::Unit::TestCase
   end
 
   def test_write_redshift_connection_error
-    def PG.connect(dbinfo)
-      return Class.new do
-        def initialize(return_keys=[]); end
-        def exec(sql)
-          raise PG::Error, "redshift connection error"
-        end
-        def close; end
-      end.new
-    end
-    setup_s3_mock(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n])
-
+    setup_mocks(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n],
+      exec_sql_proc: Proc.new {|sql, block| raise Fluent::RedshiftOutput::RedshiftError, "redshift connection error" })
     d_csv = create_driver
     emit_csv(d_csv)
-    assert_raise(PG::Error) {
+    assert_raise(Fluent::RedshiftOutput::RedshiftError) {
       d_csv.run
     }
   end
 
   def test_write_redshift_load_error
-    PG::Error.module_eval { attr_accessor :result}
-    def PG.connect(dbinfo)
-      return Class.new do
-        def initialize(return_keys=[]); end
-        def exec(sql)
-          error = PG::Error.new("ERROR:  Load into table 'apache_log' failed.  Check 'stl_load_errors' system table for details.")
-          error.result = "ERROR:  Load into table 'apache_log' failed.  Check 'stl_load_errors' system table for details."
-          raise error
-        end
-        def close; end
-      end.new
-    end
-    setup_s3_mock(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n])
+    setup_mocks(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n],
+      exec_sql_proc: Proc.new {|sql, block|
+      msg = "ERROR:  Load into table 'apache_log' failed.  Check 'stl_load_errors' system table for details."
+      raise Fluent::RedshiftOutput::RedshiftError.new(msg)
+    })
 
     d_csv = create_driver
     emit_csv(d_csv)
@@ -460,36 +436,19 @@ class RedshiftOutputTest < Test::Unit::TestCase
   end
 
   def test_write_with_json_redshift_connection_error
-    def PG.connect(dbinfo)
-      return Class.new do
-        def initialize(return_keys=[]); end
-        def exec(sql, &block)
-          error = PG::Error.new("redshift connection error")
-          raise error
-        end
-        def close; end
-      end.new
-    end
-    setup_s3_mock(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n])
+    setup_mocks(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n],
+      exec_sql_proc: Proc.new {|sql, block| raise Fluent::RedshiftOutput::RedshiftError.new("redshift connection error")})
 
     d_json = create_driver(CONFIG_JSON)
     emit_json(d_json)
-    assert_raise(PG::Error) {
+    assert_raise(Fluent::RedshiftOutput::RedshiftError) {
       d_json.run
     }
   end
 
   def test_write_with_json_no_table_on_redshift
-    def PG.connect(dbinfo)
-      return Class.new do
-        def initialize(return_keys=[]); end
-        def exec(sql, &block)
-          yield [] if block_given?
-        end
-        def close; end
-      end.new
-    end
-    setup_s3_mock(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n])
+    setup_mocks(%[val_a,val_b,val_c,val_d\nval_e,val_f,val_g,val_h\n],
+                exec_sql_proc: Proc.new {|sql, block| block.call [] if block })
 
     d_json = create_driver(CONFIG_JSON)
     emit_json(d_json)
@@ -497,15 +456,7 @@ class RedshiftOutputTest < Test::Unit::TestCase
   end
 
   def test_write_with_json_failed_to_get_columns
-    def PG.connect(dbinfo)
-      return Class.new do
-        def initialize(return_keys=[]); end
-        def exec(sql, &block)
-        end
-        def close; end
-      end.new
-    end
-    setup_s3_mock("")
+    setup_mocks("", exec_sql_proc: Proc.new {|sql, block| nil})
 
     d_json = create_driver(CONFIG_JSON)
     emit_json(d_json)
@@ -515,10 +466,8 @@ class RedshiftOutputTest < Test::Unit::TestCase
   end
 
   def test_write_with_json_fetch_column_with_schema
-    def PG.connect(dbinfo)
-      return PGConnectionMock.new(:schemaname => 'test_schema')
-    end
-    setup_s3_mock(%[val_a\tval_b\t\t\t\t\t\t\n\t\tval_c\tval_d\t\t\t\t\n])
+    setup_mocks(%[val_a\tval_b\t\t\t\t\t\t\n\t\tval_c\tval_d\t\t\t\t\n],
+               schema_name: 'test_schema')
     d_json = create_driver(CONFIG_JSON_WITH_SCHEMA)
     emit_json(d_json)
     assert_equal true, d_json.run

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: fluent-plugin-redshift
 version: !ruby/object:Gem::Version
-  version: 0.0.5
+  version: 0.0.6
 platform: ruby
 authors:
 - Masashi Miyazaki
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2015-05-27 00:00:00.000000000 Z
+date: 2015-08-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: fluentd