fluent-plugin-postgresql-csvlog 0.0.2 → 0.3.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 002b7169e3f3ac6493eb09dca5c2e3820a944f38c64a0b45641f3cdbd1717ddf
-  data.tar.gz: eac3646a2404665924c4fa30114b91ecce49dbadef3919d89eab458450568328
+  metadata.gz: 81cfb53854390b0aae43fd6f8cfaf05e11c4251a4442bf40ccdee016db11b3f6
+  data.tar.gz: 0e705c2af91af84e400e96d41b5275f0d01f4d6dc4ec0a37b3bfed0974e8a424
 SHA512:
-  metadata.gz: 354478f3573f0934dcee72305069aa8f25333087d1667fc4e127963b8f8f955b66a02462925e4672f27f80e43b772640a68a4760266f461704bd9486ad66a3ef
-  data.tar.gz: f6d89c2db73d337b1aa52e9838ab08f6f3bda80e7053307482fa5b284b4b6f81ae0389e2a20411fbcc383c878b5309a320946d5849d43fa466e469037b9d5102
+  metadata.gz: 24bb505d30f06847e16b2f631e1e2ce011559edcc11724556d62abc0089ee99b00e5b911c1f23d87740720fd4d57235af1ee104991cbff72ac8ba790c700280a
+  data.tar.gz: d42ae8976276b4392c1963d6c83e7691732f7eec8dbf87eaf270e3b4547ceba8f9f2b76056006163d6b594e0791b4a9b49183c6ae1793610e7f6555a92b9d190

data/.gitlab-ci.yml

@@ -2,9 +2,28 @@ image: "ruby:2.7"
 
 test:
   before_script:
+    - bundle config set path vendor
     - bundle install --jobs $(nproc)
   script:
     - bundle exec rake 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,15 +47,20 @@ ingest and parse PostgreSQL CSV logs:
 
 <filter postgres.postgres_csv>
   @type postgresql_slowlog
+  output_key query
 </filter>
 
 <filter postgres.postgres_csv>
   @type postgresql_redactor
+  input_key query
+  output_key sql
+  fingerprint_key fingerprint
 </filter>
 
 <filter postgres.postgres_csv>
   @type marginalia
   key sql
+  strip_comment true
 </filter>
 
 # Output resulting JSON file to a directory in /tmp

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.0.2'
+  s.version     = '0.3.1'
   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,7 +15,8 @@ Gem::Specification.new do |s|
   s.require_paths = ['lib']
 
   s.add_dependency 'fluentd', ['>= 1.0', '< 2']
-  s.add_dependency 'pg_query', '~> 1.3'
+  s.add_dependency 'pg', '~> 1.1'
+  s.add_dependency 'pg_query', '~> 2.0'
 
   s.add_development_dependency 'rake'
   s.add_development_dependency 'test-unit', '~> 3.2'

data/lib/fluent/plugin/filter_marginalia.rb

@@ -16,9 +16,11 @@ module Fluent
       desc 'Field to parse for Marginalia comments (key1:value1,key2:value2)'
       config_param :key, :string, default: 'sql'
 
-      MARGINALIA_PREPENDED_REGEXP = %r{^(?<comment>/\*.*\*/).*}m.freeze
-      MARGINALIA_APPENDED_REGEXP = %r{.*(?<comment>/\*.*\*/)$}m.freeze
-      MARGINALIA_KEY_VALUE_REGEXP = /^(?<key>.*):?(?<value>.*)$/.freeze
+      desc 'Whether to strip the comment from the record specified by key'
+      config_param :strip_comment, :bool, default: true
+
+      MARGINALIA_PREPENDED_REGEXP = %r{^(?<comment>/\*.*\*/)(?<sql>.*)}m.freeze
+      MARGINALIA_APPENDED_REGEXP = %r{(?<sql>.*)(?<comment>/\*.*\*/)$}m.freeze
 
       def filter(_tag, _time, record)
         parse_comments(record)
@@ -38,15 +40,9 @@ module Fluent
         return unless comment_match
 
         entries = extract_entries(comment_match['comment'])
+        parse_entries(entries, record)
 
-        entries.each do |component|
-          data = component.split(':', 2)
-
-          break unless data.length == 2
-
-          stored_key = store_key(record, data[0])
-          record[stored_key] = data[1]
-        end
+        record[@key] = comment_match['sql'].strip if @strip_comment
       end
 
       def match_marginalia_comment(sql)
@@ -73,7 +69,21 @@ module Fluent
         comment.gsub!(%r{\*/$}, '')
       end
 
+      def parse_entries(entries, record)
+        entries.each do |component|
+          data = component.split(':', 2)
+
+          break unless data.length == 2
+
+          stored_key = store_key(record, data[0])
+          record[stored_key] = data[1]
+        end
+      end
+
       def store_key(record, component_key)
+        # In case there is a conflict with the Marginalia key
+        # (e.g. `correlation_id`), we use the base key
+        # (`sql_correlation_id`) instead.
         if record.key?(component_key)
           "#{@key}_#{component_key}"
         else

data/lib/fluent/plugin/filter_postgresql_redactor.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'fluent/plugin/filter'
 require 'pg_query'
 
@@ -5,24 +7,28 @@ module Fluent::Plugin
   class PostgreSQLRedactor < Filter
     Fluent::Plugin.register_filter('postgresql_redactor', self)
 
-    def configure(conf)
-      super
-    end
+    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.delete('statement')
-      record['sql'] = normalized
-      record.delete('message')
+      record[@fingerprint_key] = PgQuery.parse(normalized).fingerprint if @fingerprint_key
+
+      record.delete(@input_key)
+      record[@output_key] = normalized
 
       record
     rescue PgQuery::ParseError
-      # pg_query currently only supports PostgresQL 10:
-      # https://github.com/lfittl/pg_query/issues/184
       record['pg_query_error'] = true
       record
     end

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 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,15 +22,16 @@ 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
       inputs.each { |input| d.feed(input) }
     end
 
-    assert_equal(inputs[0].merge, d.filtered[0].last)
-    assert_equal(inputs[1].merge, d.filtered[1].last)
+    assert_equal(inputs[0], d.filtered[0].last)
+    assert_equal(inputs[1], d.filtered[1].last)
     assert_equal(inputs[2].merge(
                    {
                      'application' => 'sidekiq',
@@ -40,6 +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
@@ -48,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' }
     ]
@@ -56,7 +71,7 @@ class Marginalia < Test::Unit::TestCase
       inputs.each { |input| d.feed(input) }
     end
 
-    assert_equal(inputs[0].merge, d.filtered[0].last)
+    assert_equal(inputs[0], d.filtered[0].last)
     assert_equal(inputs[1].merge(
                    {
                      'application' => 'sidekiq',
@@ -68,9 +83,45 @@ 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'])
+    assert_equal('SELECT COUNT(*) FROM "projects"', d.filtered[2].last['statement'])
+  end
+
+  test 'parses Marginalia comments with strip_comment disabled' do
+    d = create_driver(
+      <<~CONF
+        strip_comment false
+        key sql
+      CONF
+    )
+
+    sql = %(SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/)
+    appended_sql = %(SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/')
+
+    inputs = [
+      { 'sql' => sql },
+      { 'sql' => appended_sql }
+    ]
+
+    d.run(default_tag: @tag) do
+      inputs.each { |input| d.feed(input) }
+    end
+
+    assert_equal(sql, d.filtered[0].last['sql'])
+    assert_equal(appended_sql, d.filtered[1].last['sql'])
   end
 end

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 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 concurrently foo on 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.0.2
+  version: 0.3.1
 platform: ruby
 authors:
 - stanhu
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-02-20 00:00:00.000000000 Z
+date: 2021-06-28 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: fluentd
@@ -30,20 +30,34 @@ 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
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: '2.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.3'
+        version: '2.0'
 - !ruby/object:Gem::Dependency
   name: rake
   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,50 +0,0 @@
-PATH
-  remote: .
-  specs:
-    fluent-plugin-postgresql-csvlog (0.0.2)
-      fluentd (>= 1.0, < 2)
-      pg_query (~> 1.3)
-
-GEM
-  remote: https://rubygems.org/
-  specs:
-    concurrent-ruby (1.1.8)
-    cool.io (1.7.0)
-    fluentd (1.12.0)
-      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)
-    http_parser.rb (0.6.0)
-    msgpack (1.4.2)
-    pg_query (1.3.0)
-    power_assert (2.0.0)
-    rake (13.0.3)
-    serverengine (2.2.2)
-      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