fluent-plugin-postgresql-csvlog 0.3.0 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bb322f50f848a7b196962dae916a75f30a77f1461167bb314494ae203f8ef9d6
-  data.tar.gz: 077266e6c7f3589813d0c4db8ddccfde76f7d6ecd0b6d6388b67a4e6328ad858
+  metadata.gz: fad3612d27a56001c2eb7f206db3009a8735a68aec57c086f4aec7017a436987
+  data.tar.gz: ff593c807e71fc82465ea8f93844e744c33ef8efa6bee9b92fa1bd56fdb0bf56
 SHA512:
-  metadata.gz: 9b130aa84b1285a62466b3bf68ff6d02acc7e95cd53cc91ee7e34ec81ee4440e06358738790f86f8d436cb8a72ad232f12ac2c0f9d8aca16fd5c6ba19bbdfe49
-  data.tar.gz: e9480110be0cadfad86c5fcc554eb5d1d7264559714e4fd88ea396ae9bb913027b15782c70cd2f4c06e1f4cc5a5bbafe55632ca8eba02ce57e610d98f8440c09
+  metadata.gz: f608d3f94c498f29287c82e82ffc4cbc76d2b64777bd0e9f0931ca1391ce14d2cffaf6d3c523e2e1646aaef83aba5657c0a4d56c6dda9fa255a5c616f5c84fe8
+  data.tar.gz: 17405afcba62970d0a39ab0d654e5f883be745262a12e2f17eaad443c49934d35dbb464a0c69e0ffd63ff71ce88f7b799408efa7ec967ee11cc6d70168e08c35

data/.gitignore

@@ -0,0 +1 @@
+Gemfile.lock

data/.gitlab-ci.yml

@@ -27,3 +27,16 @@ itest:
   cache:
     paths:
       - vendor/ruby
+
+end_to_end_verification_test:
+  image: docker:19.03.12
+  services:
+    - docker:19.03.12-dind
+  tags:
+    - gitlab-org-docker
+  variables:
+    DOCKER_TLS_CERTDIR: ""
+  before_script:
+    - apk add --no-cache docker-compose
+  script:
+    - docker-compose run --rm verifier

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

@@ -7,6 +7,8 @@ parse PostgreSQL CSV log files and extract slow log information:
 - `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
+- `PgStatStatementsInput`: polls the [`pg_stat_statements`](https://www.postgresql.org/docs/current/pgstatstatements.html) postgres plugin and emits fluentd events.
+- `PgStatActivityInput`: polls the [`postges activity monitor`](https://www.postgresql.org/docs/current/monitoring-stats.html) and emits fluentd events.
 
 ## Installation
 
@@ -72,3 +74,17 @@ ingest and parse PostgreSQL CSV logs:
   </format>
 </match>
 ```
+
+## Developing `fluent-plugin-postgresql-csvlog`
+
+To develop and debug locally, there is a `Dockerfile` and `docker-compose.yml` that will setup a local environment,
+complete with Postgres, suitable for testing purposes.
+
+1. `docker compose rm verifier --rm` - test the current configuration
+1. `docker compose up`
+
+### Releasing a new version
+
+1. Update the version in `fluent-plugin-postgresql-csvlog.gemspec`.
+1. Create a merge request and merge the changes to `master`.
+1. Run `bundle exec rake release`.

data/docker-compose.yml

@@ -0,0 +1,37 @@
+# Docker Compose setup useful for testing and development purposes
+version: "3.3"
+services:
+  postgres:
+    image: postgres
+    restart: "no"
+    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
+
+  fluentd:
+    build: .
+    restart: "no"
+    links:
+      - postgres
+    entrypoint: /usr/bin/fluentd -vvv -c /src/example-fluentd.conf
+    volumes:
+      - ./example-fluentd.conf:/src/example-fluentd.conf
+      - log-volume:/var/log/pg/
+
+  verifier:
+    image: alpine:3.13
+    restart: "no"
+    links:
+      - fluentd
+    command: /bin/sh /src/verify-docker-compose.sh
+    volumes:
+      - ./test/verify-docker-compose.sh:/src/verify-docker-compose.sh
+      - log-volume:/var/log/pg/
+
+volumes:
+  log-volume:

data/example-fluentd.conf

@@ -0,0 +1,25 @@
+<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 file
+  path /var/log/pg/pg_stat_statements
+  time_slice_format %Y%m%d%H%M%S
+  flush_interval 1s
+  utc
+</match>
+
+<source>
+  @type pg_stat_activity
+  tag postgres.pg_stat_activity
+  host postgres
+  username testuser
+  password testpass
+  interval 1
+</source>

data/fluent-plugin-postgresql-csvlog.gemspec

@@ -2,7 +2,7 @@ $:.push File.expand_path('lib', __dir__)
 
 Gem::Specification.new do |s|
   s.name        = 'fluent-plugin-postgresql-csvlog'
-  s.version     = '0.3.0'
+  s.version     = '0.5.0'
   s.authors     = ['stanhu']
   s.email       = ['stanhu@gmail.com']
   s.homepage    = 'https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog'

data/lib/fluent/plugin/filter_marginalia.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'fluent/plugin/filter'
+require_relative './marginalia_extractor'
 
 module Fluent
   module Plugin
@@ -11,6 +12,7 @@ module Fluent
     # /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/ SELECT COUNT(*) FROM "projects"
     #
     class Marginalia < Filter
+      include MarginaliaExtractor
       Fluent::Plugin.register_filter('marginalia', self)
 
       desc 'Field to parse for Marginalia comments (key1:value1,key2:value2)'
@@ -19,77 +21,11 @@ module Fluent
       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)
+        parse_marginalia_into_record(record, @key, @strip_comment)
 
         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/in_pg_stat_activity.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+require_relative './polling_pg_input_plugin'
+require 'pg_query'
+require_relative './marginalia_extractor'
+require 'time'
+
+module Fluent::Plugin
+  # PgStatActivityInput polls the `pg_stat_activity` table
+  # emitting normalized versions of the queries currently running on
+  # the postgres server.
+  # Fingerprints of the queries are also included for easier aggregation
+  class PgStatActivityInput < PollingPostgresInputPlugin
+    include MarginaliaExtractor
+    Fluent::Plugin.register_input('pg_stat_activity', self)
+
+    ACTIVITY_QUERY = <<-SQL
+      SELECT
+        datid,
+        datname,
+        pid,
+        usesysid,
+        usename,
+        application_name,
+        host(client_addr) as client_addr,
+        client_hostname,
+        client_port,
+        xact_start,
+        extract(epoch from clock_timestamp() - xact_start) xact_age_s,
+        query_start,
+        extract(epoch from clock_timestamp() - query_start) query_age_s,
+        state_change,
+        extract(epoch from clock_timestamp() - state_change) state_age_s,
+        state,
+        query
+      FROM pg_stat_activity
+      WHERE usename IS NOT NULL
+    SQL
+
+    desc 'Name of field to store SQL query fingerprint'
+    config_param :fingerprint_key, :string, default: 'fingerprint'
+
+    protected
+
+    def on_poll
+      with_connection do |conn|
+        emit_activity_to_stream(conn)
+      end
+    end
+
+    public
+
+    # Query the database and emit statements to fluentd router
+    def emit_activity_to_stream(conn)
+      me = Fluent::MultiEventStream.new
+
+      now = Fluent::Engine.now
+      conn.exec(ACTIVITY_QUERY).each do |row|
+        record = record_for_row(row)
+        me.add(now, record)
+      end
+
+      @router.emit_stream(@tag, me)
+    end
+
+    # Returns a fluentd record for a query row
+    def record_for_row(row)
+      record = {
+        'datid' => row['datid'],
+        'datname' => row['datname'],
+        'pid' => row['pid'],
+        'usesysid' => row['usesysid'],
+        'usename' => row['usename'],
+        'application_name' => row['application_name'],
+        'client_addr' => row['client_addr'],
+        'client_hostname' => row['client_hostname'],
+        'client_port' => row['client_port'],
+        'xact_start' => row['xact_start']&.iso8601(3),
+        'xact_age_s' => row['xact_age_s'],
+        'query_start' => row['query_start']&.iso8601(3),
+        'query_age_s' => row['query_age_s'],
+        'state_change' => row['state_change']&.iso8601(3),
+        'state_age_s' => row['state_age_s'],
+        'state' => row['state'],
+        'query' => row['query'] # This will be stripped, normalized etc
+      }
+
+      # Inject marginalia into record
+      parse_marginalia_into_record(record, 'query', true)
+
+      # Normalize query and fingerprint
+      # Note that `record['query']` was updated in previous step
+      # To strip off marginalia comments
+      record.merge!(fingerprint_query(record['query']))
+
+      record
+    end
+
+    def fingerprint_query(query)
+      # We record the query_length as it will help in understanding whether unparseable
+      # queries are truncated.
+      record = { 'query_length' => query&.length, 'query' => nil }
+
+      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
+  end
+end

data/lib/fluent/plugin/in_pg_stat_statements.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require_relative './polling_pg_input_plugin'
+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 < PollingPostgresInputPlugin
+    Fluent::Plugin.register_input('pg_stat_statements', self)
+
+    desc 'Name of field to store SQL query fingerprint'
+    config_param :fingerprint_key, :string, default: 'fingerprint'
+
+    protected
+
+    def on_poll
+      with_connection do |conn|
+        emit_statements_to_stream(conn)
+      end
+    end
+
+    public
+
+    # 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
+
+    # 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
+  end
+end

data/lib/fluent/plugin/marginalia_extractor.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require 'fluent/plugin/filter'
+
+module Fluent::Plugin
+  # MarginaliaExtractor provides the parse_marginalia_into_record
+  # utility method, useful for extracting marginalia into fluentd records
+  module MarginaliaExtractor
+    MARGINALIA_PREPENDED_REGEXP = %r{^(?<comment>/\*.*\*/)(?<sql>.*)}m.freeze
+    MARGINALIA_APPENDED_REGEXP = %r{(?<sql>.*)(?<comment>/\*.*\*/)$}m.freeze
+
+    # Injects marginalia into a fluentd record
+    def parse_marginalia_into_record(record, key, strip_comment)
+      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, key, 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, key, record)
+      entries.each do |component|
+        data = component.split(':', 2)
+
+        break unless data.length == 2
+
+        stored_key = store_key(record, key, data[0])
+        record[stored_key] = data[1]
+      end
+    end
+
+    def store_key(record, key, 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

data/lib/fluent/plugin/{input_pg_stat_statements.rb → polling_pg_input_plugin.rb}

@@ -1,22 +1,15 @@
 # frozen_string_literal: true
 
-require 'fluent/input'
+require 'fluent/plugin/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)
-
+module Fluent::Plugin
+  # PollingPostgresInputPlugin is intended to be used as an base class
+  # for input plugins that poll postgres.
+  #
+  # Child classes should implement the `on_poll` method
+  class PollingPostgresInputPlugin < Input
     desc 'PostgreSQL host'
     config_param :host, :string
 
@@ -43,14 +36,12 @@ module Fluent
     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
 
+    # Fluentd shutdown method, called to terminate and cleanup plugin
     def shutdown
       @stop_flag = true
 
@@ -59,15 +50,14 @@ module Fluent
       @thread.join
     end
 
+    # Main polling loop on thread
     def thread_main
       until @stop_flag
         sleep @interval
         break if @stop_flag
 
         begin
-          with_connection do |conn|
-            emit_statements_to_stream(conn)
-          end
+          on_poll
         rescue StandardError => e
           log.error 'unexpected error', error: e.message, error_class: e.class
           log.error_backtrace e.backtrace
@@ -75,41 +65,11 @@ module Fluent
       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
+    protected
 
-      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)
+    # Child-classes should implement this method
+    def on_poll
+      raise 'on_poll must be implemented by descendents of PollingPostgresInputPlugin'
     end
 
     # Since this query is very infrequent, and it may be communicating directly
@@ -123,7 +83,11 @@ module Fluent
         user: @username,
         password: @password
       )
-      conn.type_map_for_results = PG::BasicTypeMapForResults.new conn
+
+      map = PG::BasicTypeMapForResults.new(conn)
+      map.default_type_map = PG::TypeMapAllStrings.new
+
+      conn.type_map_for_results = map
 
       begin
         block.call(conn)

data/sql/create_extension.sql

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

data/test/helper.rb

@@ -16,4 +16,6 @@ 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'
+require 'fluent/plugin/marginalia_extractor'
+require 'fluent/plugin/in_pg_stat_statements'
+require 'fluent/plugin/in_pg_stat_activity'

data/test/plugin/itest_in_pg_stat_activity.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class PgStatActivityInputIntegrationTest < 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
+    )
+  end
+
+  def teardown
+    @conn&.finish
+  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::PgStatActivityInput).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
+
+      first_with_query = emits.find do |event|
+        record = event[2]
+
+        record['usename'] == USERNAME &&
+          !record['datid'].nil? &&
+          !record['query'].nil? &&
+          record['state'] == 'active'
+      end
+
+      assert_false first_with_query.nil?
+      record = first_with_query[2]
+
+      assert_false record['datname'].nil?
+      assert_false record['pid'].nil?
+      assert_false record['usesysid'].nil?
+      assert_false record['application_name'].nil?
+      assert_false record['client_addr'].nil?
+      assert_false record['client_port'].nil?
+      assert_false record['xact_start'].nil?
+      assert_false record['xact_age_s'].nil?
+      assert_false record['xact_start'].nil?
+      assert_false record['xact_age_s'].nil?
+      assert_false record['query_start'].nil?
+      assert_false record['query_age_s'].nil?
+      assert_false record['state_change'].nil?
+      assert_false record['state_age_s'].nil?
+      assert_false record['query_length'].nil?
+      assert_false record['query'].nil?
+      assert_false record['fingerprint'].nil?
+    end
+  end
+end

data/test/plugin/{itest_input_pg_stat_statements.rb → itest_in_pg_stat_statements.rb}

@@ -3,10 +3,11 @@
 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'
+  # 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
@@ -53,7 +54,7 @@ class PgStatStatementsInputIntegrationTest < Test::Unit::TestCase
   )
 
   def create_driver(config)
-    Fluent::Test::InputTestDriver.new(Fluent::PgStatStatementsInput).configure(config)
+    Fluent::Test::InputTestDriver.new(Fluent::Plugin::PgStatStatementsInput).configure(config)
   end
 
   sub_test_case 'configuration' do

data/test/plugin/test_in_pg_stat_activity.rb

@@ -0,0 +1,223 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+require 'date'
+
+class PgStatActivityInputTest < Test::Unit::TestCase
+  def setup
+    Fluent::Test.setup
+  end
+
+  CONFIG = %(
+    tag postgres.pg_stat_activity
+    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::PgStatActivityInput).configure(CONFIG)
+  end
+
+  sub_test_case 'configuration' do
+    test 'basic configuration' do
+      d = create_driver
+
+      assert_equal 'postgres.pg_stat_activity', 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
+      row = {
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'pid' => 376,
+        'usesysid' => 10,
+        'usename' => 'testuser',
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'xact_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'xact_age_s' => 0.001884,
+        'query_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'query_age_s' => 0.001894,
+        'state_change' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'state_age_s' => 0.001894,
+        'state' => 'active',
+        'query' => "SELECT * FROM users WHERE user_secret = 's3cr3t'"
+      }
+
+      record = d.instance.record_for_row(row)
+
+      expected = {
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'fingerprint' => '5c4a61e156c7d822',
+        'pid' => 376,
+        'query' => 'SELECT * FROM users WHERE user_secret = $1', # NOTE: secret removed
+        'query_age_s' => 0.001894,
+        'query_length' => 48,
+        'query_start' => '2021-07-23T12:55:25.000+00:00',
+        'state' => 'active',
+        'state_age_s' => 0.001894,
+        'state_change' => '2021-07-23T12:55:25.000+00:00',
+        'usename' => 'testuser',
+        'usesysid' => 10,
+        'xact_age_s' => 0.001884,
+        'xact_start' => '2021-07-23T12:55:25.000+00:00'
+      }
+
+      assert_equal expected, record
+    end
+
+    # This test mostly checks that the code is nil safe
+    test 'nil query' do
+      d = create_driver
+      record = d.instance.record_for_row({})
+
+      expected = {
+        'application_name' => nil,
+        'client_addr' => nil,
+        'client_hostname' => nil,
+        'client_port' => nil,
+        'datid' => nil,
+        'datname' => nil,
+        'pid' => nil,
+        'query' => nil,
+        'query_age_s' => nil,
+        'query_length' => nil,
+        'query_start' => nil,
+        'state' => nil,
+        'state_age_s' => nil,
+        'state_change' => nil,
+        'usename' => nil,
+        'usesysid' => nil,
+        'xact_age_s' => nil,
+        'xact_start' => nil
+      }
+
+      assert_equal expected, record
+    end
+
+    test 'unparseable sql' do
+      d = create_driver
+      row = {
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'pid' => 376,
+        'usesysid' => 10,
+        'usename' => 'testuser',
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'xact_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'xact_age_s' => 0.001884,
+        'query_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'query_age_s' => 0.001894,
+        'state_change' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'state_age_s' => 0.001894,
+        'state' => 'active',
+        'query' => "SELECT * FROM users WHERE user_se="
+      }
+
+      record = d.instance.record_for_row(row)
+
+      expected = {
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'pid' => 376,
+        'query' => nil,
+        'query_age_s' => 0.001894,
+        'query_length' => 34,
+        'query_start' => '2021-07-23T12:55:25.000+00:00',
+        'query_unparseable' => true,
+        'state' => 'active',
+        'state_age_s' => 0.001894,
+        'state_change' => '2021-07-23T12:55:25.000+00:00',
+        'usename' => 'testuser',
+        'usesysid' => 10,
+        'xact_age_s' => 0.001884,
+        'xact_start' => '2021-07-23T12:55:25.000+00:00'
+      }
+
+      assert_equal expected, record
+    end
+
+    test 'marginalia prepended' do
+      d = create_driver
+      row = {
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'pid' => 376,
+        'usesysid' => 10,
+        'usename' => 'testuser',
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'xact_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'xact_age_s' => 0.001884,
+        'query_start' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'query_age_s' => 0.001894,
+        'state_change' => Time.parse('2021-07-23 12:55:25 +0000'),
+        'state_age_s' => 0.001894,
+        'state' => 'active',
+        'query' => "/*application:web,correlation_id:01F1D2T1SC9DM82A4865ATG1CP,endpoint_id:POST /api/:version/groups/:id/-/packages/mavenpath/:file_name*/ SELECT * FROM users WHERE user_secret = 's3cr3t'"
+      }
+
+      record = d.instance.record_for_row(row)
+
+      expected = {
+        'application' => 'web',
+        'application_name' => 'psql',
+        'client_addr' => '172.17.0.1',
+        'client_hostname' => nil,
+        'client_port' => 60182,
+        'correlation_id' => '01F1D2T1SC9DM82A4865ATG1CP',
+        'datid' => 16384,
+        'datname' => 'testuser',
+        'endpoint_id' => 'POST /api/:version/groups/:id/-/packages/mavenpath/:file_name',
+        'fingerprint' => '5c4a61e156c7d822',
+        'pid' => 376,
+        'query' => 'SELECT * FROM users WHERE user_secret = $1', # Secret removed
+        'query_age_s' => 0.001894,
+        'query_length' => 48,
+        'query_start' => '2021-07-23T12:55:25.000+00:00',
+        'state' => 'active',
+        'state_age_s' => 0.001894,
+        'state_change' => '2021-07-23T12:55:25.000+00:00',
+        'usename' => 'testuser',
+        'usesysid' => 10,
+        'xact_age_s' => 0.001884,
+        'xact_start' => '2021-07-23T12:55:25.000+00:00'
+      }
+
+      assert_equal expected, record
+    end
+  end
+end

data/test/plugin/{test_input_pg_stat_statements.rb → test_in_pg_stat_statements.rb}

@@ -20,7 +20,7 @@ class PgStatStatementsInputTest < Test::Unit::TestCase
   )
 
   def create_driver
-    Fluent::Test::InputTestDriver.new(Fluent::PgStatStatementsInput).configure(CONFIG)
+    Fluent::Test::InputTestDriver.new(Fluent::Plugin::PgStatStatementsInput).configure(CONFIG)
   end
 
   sub_test_case 'configuration' do

data/test/plugin/test_marginalia_extractor.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require_relative '../helper'
+
+class Marginalia < Test::Unit::TestCase
+  include Fluent::Plugin::MarginaliaExtractor
+
+  def test_parse(sql, record, key, strip_comment, expected)
+    record[key] = sql
+    parse_marginalia_into_record(record, key, strip_comment)
+    assert_equal(expected, record)
+  end
+
+  test 'no marginalia' do
+    sql = 'SELECT * FROM projects'
+    expected = { 'sql' => 'SELECT * FROM projects' }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'normal comment appended' do
+    sql = 'SELECT COUNT(*) FROM "projects" /* this is just a comment */'
+    expected = {
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'marginalia appended for sidekiq' do
+    sql = 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/'
+    expected = {
+      'application' => 'sidekiq',
+      'correlation_id' => 'd67cae54c169e0cab7d73389e2934f0e',
+      'jid' => '52a1c8a9e4c555ea573f20f0',
+      'job_class' => 'Geo::MetricsUpdateWorker',
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'marginalia appended for web' do
+    sql = 'SELECT COUNT(*) FROM "projects" /*application:web,correlation_id:01F1D2T1SC9DM82A4865ATG1CP,endpoint_id:POST /api/:version/groups/:id/-/packages/mavenpath/:file_name*/'
+    expected = {
+      'application' => 'web',
+      'correlation_id' => '01F1D2T1SC9DM82A4865ATG1CP',
+      'endpoint_id' => 'POST /api/:version/groups/:id/-/packages/mavenpath/:file_name',
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'normal comment prepended' do
+    sql = '/* this is just a comment */ SELECT COUNT(*) FROM "projects"'
+    expected = {
+      "sql" => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'marginalia prepended for sidekiq' do
+    sql = '/*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/ SELECT COUNT(*) FROM "projects"'
+    expected = {
+      'application' => 'sidekiq',
+      'correlation_id' => 'd67cae54c169e0cab7d73389e2934f0e',
+      'jid' => '52a1c8a9e4c555ea573f20f0',
+      'job_class' => 'Geo::MetricsUpdateWorker',
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'marginalia prepended for web' do
+    sql = '/*application:web,correlation_id:01F1D2T1SC9DM82A4865ATG1CP,endpoint_id:POST /api/:version/groups/:id/-/packages/mavenpath/:file_name*/ SELECT COUNT(*) FROM "projects"'
+    expected = {
+      'application' => 'web',
+      'correlation_id' => '01F1D2T1SC9DM82A4865ATG1CP',
+      'endpoint_id' => 'POST /api/:version/groups/:id/-/packages/mavenpath/:file_name',
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, {}, 'sql', true, expected)
+  end
+
+  test 'marginalia prepended for web, comment_strip disabled' do
+    sql = 'SELECT COUNT(*) FROM "projects" /*application:sidekiq,correlation_id:d67cae54c169e0cab7d73389e2934f0e,jid:52a1c8a9e4c555ea573f20f0,job_class:Geo::MetricsUpdateWorker*/'
+    expected = {
+      'application' => 'sidekiq',
+      'correlation_id' => 'd67cae54c169e0cab7d73389e2934f0e',
+      'jid' => '52a1c8a9e4c555ea573f20f0',
+      'job_class' => 'Geo::MetricsUpdateWorker',
+      'sql' => sql
+    }
+    test_parse(sql, { 'sql' => sql }, 'sql', false, expected)
+  end
+
+  test 'avoid clash' do
+    sql = '/*clash_key:bbb*/ SELECT COUNT(*) FROM "projects"'
+    expected = {
+      'clash_key' => 'aaa',
+      'sql_clash_key' => 'bbb',
+      'sql' => 'SELECT COUNT(*) FROM "projects"'
+    }
+    test_parse(sql, { 'clash_key' => 'aaa' }, 'sql', true, expected)
+  end
+end

data/test/verify-docker-compose.sh

@@ -0,0 +1,22 @@
+#!/bin/sh
+
+# This script is use by the end-to-end fluent test
+# See the docker-compose.yml for more details
+
+cleanup() {
+  echo "# removing all logs"
+  rm -rf /var/log/pg/*
+}
+
+die() {
+  cleanup
+  echo "$1"
+  exit 1
+}
+
+cleanup
+echo "# sleeping 10, awaiting logs"
+sleep 10;
+find /var/log/pg/ -name "pg_stat_statements.*.log" || die "No pg_stat_statements files created"
+
+cleanup

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: fluent-plugin-postgresql-csvlog
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.5.0
 platform: ruby
 authors:
 - stanhu
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-06-24 00:00:00.000000000 Z
+date: 2021-07-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: fluentd
@@ -93,23 +93,35 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- ".gitignore"
 - ".gitlab-ci.yml"
+- Dockerfile
 - Gemfile
 - 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/input_pg_stat_statements.rb
+- lib/fluent/plugin/in_pg_stat_activity.rb
+- lib/fluent/plugin/in_pg_stat_statements.rb
+- lib/fluent/plugin/marginalia_extractor.rb
 - lib/fluent/plugin/parser_multiline_csv.rb
+- lib/fluent/plugin/polling_pg_input_plugin.rb
+- sql/create_extension.sql
 - test/helper.rb
-- test/plugin/itest_input_pg_stat_statements.rb
+- test/plugin/itest_in_pg_stat_activity.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_input_pg_stat_statements.rb
+- test/plugin/test_in_pg_stat_activity.rb
+- test/plugin/test_in_pg_stat_statements.rb
+- test/plugin/test_marginalia_extractor.rb
+- test/verify-docker-compose.sh
 homepage: https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog
 licenses: []
 metadata: {}
@@ -134,8 +146,12 @@ specification_version: 4
 summary: fluentd plugins to work with PostgreSQL CSV logs
 test_files:
 - test/helper.rb
-- test/plugin/itest_input_pg_stat_statements.rb
+- test/plugin/itest_in_pg_stat_activity.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_input_pg_stat_statements.rb
+- test/plugin/test_in_pg_stat_activity.rb
+- test/plugin/test_in_pg_stat_statements.rb
+- test/plugin/test_marginalia_extractor.rb
+- test/verify-docker-compose.sh