Experimental PG patch

.github/workflows/ci.yml

@@ -35,6 +35,14 @@ jobs:
         ports:
           - 9200:9200
         options: -e="discovery.type=single-node" --health-cmd="curl http://localhost:9200/_cluster/health" --health-interval=3s --health-timeout=5s --health-retries=20
+      postgres:
+        image: postgres:12.6
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+        ports:
+          - 5432:5432
+        options: --health-cmd="pg_isready -U postgres" --health-interval=3s --health-timeout=5s --health-retries=20
     steps:
       - uses: actions/checkout@v3
       - uses: ruby/setup-ruby@v1

Gemfile

@@ -16,6 +16,7 @@ gem 'byebug', platforms: not_jruby
 gem 'elasticsearch', '> 7', '< 7.14'
 gem 'honeybadger', '>= 2.0'
 gem 'irb', '~> 1.0'
+gem 'pg', '~> 1.1'
 # Minimum of 0.5.0 for specific error classes
 gem 'mysql2', '>= 0.5.0', platforms: not_jruby
 gem 'redcarpet', '~> 3.5', platforms: not_jruby

README.md

@@ -1047,6 +1047,34 @@ mysql.query('SELECT * FROM users') # raises Faulty::CircuitError if connection f
 mysql = Mysql2::Client.new(host: '127.0.0.1')
 mysql.query('SELECT * FROM users') # not protected by a circuit
 ```
+### Patch::Postgres
+
+[`Faulty::Patch::Postgres`](https://www.rubydoc.info/gems/faulty/Faulty/Patch/Postgres)
+protects a `PG::Connection` with an internal circuit. Pass a `:faulty` key along
+with your connection options to enable the circuit breaker.
+
+Faulty supports the pg gem versions 1.0 and greater.
+
+```ruby
+require 'faulty/patch/postgres'
+
+pg = PG::Connection.new(host: 'localhost', faulty: {
+  # The name for the Postgres circuit
+  name: 'postgres'
+
+  # The faulty instance to use
+  # This can also be a registered faulty instance or a constant name. See API
+  # docs for more details
+  instance: Faulty.default
+
+  # By default, circuit errors will be subclasses of PG::Error
+  # To disable this behavior, set patch_errors to false and Faulty
+  # will raise its default errors
+  patch_errors: true
+})
+```
+
+
 
 ### Patch::Elasticsearch
 

lib/faulty/patch/postgres.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'pg'
+
+class Faulty
+  module Patch
+    # Patch for the Postgres gem
+    module Postgres
+      include Base
+
+      Patch.define_circuit_errors(self, ::PG::ConnectionBad)
+
+      QUERY_WHITELIST = [
+        %r{\A(?:/\*.*?\*/)?\s*ROLLBACK}i,
+        %r{\A(?:/\*.*?\*/)?\s*COMMIT}i,
+        %r{\A(?:/\*.*?\*/)?\s*RELEASE\s+SAVEPOINT}i
+      ].freeze
+
+      def initialize(opts = {})
+        @faulty_circuit = Patch.circuit_from_hash(
+          'pg',
+          opts[:faulty],
+          errors: [
+            ::PG::ConnectionBad,
+            ::PG::UnableToSend
+          ],
+          patched_error_mapper: Faulty::Patch::Postgres
+        )
+
+        super
+      end
+
+      def ping
+        faulty_run { super }
+      rescue Faulty::Patch::Postgres::FaultyError
+        false
+      end
+
+      def connect(*args)
+        faulty_run { super }
+      end
+
+      def query(*args)
+        return super if QUERY_WHITELIST.any? { |r| !r.match(args.first).nil? }
+
+        faulty_run { super }
+      end
+    end
+  end
+end
+
+module PG
+  class Connection
+    prepend Faulty::Patch::Postgres
+  end
+end

spec/patch/postgres_spec.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+RSpec.describe 'Faulty::Patch::Postgres', if: defined?(PG) do
+  def new_client(options = {})
+    PG::Connection.new({
+      username: ENV.fetch('POSTGRES_USER', nil),
+      password: ENV.fetch('POSTGRES_PASSWORD', nil),
+      host: ENV.fetch('POSTGRES_HOST', nil),
+      port: ENV.fetch('POSTGRES_PORT', nil),
+      socket: ENV.fetch('POSTGRES_SOCKET', nil)
+    }.merge(options))
+  end
+
+  def create_table(client, table_name)
+    client.exec("CREATE TABLE #{table_name} (id serial PRIMARY KEY, name text)")
+  end
+
+  def trip_circuit
+    client
+    4.times do
+      begin
+        new_client(host: '127.0.0.1', port: 9999, faulty: { instance: 'faulty' })
+      rescue PG::ConnectionBad
+        # expected
+      end
+    end
+  end
+
+  let(:client) { new_client(database: db_name, faulty: { instance: 'faulty' }) }
+  let(:bad_client) { new_client(host: '127.0.0.1', port: 9999, faulty: { instance: 'faulty' }) }
+  let(:bad_unpatched_client) { new_client(host: '127.0.0.1', port: 9999) }
+  let(:faulty) { Faulty.new(listeners: [], circuit_defaultts: { sample_threshold: 2 }) }
+
+  before do
+    new_client.exec("CREATE DATABASE #{db_name}")
+  end
+
+  after do
+    new_client.exec("DROP DATABASE #{db_name}")
+  end
+
+  it 'captures connection error' do
+    expect { bad_client.query('SELECT 1 FROM dual') }.to raise_error do |error|
+      expect(error).to be_a(Faulty::Patch::PG::ConnectionError)
+      expect(error.cause).to be_a(PG::Error::ConnectionBad)
+    end
+    expect(faulty.circuit('postgres').status.failure_rate).to eq(1)
+  end
+
+  it 'does not capture unpatched client errors' do
+    expect { bad_unpatched_client.query('SELECT 1 FROM dual') }.to raise_error(PG::Error::ConnectionBad)
+    expect(faulty.circuit('postgres').status.failure_rate).to eq(0)
+  end
+
+  it 'does not capture application errors' do
+    expect { client.query('SELECT * FROM not_a_table') }.to raise_error(PG::Error)
+    expect(faulty.circuit('postgres').status.failure_rate).to eq(0)
+  end
+
+  it 'successfully executes query' do
+    create_table(client, 'test')
+    client.query('INSERT INTO test VALUES(1)')
+    expect(client.query('SELECT * FROM test').to_a).to eq([{ 'id' => '1' }])
+    expect(faulty.circuit('postgres').status.failure_rate).to eq(0)
+  end
+
+  it 'prevents additional queries when tripped' do
+    trip_circuit
+    expect { client.query('SELECT 1 FROM dual') }.to raise_error(Faulty::Patch::PG::ConnectionError)
+  end
+
+  it 'allows COMMIT when tripped' do
+    create_table(client, 'test')
+    client.query('BEGIN')
+    client.query('INSERT INTO test VALUES(1)')
+    trip_circuit
+    expect { client.query('COMMIT') }.to be_nil
+    expect(client.query('SELECT * FROM test')).to raise_error(Faulty::Patch::PG::ConnectionError)
+    faulty.circuit('postgres').reset
+    expect(client.query('SELECT * FROM test').to_a).to eq([{ 'id' => '1' }])
+  end
+
+  it 'allows ROLLBACK with a leading comment when tripped' do
+    create_table(client, 'test')
+    client.query('BEGIN')
+    client.query('INSERT INTO test VALUES(1)')
+    trip_circuit
+    expect { client.query('/* hi there */ ROLLBACK') }.to be_nil
+    expect { client.query('SELECT * FROM test') }.to raise_error(Faulty::Patch::PG::ConnectionError)
+    faulty.circuit('postgres').reset
+    expect(client.query('SELECT * FROM test').to_a).to eq([])
+  end
+end

spec/spec_helper.rb

@@ -19,6 +19,7 @@
 require 'faulty'
 require 'faulty/patch/redis'
 require 'faulty/patch/elasticsearch'
+require 'faulty/patch/postgres'
 require 'timecop'
 require 'redis'
 require 'json'