fluent-plugin-postgresql-csvlog 0.0.1 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9fd4d916e4d8cc9ebb39d62e847de1c2dca1ce0175d738ca55430a962b3283ab
-  data.tar.gz: b3e6baacf5bd683c39eca587c55ea0326ffe18a55d3b4f8b23ba435642641b8c
+  metadata.gz: bb322f50f848a7b196962dae916a75f30a77f1461167bb314494ae203f8ef9d6
+  data.tar.gz: 077266e6c7f3589813d0c4db8ddccfde76f7d6ecd0b6d6388b67a4e6328ad858
 SHA512:
-  metadata.gz: 2e783b77515eb3ec55b4684e11d3014fdea761b057d22fb3935eafd7064d9fcb7038b8f78f7fa7e9aa238d339c8a5c3d1cac358a430c207f6ba9cd19bf8ae55a
-  data.tar.gz: 77dace38e2de7a851547ee25502943f0d610cceefa8ee4d1194e084301852bf3b1e8fd93e79144f6c81865d7aadfe3c3705a210f3ca4dee2029e2daa5fbdf989
+  metadata.gz: 9b130aa84b1285a62466b3bf68ff6d02acc7e95cd53cc91ee7e34ec81ee4440e06358738790f86f8d436cb8a72ad232f12ac2c0f9d8aca16fd5c6ba19bbdfe49
+  data.tar.gz: e9480110be0cadfad86c5fcc554eb5d1d7264559714e4fd88ea396ae9bb913027b15782c70cd2f4c06e1f4cc5a5bbafe55632ca8eba02ce57e610d98f8440c09

data/.gitlab-ci.yml

@@ -0,0 +1,29 @@
+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/README.md

@@ -6,6 +6,7 @@ parse PostgreSQL CSV log files and extract slow log information:
 - `MultilineCSVParser`: Parses CSV files that span multiple lines
 - `PostgreSQLSlowLog`: Extracts slow log entries into `duration_s` and `statement` fields
 - `PostgreSQLRedactor`: Normalizes the SQL query and redacts sensitive information
+- `Marginalia`: Parses [Marginalia comments](https://github.com/basecamp/marginalia) into key-value pairs and stores them
 
 ## Installation
 
@@ -28,7 +29,6 @@ The configuration below shows how you might use these filters to
 ingest and parse PostgreSQL CSV logs:
 
 ```conf
-## PostgreSQL csvlog (enabled with 
 <source>
   @type tail
   tag postgres.postgres_csv
@@ -47,10 +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

@@ -0,0 +1,18 @@
+require 'bundler'
+Bundler::GemHelper.install_tasks
+
+require 'rake/testtask'
+
+Rake::TestTask.new(:test) do |test|
+  test.libs << 'lib' << 'test'
+  test.test_files = FileList['test/**/test_*.rb']
+  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/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.1'
+  s.version     = '0.3.0'
   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

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require 'fluent/plugin/filter'
+
+module Fluent
+  module Plugin
+    # Filters SQL statements for Marginalia comments.
+    #
+    # Examples:
+    # SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/
+    # /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/ SELECT COUNT(*) FROM "projects"
+    #
+    class Marginalia < Filter
+      Fluent::Plugin.register_filter('marginalia', self)
+
+      desc 'Field to parse for Marginalia comments (key1:value1,key2:value2)'
+      config_param :key, :string, default: 'sql'
+
+      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)
+
+        record
+      end
+
+      private
+
+      def parse_comments(record)
+        sql = record[@key]
+
+        return unless sql
+
+        comment_match = match_marginalia_comment(sql)
+
+        return unless comment_match
+
+        entries = extract_entries(comment_match['comment'])
+        parse_entries(entries, record)
+
+        record[@key] = comment_match['sql'].strip if @strip_comment
+      end
+
+      def match_marginalia_comment(sql)
+        matched = MARGINALIA_PREPENDED_REGEXP.match(sql)
+
+        return matched if matched
+
+        MARGINALIA_APPENDED_REGEXP.match(sql)
+      end
+
+      def extract_entries(comment)
+        comment = scrub_comment(comment)
+
+        return [] unless comment
+
+        comment.split(',')
+      end
+
+      def scrub_comment(comment)
+        return unless comment
+
+        comment.strip!
+        comment.gsub!(%r{^/\*}, '')
+        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
+          component_key
+        end
+      end
+    end
+  end
+end

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/input_pg_stat_statements.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require 'fluent/input'
+require 'pg'
+require 'pg_query'
+
+module Fluent
+  # 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 = MultiEventStream.new
+
+      now = 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/test/helper.rb

@@ -7,9 +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/input_pg_stat_statements'

data/test/plugin/itest_input_pg_stat_statements.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class PgStatStatementsInputIntegrationTest < Test::Unit::TestCase
+  # These items are configured in .gitlab-ci.yml on the postgres service
+  HOST = 'postgres'
+  USERNAME = 'testuser'
+  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::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

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class Marginalia < Test::Unit::TestCase
+  def setup
+    Fluent::Test.setup
+    @tag = 'test.tag'
+  end
+
+  CONFIG = '
+    key statement
+  '
+
+  def create_driver(conf = CONFIG)
+    Fluent::Test::Driver::Filter.new(Fluent::Plugin::Marginalia).configure(conf)
+  end
+
+  test 'parses appended Marginalia comments' do
+    d = create_driver
+
+    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: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], d.filtered[0].last)
+    assert_equal(inputs[1], d.filtered[1].last)
+    assert_equal(inputs[2].merge(
+                   {
+                     'application' => 'sidekiq',
+                     'correlation_id' => 'd67cae54c169e0cab7d73389e2934f0e',
+                     'jid' => '52a1c8a9e4c555ea573f20f0',
+                     'job_class' => 'Geo::MetricsUpdateWorker'
+                   }
+                 ),
+                 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
+    d = create_driver
+
+    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' }
+    ]
+
+    d.run(default_tag: @tag) do
+      inputs.each { |input| d.feed(input) }
+    end
+
+    assert_equal(inputs[0], d.filtered[0].last)
+    assert_equal(inputs[1].merge(
+                   {
+                     'application' => 'sidekiq',
+                     'correlation_id' => 'd67cae54c169e0cab7d73389e2934f0e',
+                     'jid' => '52a1c8a9e4c555ea573f20f0',
+                     'job_class' => 'Geo::MetricsUpdateWorker'
+                   }
+                 ),
+                 d.filtered[1].last)
+    assert_equal(inputs[2].merge(
+                   {
+                     '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_input_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::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.1
+  version: 0.3.0
 platform: ruby
 authors:
 - stanhu
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-02-10 00:00:00.000000000 Z
+date: 2021-06-24 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
@@ -79,18 +93,24 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- ".gitlab-ci.yml"
 - Gemfile
-- Gemfile.lock
 - LICENSE
 - README.md
+- Rakefile
 - 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/input_pg_stat_statements.rb
 - lib/fluent/plugin/parser_multiline_csv.rb
 - test/helper.rb
+- test/plugin/itest_input_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_input_pg_stat_statements.rb
+homepage: https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog
 licenses: []
 metadata: {}
 post_install_message: 
@@ -112,4 +132,10 @@ rubygems_version: 3.1.4
 signing_key: 
 specification_version: 4
 summary: fluentd plugins to work with PostgreSQL CSV logs
-test_files: []
+test_files:
+- test/helper.rb
+- test/plugin/itest_input_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_input_pg_stat_statements.rb

data/Gemfile.lock

@@ -1,50 +0,0 @@
-PATH
-  remote: .
-  specs:
-    fluent-plugin-postgresql-csvlog (0.0.1)
-      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