fluent-plugin-postgresql-csvlog 0.3.2 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5d61fc31718e43c6d1dff46139a1b03d56384405d45b69c9772dc6f7b6a66dbf
-  data.tar.gz: 5d2d23a4a7b5f277b19f181aa515ab74718f6ef93824e2c958264d9bbdf9c9aa
+  metadata.gz: e1cc2264997378ec801c6623768d57ca27ce1c28d6ac8ffcb6a905add2ce9f16
+  data.tar.gz: 79aab5d187209db5790f9bca181ea0ca8c909c808809590ad46cef0f6645902c
 SHA512:
-  metadata.gz: f2106f60749b6fa8fc931ccd3d85f51595b0ec60eb9f425935072b1749ed068e8a068c6d513f1557a20e8c5a0613acad80c6690c7ce2aa20cca3027d686c388a
-  data.tar.gz: 9ceef623cbd5256e047dea817d487d2711cc8b387ff068a10913100fe649e77174a600ac019b4729403390152c1a3e6f162accf4cc50f020561032b6493d965f
+  metadata.gz: e38e47908ac33706c4b6679fabc97a395bb63f8a110cb6a836c93f984dbfaf1f77f2204209cfe8a0609ddf969ab960cd855ce363398388084a0026e5df0c4970
+  data.tar.gz: 94505a6d7ceee368c22b0c32081b85aaf215d28217ceed1ff8f58bab840367f5afe19ed725056394e57fd1d6ebfe3ae98fba66b39eb781cfa2f1f80a602385c6

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/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

@@ -1,14 +1,9 @@
 # Docker Compose setup useful for testing and development purposes
-version: "3.9"
+version: "3.3"
 services:
-  fluentd:
-    build: .
-    links:
-      - postgres
-    entrypoint: /usr/bin/fluentd -vvv -c /src/example-fluentd.conf
   postgres:
     image: postgres
-    restart: always
+    restart: "no"
     environment:
       - POSTGRES_USER=testuser
       - POSTGRES_PASSWORD=testpass
@@ -17,3 +12,26 @@ services:
     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

@@ -8,5 +8,18 @@
 </source>
 
 <match postgres.pg_stat_statements>
-  @type stdout
+  @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.2'
+  s.version     = '0.4.0'
   s.authors     = ['stanhu']
   s.email       = ['stanhu@gmail.com']
   s.homepage    = 'https://gitlab.com/gitlab-org/fluent-plugins/fluent-plugin-postgresql-csvlog'
@@ -16,7 +16,7 @@ Gem::Specification.new do |s|
 
   s.add_dependency 'fluentd', ['>= 1.0', '< 2']
   s.add_dependency 'pg', '~> 1.1'
-  s.add_dependency 'pg_query', '~> 2.0'
+  s.add_dependency 'pg_query', '~> 2.1'
 
   s.add_development_dependency 'rake'
   s.add_development_dependency 'test-unit', '~> 3.2'

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

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
-require 'fluent/plugin/input'
-require 'pg'
+require_relative './polling_pg_input_plugin'
 require 'pg_query'
 
 module Fluent::Plugin
@@ -14,67 +13,22 @@ module Fluent::Plugin
   #  'query_length' => 58,
   #  'queryid' => 3239318621761098074
   # }
-  class PgStatStatementsInput < Input
+  class PgStatStatementsInput < PollingPostgresInputPlugin
     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
+    protected
 
-    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
+    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']
@@ -97,8 +51,6 @@ module Fluent::Plugin
       record
     end
 
-    private
-
     # Query the database and emit statements to fluentd router
     def emit_statements_to_stream(conn)
       me = Fluent::MultiEventStream.new
@@ -111,26 +63,5 @@ module Fluent::Plugin
 
       @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/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/polling_pg_input_plugin.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require 'fluent/plugin/input'
+require 'pg'
+require 'pg_query'
+
+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
+
+    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
+
+    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
+
+      # Interrupt thread and wait for it to finish
+      Thread.new { @thread.run } if @thread
+      @thread.join
+    end
+
+    # Main polling loop on thread
+    def thread_main
+      until @stop_flag
+        sleep @interval
+        break if @stop_flag
+
+        begin
+          on_poll
+        rescue StandardError => e
+          log.error 'unexpected error', error: e.message, error_class: e.class
+          log.error_backtrace e.backtrace
+        end
+      end
+    end
+
+    protected
+
+    # 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
+    # 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
+      )
+
+      map = PG::BasicTypeMapForResults.new(conn)
+      map.default_type_map = PG::TypeMapAllStrings.new
+
+      conn.type_map_for_results = map
+
+      begin
+        block.call(conn)
+      ensure
+        # Always close the connection
+        conn.finish
+      end
+    end
+  end
+end

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/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/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_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.2
+  version: 0.4.0
 platform: ruby
 authors:
 - stanhu
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-07-12 00:00:00.000000000 Z
+date: 2021-07-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: fluentd
@@ -50,14 +50,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.0'
+        version: '2.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.0'
+        version: '2.1'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -93,6 +93,7 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- ".gitignore"
 - ".gitlab-ci.yml"
 - Dockerfile
 - Gemfile
@@ -105,15 +106,22 @@ files:
 - 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_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_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_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: {}
@@ -138,8 +146,12 @@ specification_version: 4
 summary: fluentd plugins to work with PostgreSQL CSV logs
 test_files:
 - test/helper.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_in_pg_stat_activity.rb
 - test/plugin/test_in_pg_stat_statements.rb
+- test/plugin/test_marginalia_extractor.rb
+- test/verify-docker-compose.sh