From 39cfce307d972bdafae81763e0d3ba9bd2008910 Mon Sep 17 00:00:00 2001
From: Brandon Max <bmax1337@gmail.com>
Date: Sat, 25 Apr 2020 01:43:46 -0400
Subject: [PATCH] Fix jdbc_fetch_size with postgresql

---
 lib/logstash/plugin_mixins/jdbc/jdbc.rb        |  2 ++
 .../plugin_mixins/jdbc/statement_handler.rb    | 18 +++++++++++-------
 spec/inputs/jdbc_spec.rb                       |  3 +--
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/lib/logstash/plugin_mixins/jdbc/jdbc.rb b/lib/logstash/plugin_mixins/jdbc/jdbc.rb
index 750a699..d29cd5f 100644
--- a/lib/logstash/plugin_mixins/jdbc/jdbc.rb
+++ b/lib/logstash/plugin_mixins/jdbc/jdbc.rb
@@ -144,6 +144,7 @@ def load_driver
         require "java"
         require "sequel"
         require "sequel/adapters/jdbc"
+        require "sequel/adapters/jdbc/transactions"
 
         load_driver_jars
         begin
@@ -193,6 +194,7 @@ def open_jdbc_connection
 
       @database = jdbc_connect()
       @database.extension(:pagination)
+      @database.extend(Sequel::JDBC::Transactions)
       if @jdbc_default_timezone
         @database.extension(:named_timezones)
         @database.timezone = @jdbc_default_timezone
diff --git a/lib/logstash/plugin_mixins/jdbc/statement_handler.rb b/lib/logstash/plugin_mixins/jdbc/statement_handler.rb
index 7b146a6..65e5dcd 100644
--- a/lib/logstash/plugin_mixins/jdbc/statement_handler.rb
+++ b/lib/logstash/plugin_mixins/jdbc/statement_handler.rb
@@ -31,16 +31,20 @@ class NormalStatementHandler < StatementHandler
     # @yieldparam row [Hash{Symbol=>Object}]
     def perform_query(db, sql_last_value, jdbc_paging_enabled, jdbc_page_size)
       query = build_query(db, sql_last_value)
-      if jdbc_paging_enabled
-        query.each_page(jdbc_page_size) do |paged_dataset|
-          paged_dataset.each do |row|
+      # Execute query in transaction cause PG driver require autocommit off for set fetch count
+      # See: https://jdbc.postgresql.org/documentation/head/query.html
+      db.transaction(rollback: :always) do
+        if jdbc_paging_enabled
+          query.each_page(jdbc_page_size) do |paged_dataset|
+            paged_dataset.each do |row|
+              yield row
+            end
+          end
+        else
+          query.each do |row|
             yield row
           end
         end
-      else
-        query.each do |row|
-          yield row
-        end
       end
     end
 
diff --git a/spec/inputs/jdbc_spec.rb b/spec/inputs/jdbc_spec.rb
index 919cc70..ad71309 100755
--- a/spec/inputs/jdbc_spec.rb
+++ b/spec/inputs/jdbc_spec.rb
@@ -1020,7 +1020,6 @@
       {
         "statement" => "SELECT * FROM test_table",
         "jdbc_pool_timeout" => 0,
-        "jdbc_connection_string" => 'mock://localhost:1527/db',
         "sequel_opts" => {
           "max_connections" => 1
         }
@@ -1076,7 +1075,7 @@
     end
 
     it "should report the statements to logging" do
-      expect(plugin.logger).to receive(:debug).once
+      expect(plugin.logger).to receive(:debug).thrice
       plugin.run(queue)
     end
   end