fluent-plugin-postgresql-csvlog 0.1.0 → 0.3.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 3f63769c3a4c78fc3db074b812e6f75e77c46b355563b403c3de620c49935392
-  data.tar.gz: 2fd8d1acd1fac5ce24b72ed4d488fc5f80f0a8752084784cf16f419389f34499
+  metadata.gz: 5d61fc31718e43c6d1dff46139a1b03d56384405d45b69c9772dc6f7b6a66dbf
+  data.tar.gz: 5d2d23a4a7b5f277b19f181aa515ab74718f6ef93824e2c958264d9bbdf9c9aa
 SHA512:
-  metadata.gz: 34fefa81b79223bf1840481f6ffd232e4d343a1eed9876f8f58893f686c8c4a1c1fbd37679c1a8960bcddef3a992733fe4b53158ee7558c10b49ac42a1793996
-  data.tar.gz: d105daf33a577b40e0035bc2a6d512390c361bb42bb9a11c2290e7489f04f813338779b45682969fa7616e8d7bb319db04dfec1eaf624586e50709c4a71eba83
+  metadata.gz: f2106f60749b6fa8fc931ccd3d85f51595b0ec60eb9f425935072b1749ed068e8a068c6d513f1557a20e8c5a0613acad80c6690c7ce2aa20cca3027d686c388a
+  data.tar.gz: 9ceef623cbd5256e047dea817d487d2711cc8b387ff068a10913100fe649e77174a600ac019b4729403390152c1a3e6f162accf4cc50f020561032b6493d965f

data/.gitlab-ci.yml

@@ -9,3 +9,21 @@ test:
   cache:
     paths:
       - vendor/ruby
+
+# integration tests
+itest:
+  services:
+    - name: postgres:12
+      alias: postgres
+      command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all"]
+  variables:
+    POSTGRES_USER: testuser
+    POSTGRES_PASSWORD: testpass
+  before_script:
+    - bundle config set path vendor
+    - bundle install --jobs $(nproc)
+  script:
+    - bundle exec rake itest
+  cache:
+    paths:
+      - vendor/ruby

data/Dockerfile

@@ -0,0 +1,18 @@
+# Dockerfile useful for manual testing purposes
+FROM fluent/fluentd:v1.13-1
+USER root
+RUN apk add bash alpine-sdk postgresql-dev postgresql-client ruby-dev
+
+WORKDIR /src/
+
+ADD Gemfile /src/
+ADD fluent-plugin-postgresql-csvlog.gemspec /src/
+
+ADD . /src/
+
+RUN bundle install
+RUN rake build
+
+RUN fluent-gem install pkg/*.gem
+
+ENTRYPOINT [ "/bin/bash"]

data/README.md

@@ -47,11 +47,13 @@ ingest and parse PostgreSQL CSV logs:
 
 <filter postgres.postgres_csv>
   @type postgresql_slowlog
+  output_key query
 </filter>
 
 <filter postgres.postgres_csv>
   @type postgresql_redactor
-  key sql
+  input_key query
+  output_key sql
   fingerprint_key fingerprint
 </filter>
 

data/Rakefile

@@ -9,4 +9,10 @@ Rake::TestTask.new(:test) do |test|
   test.verbose = true
 end
 
+Rake::TestTask.new(:itest) do |test|
+  test.libs << 'lib' << 'test'
+  test.test_files = FileList['test/**/itest_*.rb']
+  test.verbose = true
+end
+
 task :default => [:build]

data/docker-compose.yml

@@ -0,0 +1,19 @@
+# Docker Compose setup useful for testing and development purposes
+version: "3.9"
+services:
+  fluentd:
+    build: .
+    links:
+      - postgres
+    entrypoint: /usr/bin/fluentd -vvv -c /src/example-fluentd.conf
+  postgres:
+    image: postgres
+    restart: always
+    environment:
+      - POSTGRES_USER=testuser
+      - POSTGRES_PASSWORD=testpass
+    ports:
+      - '5438:5432'
+    command: postgres -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all
+    volumes:
+      - ./sql/create_extension.sql:/docker-entrypoint-initdb.d/create_extension.sql

data/example-fluentd.conf

@@ -0,0 +1,12 @@
+<source>
+  @type pg_stat_statements
+  tag postgres.pg_stat_statements
+  host postgres
+  username testuser
+  password testpass
+  interval 1
+</source>
+
+<match postgres.pg_stat_statements>
+  @type stdout
+</match>

data/fluent-plugin-postgresql-csvlog.gemspec

@@ -2,10 +2,10 @@ $:.push File.expand_path('lib', __dir__)
 
 Gem::Specification.new do |s|
   s.name        = 'fluent-plugin-postgresql-csvlog'
-  s.version     = '0.1.0'
+  s.version     = '0.3.2'
   s.authors     = ['stanhu']
   s.email       = ['stanhu@gmail.com']
-  s.homepage    = 'https://gitlab.com/gitlab-org/fluent-plugin-postgresql-csvlog'
+  s.homepage    = 'https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog'
   s.summary     = 'fluentd plugins to work with PostgreSQL CSV logs'
   s.description = 'fluentd plugins to work with PostgreSQL CSV logs'
 
@@ -15,6 +15,7 @@ Gem::Specification.new do |s|
   s.require_paths = ['lib']
 
   s.add_dependency 'fluentd', ['>= 1.0', '< 2']
+  s.add_dependency 'pg', '~> 1.1'
   s.add_dependency 'pg_query', '~> 2.0'
 
   s.add_development_dependency 'rake'

data/lib/fluent/plugin/filter_postgresql_redactor.rb

@@ -7,23 +7,25 @@ module Fluent::Plugin
   class PostgreSQLRedactor < Filter
     Fluent::Plugin.register_filter('postgresql_redactor', self)
 
-    desc 'Field to parse for SQL queries'
-    config_param :key, :string, default: 'sql'
+    desc 'Input field to parse for SQL queries'
+    config_param :input_key, :string, default: 'query'
+
+    desc 'Output field to store SQL queries'
+    config_param :output_key, :string, default: 'sql'
 
     desc 'Name of field to store SQL query fingerprint'
     config_param :fingerprint_key, :string, default: 'fingerprint'
 
     def filter(_tag, _time, record)
-      statement = record['statement']
+      statement = record[@input_key]
 
       return record unless statement
 
       normalized = PgQuery.normalize(statement)
       record[@fingerprint_key] = PgQuery.parse(normalized).fingerprint if @fingerprint_key
 
-      record.delete('statement')
-      record[@key] = normalized
-      record.delete('message')
+      record.delete(@input_key)
+      record[@output_key] = normalized
 
       record
     rescue PgQuery::ParseError

data/lib/fluent/plugin/filter_postgresql_slowlog.rb

@@ -12,6 +12,9 @@ module Fluent
     class PostgreSQLSlowLog < Filter
       Fluent::Plugin.register_filter('postgresql_slowlog', self)
 
+      desc 'Field to output SQL queries'
+      config_param :output_key, :string, default: 'query'
+
       SLOWLOG_REGEXP = /^duration: (\d+(?:\.\d+)?) ms .*?:\s*(.*)/m.freeze
 
       def filter(_tag, _time, record)
@@ -20,7 +23,8 @@ module Fluent
         # rubocop:disable Style/PerlBackrefs
         if record['message'] =~ SLOWLOG_REGEXP
           record['duration_s'] = $1.to_f / 1000.0
-          record['statement'] = $2
+          record[@output_key] = $2
+          record.delete('message')
         end
         # rubocop:enable Style/PerlBackrefs
 

data/lib/fluent/plugin/in_pg_stat_statements.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require 'fluent/plugin/input'
+require 'pg'
+require 'pg_query'
+
+module Fluent::Plugin
+  # PgStatStatementsInput will periodically poll postgres, querying pg_stat_statements
+  # for queryid to query mappings. These are then normalized for security purposes
+  # fingerprinted and emitted as records with the following format:
+  # {
+  #  'fingerprint' => '8a6e9896bd9048a2',
+  #  'query' => 'SELECT * FROM table ORDER BY queryid LIMIT $1',
+  #  'query_length' => 58,
+  #  'queryid' => 3239318621761098074
+  # }
+  class PgStatStatementsInput < Input
+    Fluent::Plugin.register_input('pg_stat_statements', self)
+
+    desc 'PostgreSQL host'
+    config_param :host, :string
+
+    desc 'RDBMS port (default: 5432)'
+    config_param :port, :integer, default: 5432
+
+    desc 'login user name'
+    config_param :username, :string, default: nil
+
+    desc 'postgres db'
+    config_param :dbname, :string, default: nil
+
+    desc 'login password'
+    config_param :password, :string, default: nil, secret: true
+
+    # See https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-SSLMODE
+    # for options
+    desc 'postgres sslmode'
+    config_param :sslmode, :string, default: 'prefer'
+
+    desc 'tag'
+    config_param :tag, :string, default: nil
+
+    desc 'interval in second to run query'
+    config_param :interval, :time, default: 300
+
+    desc 'Name of field to store SQL query fingerprint'
+    config_param :fingerprint_key, :string, default: 'fingerprint'
+
+    def start
+      @stop_flag = false
+      @thread = Thread.new(&method(:thread_main))
+    end
+
+    def shutdown
+      @stop_flag = true
+
+      # Interrupt thread and wait for it to finish
+      Thread.new { @thread.run } if @thread
+      @thread.join
+    end
+
+    def thread_main
+      until @stop_flag
+        sleep @interval
+        break if @stop_flag
+
+        begin
+          with_connection do |conn|
+            emit_statements_to_stream(conn)
+          end
+        rescue StandardError => e
+          log.error 'unexpected error', error: e.message, error_class: e.class
+          log.error_backtrace e.backtrace
+        end
+      end
+    end
+
+    # Returns a fluentd record for a query row
+    def record_for_row(row)
+      query = row['query']
+
+      # We record the query_length as it will help in understanding whether unparseable
+      # queries are truncated.
+      record = { 'queryid' => row['queryid'], 'query_length' => query&.length }
+
+      return record unless query
+
+      normalized = PgQuery.normalize(query)
+      record['query'] = normalized
+
+      record[@fingerprint_key] = PgQuery.parse(normalized).fingerprint if @fingerprint_key
+
+      record
+    rescue PgQuery::ParseError
+      record['query_unparseable'] = true
+
+      record
+    end
+
+    private
+
+    # Query the database and emit statements to fluentd router
+    def emit_statements_to_stream(conn)
+      me = Fluent::MultiEventStream.new
+
+      now = Fluent::Engine.now
+      conn.exec('SELECT queryid, query FROM public.pg_stat_statements').each do |row|
+        record = record_for_row(row)
+        me.add(now, record)
+      end
+
+      @router.emit_stream(@tag, me)
+    end
+
+    # Since this query is very infrequent, and it may be communicating directly
+    # with postgres without pgbouncer, don't use a persistent connection and
+    # ensure that it is properly closed
+    def with_connection(&block)
+      conn = PG.connect(
+        host: @host,
+        dbname: @dbname,
+        sslmode: @sslmode,
+        user: @username,
+        password: @password
+      )
+      conn.type_map_for_results = PG::BasicTypeMapForResults.new conn
+
+      begin
+        block.call(conn)
+      ensure
+        # Always close the connection
+        conn.finish
+      end
+    end
+  end
+end

data/sql/create_extension.sql

@@ -0,0 +1 @@
+CREATE EXTENSION pg_stat_statements;

data/test/helper.rb

@@ -7,10 +7,13 @@ $LOAD_PATH.unshift(File.join(__dir__, '..', 'lib'))
 $LOAD_PATH.unshift(__dir__)
 require 'fluent/test'
 require 'fluent/test/driver/filter'
+require 'fluent/test/driver/input'
 require 'fluent/test/helpers'
 
 Test::Unit::TestCase.include(Fluent::Test::Helpers)
+Test::Unit::TestCase.extend(Fluent::Test::Helpers)
 
 require 'fluent/plugin/filter_postgresql_slowlog'
 require 'fluent/plugin/filter_postgresql_redactor'
 require 'fluent/plugin/filter_marginalia'
+require 'fluent/plugin/in_pg_stat_statements'

data/test/plugin/itest_in_pg_stat_statements.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class PgStatStatementsInputIntegrationTest < Test::Unit::TestCase
+  # The defaults values work with the configuration in .gitlab-ci.yml on the postgres service
+  # Override with env vars for local development
+  HOST = ENV.fetch('PG_TEST_HOST', 'postgres')
+  USERNAME = ENV.fetch('PG_TEST_USER', 'testuser')
+  PASSWORD = ENV.fetch('PG_TEST_PASSWORD', 'testpass')
+
+  def setup
+    Fluent::Test.setup
+
+    @conn = PG.connect(
+      host: HOST,
+      user: USERNAME,
+      password: PASSWORD
+    )
+
+    try_setup_extension
+    create_known_statement
+  end
+
+  def teardown
+    @conn&.finish
+  end
+
+  # Setup pg_stat_statements extension
+  def try_setup_extension
+    @conn.exec('CREATE EXTENSION pg_stat_statements')
+  rescue PG::DuplicateObject
+  end
+
+  # This statement gives us something to look for in the emitted stream
+  def create_known_statement
+    @conn.exec('SELECT * FROM pg_stat_statements ORDER BY queryid LIMIT 1')
+  end
+
+  VALID_CONFIG = %(
+    tag postgres.pg_stat_statements
+    host #{HOST}
+    username #{USERNAME}
+    password #{PASSWORD}
+    interval 1
+  )
+
+  INVALID_CONFIG = %(
+    host 'invalid_host.dne'
+    port 1234
+    username #{USERNAME}
+    password #{PASSWORD}
+    interval 1
+  )
+
+  def create_driver(config)
+    Fluent::Test::InputTestDriver.new(Fluent::Plugin::PgStatStatementsInput).configure(config)
+  end
+
+  sub_test_case 'configuration' do
+    test 'connects' do
+      d = create_driver(VALID_CONFIG)
+
+      emits = []
+      # wait 50 * 0.05, "see fluentd/lib/fluent/test/base.rb:79 num_waits.times { sleep 0.05 }
+      d.run(num_waits = 50) do
+        emits = d.emits
+      end
+
+      assert_false emits.empty?
+    end
+
+    # Why do we have this test? If postgres is still starting up, we don't want to cause the
+    # the fluentd configuration to fail. We would rather retry until we get a connection
+    test 'connects for an invalid config' do
+      d = create_driver(INVALID_CONFIG)
+
+      emits = []
+      # wait 50 * 0.05, "see fluentd/lib/fluent/test/base.rb:79 num_waits.times { sleep 0.05 }
+      d.run(num_waits = 50) do
+        emits = d.emits
+      end
+
+      assert_true emits.empty?
+    end
+  end
+
+  sub_test_case 'execution' do
+    test 'connects' do
+      d = create_driver(VALID_CONFIG)
+
+      emits = []
+      # wait 50 * 0.05, "see fluentd/lib/fluent/test/base.rb:79 num_waits.times { sleep 0.05 }
+      d.run(num_waits = 50) do
+        emits = d.emits
+      end
+
+      expected_record = {
+        'fingerprint' => '8a6e9896bd9048a2',
+        'query' => 'SELECT * FROM pg_stat_statements ORDER BY queryid LIMIT $1',
+        'query_length' => 58,
+        'queryid' => 3_239_318_621_761_098_074
+      }
+      known_statement_event = emits.find do |event|
+        record = event[2]
+        record['query'] == expected_record['query']
+      end
+
+      assert_false known_statement_event.nil?
+
+      tag = known_statement_event[0]
+      record = known_statement_event[2]
+
+      assert_equal 'postgres.pg_stat_statements', tag
+      assert_equal expected_record['fingerprint'], record['fingerprint']
+      assert_equal expected_record['query_length'], record['query_length']
+      assert_true expected_record.include? 'queryid'
+    end
+  end
+end

data/test/plugin/test_filter_marginalia.rb

@@ -22,7 +22,8 @@ class Marginalia < Test::Unit::TestCase
     inputs = [
       { 'statement' => 'SELECT * FROM projects' },
       { 'statement' => 'SELECT COUNT(*) FROM "projects" /* this is just a comment */' },
-      { 'statement' => 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/' }
+      { 'statement' => 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/' },
+      { 'statement' => 'SELECT COUNT(*) FROM "projects" /*application:web,correlation_id:01F1D2T1SC9DM82A4865ATG1CP,endpoint_id:POST /api/:version/groups/:id/-/packages/mavenpath/:file_name*/' }
     ]
 
     d.run(default_tag: @tag) do
@@ -40,10 +41,19 @@ class Marginalia < Test::Unit::TestCase
                    }
                  ),
                  d.filtered[2].last)
+    assert_equal(inputs[3].merge(
+                   {
+                     'application' => 'web',
+                     'correlation_id' => '01F1D2T1SC9DM82A4865ATG1CP',
+                     'endpoint_id' => 'POST /api/:version/groups/:id/-/packages/mavenpath/:file_name'
+                   }
+                 ),
+                 d.filtered[3].last)
 
     assert_equal('SELECT * FROM projects', d.filtered[0].last['statement'])
     assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[1].last['statement'])
     assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[2].last['statement'])
+    assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[3].last['statement'])
   end
 
   test 'parses prepended Marginalia comments' do
@@ -52,6 +62,7 @@ class Marginalia < Test::Unit::TestCase
     inputs = [
       { 'statement' => '/* this is just a comment */ SELECT COUNT(*) FROM "projects"' },
       { 'statement' => '/*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/ SELECT COUNT(*) FROM "projects"' },
+      { 'statement' => '/*application:web,correlation_id:01F1D2T1SC9DM82A4865ATG1CP,endpoint_id:POST /api/:version/groups/:id/-/packages/mavenpath/:file_name*/ SELECT COUNT(*) FROM "projects"' },
       { 'statement' => '/*application:sidekiq*/ SELECT COUNT(*) FROM "projects"',
         'application' => 'test-conflict' }
     ]
@@ -72,10 +83,18 @@ class Marginalia < Test::Unit::TestCase
                  d.filtered[1].last)
     assert_equal(inputs[2].merge(
                    {
-                     'statement_application' => 'sidekiq'
+                     'application' => 'web',
+                     'correlation_id' => '01F1D2T1SC9DM82A4865ATG1CP',
+                     'endpoint_id' => 'POST /api/:version/groups/:id/-/packages/mavenpath/:file_name'
                    }
                  ),
                  d.filtered[2].last)
+    assert_equal(inputs[3].merge(
+                   {
+                     'statement_application' => 'sidekiq'
+                   }
+                 ),
+                 d.filtered[3].last)
 
     assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[0].last['statement'])
     assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[1].last['statement'])

data/test/plugin/test_filter_postgresql_redactor.rb

@@ -8,11 +8,7 @@ class PostgreSQLRedactorTest < Test::Unit::TestCase
     @tag = 'test.tag'
   end
 
-  CONFIG = '
-    <filter test.tag>
-      @type postgresql_redactor
-    </filter>
-  '
+  CONFIG = ''
 
   def create_driver(conf = CONFIG)
     Fluent::Test::Driver::Filter.new(Fluent::Plugin::PostgreSQLRedactor).configure(conf)
@@ -23,7 +19,7 @@ class PostgreSQLRedactorTest < Test::Unit::TestCase
 
     inputs = [
       { 'message' => 'duration: 2357.1 ms  execute <unnamed>: SELECT * FROM projects WHERE id = 1',
-        'statement' => %(SELECT * FROM projects WHERE id = 1),
+        'query' => %(SELECT * FROM projects WHERE id = 1),
         'duration_s' => 2.3571 }
     ]
 
@@ -31,20 +27,43 @@ class PostgreSQLRedactorTest < Test::Unit::TestCase
       inputs.each { |input| d.feed(input) }
     end
 
-    assert_equal(%w[duration_s fingerprint sql], d.filtered[0].last.keys.sort)
+    assert_equal(%w[duration_s fingerprint message sql], d.filtered[0].last.keys.sort)
     assert_equal('SELECT * FROM projects WHERE id = $1', d.filtered[0].last['sql'])
   end
 
   test 'handles parse errors' do
     d = create_driver
 
-    input = { 'statement' => 'create index something test (bla) include (bar)' }
+    input = { 'query' => 'create index something test (bla) include (bar)' }
 
     d.run(default_tag: @tag) do
       d.feed(input)
     end
 
-    assert_equal(%w[pg_query_error statement], d.filtered[0].last.keys.sort)
-    assert_equal(input['statement'], d.filtered[0].last['statement'])
+    assert_equal(%w[pg_query_error query], d.filtered[0].last.keys.sort)
+    assert_equal(input['query'], d.filtered[0].last['query'])
+  end
+
+  test 'uses configured input and output keys' do
+    d = create_driver(<<~CONF
+      input_key sql
+      output_key out_sql
+    CONF
+                     )
+
+    inputs = [
+      {
+        'message' => 'duration: 2357.1 ms  execute <unnamed>: SELECT * FROM projects WHERE id = 1',
+        'sql' => %(SELECT * FROM projects WHERE id = 1),
+        'duration_s' => 2.3571
+      }
+    ]
+
+    d.run(default_tag: @tag) do
+      inputs.each { |input| d.feed(input) }
+    end
+
+    assert_equal(%w[duration_s fingerprint message out_sql], d.filtered[0].last.keys.sort)
+    assert_equal('SELECT * FROM projects WHERE id = $1', d.filtered[0].last['out_sql'])
   end
 end

data/test/plugin/test_filter_postgresql_slowlog.rb

@@ -32,18 +32,21 @@ class PostgreSQLSlowLogTest < Test::Unit::TestCase
 
     assert_equal(inputs[0].merge(
                    {
-                     'statement' => 'SELECT * FROM projects',
+                     'query' => 'SELECT * FROM projects',
                      'duration_s' => 2.3571
                    }
                  ),
                  d.filtered[0].last)
     assert_equal(inputs[1].merge(
                    {
-                     'statement' => 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/',
+                     'query' => 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/',
                      'duration_s' => 1.873345
                    }
                  ),
                  d.filtered[1].last)
+
+    assert_equal(%w[duration_s query], d.filtered[0].last.keys.sort)
+    assert_equal(%w[duration_s query], d.filtered[1].last.keys.sort)
   end
 
   test 'ignores messages not having to do with slow logs' do
@@ -55,5 +58,32 @@ class PostgreSQLSlowLogTest < Test::Unit::TestCase
     end
 
     assert_equal(input, d.filtered[0].last)
+    assert_equal(%w[message], d.filtered[0].last.keys.sort)
+  end
+
+  test 'outputs slow log entries to configured output key' do
+    d = create_driver(
+      <<~CONF
+        output_key my_key
+      CONF
+    )
+
+    inputs = [
+      { 'message' => 'duration: 2357.1 ms  execute <unnamed>: SELECT * FROM projects' }
+    ]
+
+    d.run(default_tag: @tag) do
+      inputs.each { |input| d.feed(input) }
+    end
+
+    assert_equal(inputs[0].merge(
+                   {
+                     'my_key' => 'SELECT * FROM projects',
+                     'duration_s' => 2.3571
+                   }
+                 ),
+                 d.filtered[0].last)
+
+    assert_equal(%w[duration_s my_key], d.filtered[0].last.keys.sort)
   end
 end

data/test/plugin/test_in_pg_stat_statements.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class PgStatStatementsInputTest < Test::Unit::TestCase
+  def setup
+    Fluent::Test.setup
+  end
+
+  CONFIG = %(
+    tag postgres.pg_stat_statements
+    host localhost
+    port 1234
+    dbname gitlab
+    sslmode require
+    username moo
+    password secret
+    interval 600
+    fingerprint_key fingerprint
+  )
+
+  def create_driver
+    Fluent::Test::InputTestDriver.new(Fluent::Plugin::PgStatStatementsInput).configure(CONFIG)
+  end
+
+  sub_test_case 'configuration' do
+    test 'basic configuration' do
+      d = create_driver
+
+      assert_equal 'postgres.pg_stat_statements', d.instance.tag
+      assert_equal 'localhost', d.instance.host
+      assert_equal 1234, d.instance.port
+      assert_equal 'gitlab', d.instance.dbname
+      assert_equal 'require', d.instance.sslmode
+      assert_equal 'moo', d.instance.username
+      assert_equal 'secret', d.instance.password
+      assert_equal 600, d.instance.interval
+      assert_equal 'fingerprint', d.instance.fingerprint_key
+    end
+  end
+
+  sub_test_case 'execution' do
+    test 'sql' do
+      d = create_driver
+      record = d.instance.record_for_row({ 'queryid' => '1234', 'query' => 'SELECT * FROM users WHERE user_id = ?' })
+
+      expected = {
+        'fingerprint' => 'c071dee80d466e7d',
+        'query' => 'SELECT * FROM users WHERE user_id = ?',
+        'query_length' => 37,
+        'queryid' => '1234'
+      }
+
+      assert_equal expected, record
+    end
+
+    test 'nil query' do
+      d = create_driver
+      record = d.instance.record_for_row({ 'queryid' => '1234', 'query' => nil })
+
+      expected = { 'query_length' => nil, 'queryid' => '1234' }
+      assert_equal expected, record
+    end
+
+    test 'ddl query' do
+      d = create_driver
+      ddl_sql = <<-SQL
+        CREATE TABLE accounts (
+          user_id serial PRIMARY KEY,
+          username VARCHAR(50) UNIQUE NOT NULL,
+          password VARCHAR(50) NOT NULL,
+          email VARCHAR(255) UNIQUE NOT NULL,
+          created_on TIMESTAMP NOT NULL,
+          last_login TIMESTAMP
+        )
+      SQL
+
+      record = d.instance.record_for_row({ 'queryid' => '1234', 'query' => ddl_sql })
+
+      expected = {
+        'fingerprint' => 'fa9c9d26757c4f9b',
+        'query' => ddl_sql,
+        'query_length' => 287,
+        'queryid' => '1234'
+      }
+      assert_equal expected, record
+    end
+
+    test 'set command' do
+      d = create_driver
+      record = d.instance.record_for_row({ 'queryid' => '1234', 'query' => "SET TIME ZONE 'PST8PDT'" })
+
+      expected = {
+        'fingerprint' => '23f8d6eb1d3125c3',
+        'query' => 'SET TIME ZONE $1',
+        'query_length' => 23,
+        'queryid' => '1234'
+      }
+
+      assert_equal expected, record
+    end
+
+    test 'unparseable sql' do
+      d = create_driver
+      record = d.instance.record_for_row({ 'queryid' => '1234', 'query' => 'SELECT * FROM' })
+
+      expected = { 'query_length' => 13, 'query_unparseable' => true, 'queryid' => '1234' }
+      assert_equal expected, record
+    end
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: fluent-plugin-postgresql-csvlog
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.3.2
 platform: ruby
 authors:
 - stanhu
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-03-21 00:00:00.000000000 Z
+date: 2021-07-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: fluentd
@@ -30,6 +30,20 @@ dependencies:
     - - "<"
       - !ruby/object:Gem::Version
         version: '2'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.1'
 - !ruby/object:Gem::Dependency
   name: pg_query
   requirement: !ruby/object:Gem::Requirement
@@ -80,21 +94,27 @@ extensions: []
 extra_rdoc_files: []
 files:
 - ".gitlab-ci.yml"
+- Dockerfile
 - Gemfile
-- Gemfile.lock
 - LICENSE
 - README.md
 - Rakefile
+- docker-compose.yml
+- example-fluentd.conf
 - fluent-plugin-postgresql-csvlog.gemspec
 - lib/fluent/plugin/filter_marginalia.rb
 - lib/fluent/plugin/filter_postgresql_redactor.rb
 - lib/fluent/plugin/filter_postgresql_slowlog.rb
+- lib/fluent/plugin/in_pg_stat_statements.rb
 - lib/fluent/plugin/parser_multiline_csv.rb
+- sql/create_extension.sql
 - test/helper.rb
+- test/plugin/itest_in_pg_stat_statements.rb
 - test/plugin/test_filter_marginalia.rb
 - test/plugin/test_filter_postgresql_redactor.rb
 - test/plugin/test_filter_postgresql_slowlog.rb
-homepage: https://gitlab.com/gitlab-org/fluent-plugin-postgresql-csvlog
+- test/plugin/test_in_pg_stat_statements.rb
+homepage: https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog
 licenses: []
 metadata: {}
 post_install_message: 
@@ -118,6 +138,8 @@ specification_version: 4
 summary: fluentd plugins to work with PostgreSQL CSV logs
 test_files:
 - test/helper.rb
+- test/plugin/itest_in_pg_stat_statements.rb
 - test/plugin/test_filter_marginalia.rb
 - test/plugin/test_filter_postgresql_redactor.rb
 - test/plugin/test_filter_postgresql_slowlog.rb
+- test/plugin/test_in_pg_stat_statements.rb

data/Gemfile.lock

@@ -1,52 +0,0 @@
-PATH
-  remote: .
-  specs:
-    fluent-plugin-postgresql-csvlog (0.1.0)
-      fluentd (>= 1.0, < 2)
-      pg_query (~> 2.0)
-
-GEM
-  remote: https://rubygems.org/
-  specs:
-    concurrent-ruby (1.1.8)
-    cool.io (1.7.1)
-    fluentd (1.12.1)
-      bundler
-      cool.io (>= 1.4.5, < 2.0.0)
-      http_parser.rb (>= 0.5.1, < 0.7.0)
-      msgpack (>= 1.3.1, < 2.0.0)
-      serverengine (>= 2.2.2, < 3.0.0)
-      sigdump (~> 0.2.2)
-      strptime (>= 0.2.2, < 1.0.0)
-      tzinfo (>= 1.0, < 3.0)
-      tzinfo-data (~> 1.0)
-      yajl-ruby (~> 1.0)
-    google-protobuf (3.15.6-universal-darwin)
-    http_parser.rb (0.6.0)
-    msgpack (1.4.2)
-    pg_query (2.0.1)
-      google-protobuf (~> 3.15.5)
-    power_assert (2.0.0)
-    rake (13.0.3)
-    serverengine (2.2.3)
-      sigdump (~> 0.2.2)
-    sigdump (0.2.4)
-    strptime (0.2.5)
-    test-unit (3.4.0)
-      power_assert
-    tzinfo (2.0.4)
-      concurrent-ruby (~> 1.0)
-    tzinfo-data (1.2021.1)
-      tzinfo (>= 1.0.0)
-    yajl-ruby (1.4.1)
-
-PLATFORMS
-  ruby
-
-DEPENDENCIES
-  fluent-plugin-postgresql-csvlog!
-  rake
-  test-unit (~> 3.2)
-
-BUNDLED WITH
-   2.1.4