Poorly behaved client fix (#65)

* Poorly behaved client fix

* yes officer

* fix tests

* no useless rescue

* Looks ok

Author: levkk

src/client.rs

@@ -359,7 +359,8 @@ impl Client {
                             // Clean up the server and re-use it.
                             // This prevents connection thrashing by bad clients.
                             if server.in_transaction() {
-                                server.query("ROLLBACK; DISCARD ALL;").await?;
+                                server.query("ROLLBACK").await?;
+                                server.query("DISCARD ALL").await?;
                             }
 
                             return Err(err);
@@ -429,7 +430,8 @@ impl Client {
                         // Pgbouncer closes the connection which leads to
                         // connection thrashing when clients misbehave.
                         if server.in_transaction() {
-                            server.query("ROLLBACK; DISCARD ALL;").await?;
+                            server.query("ROLLBACK").await?;
+                            server.query("DISCARD ALL").await?;
                         }
 
                         return Ok(());

tests/ruby/Gemfile

@@ -2,3 +2,4 @@ source "https://rubygems.org"
 
 gem "pg"
 gem "activerecord"
+gem "rubocop"

tests/ruby/Gemfile.lock

@@ -11,20 +11,41 @@ GEM
       i18n (>= 1.6, < 2)
       minitest (>= 5.1)
       tzinfo (~> 2.0)
+    ast (2.4.2)
     concurrent-ruby (1.1.9)
     i18n (1.10.0)
       concurrent-ruby (~> 1.0)
     minitest (5.15.0)
+    parallel (1.22.1)
+    parser (3.1.2.0)
+      ast (~> 2.4.1)
     pg (1.3.2)
+    rainbow (3.1.1)
+    regexp_parser (2.3.1)
+    rexml (3.2.5)
+    rubocop (1.29.0)
+      parallel (~> 1.10)
+      parser (>= 3.1.0.0)
+      rainbow (>= 2.2.2, < 4.0)
+      regexp_parser (>= 1.8, < 3.0)
+      rexml (>= 3.2.5, < 4.0)
+      rubocop-ast (>= 1.17.0, < 2.0)
+      ruby-progressbar (~> 1.7)
+      unicode-display_width (>= 1.4.0, < 3.0)
+    rubocop-ast (1.17.0)
+      parser (>= 3.1.1.0)
+    ruby-progressbar (1.11.0)
     tzinfo (2.0.4)
       concurrent-ruby (~> 1.0)
+    unicode-display_width (2.1.0)
 
 PLATFORMS
   x86_64-linux
 
 DEPENDENCIES
   activerecord
   pg
+  rubocop
 
 BUNDLED WITH
    2.3.7

tests/ruby/tests.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require 'active_record'
+require 'pg'
+
+$stdout.sync = true
 
 # Uncomment these two to see all queries.
 # ActiveRecord.verbose_query_logs = true
@@ -110,3 +113,17 @@ def down
 rescue ActiveRecord::StatementInvalid
   puts 'OK'
 end
+
+# Test evil clients
+def poorly_behaved_client
+  conn = PG::connect("postgres://sharding_user:[email protected]:6432/rails_dev")
+  conn.async_exec 'BEGIN'
+  conn.async_exec 'SELECT 1'
+
+  conn.close
+  puts 'Bad client ok'
+end
+
+25.times do
+  poorly_behaved_client
+end