Add Ruby 2pc tests (#448)

* Add Ruby 2pc tests

* 2pc setup

* fix test warnings

* tag 0.1.4

* prepared transactions note

Author: levkk

CONTRIBUTING.md

@@ -3,16 +3,18 @@
 Contributions are welcome. If you see a bug, feel free to submit a PR with a fix or an issue to discuss. For any features, please open an issue to discuss first.
 
 ## Necessary crates - cargo install <name>
+
 - cargo-nextest
 - cargo-watch
 
 ## Dev setup
 
 1. Run cargo build in the project directory.
-2. Install Postgres (v17 currently supported).
-3. Run the setup script `bash integration/setup.sh`.
-4. Launch pgdog configured for integration: `bash integration/dev-server.sh`.
-5. Run the tests `cargo nextest run --test-threads=1`. If a test fails, try running it directly.
+2. Install Postgres (all Pg versions supported).
+3. Some tests used prepared transactions. Enable them with `ALTER SYSTEM SET max_prepared_transactions TO 1000` and restart Postgres.
+4. Run the setup script `bash integration/setup.sh`.
+5. Launch pgdog configured for integration: `bash integration/dev-server.sh`.
+6. Run the tests `cargo nextest run --test-threads=1`. If a test fails, try running it directly.
 
 ## Coding
 

Cargo.lock

@@ -29,11 +29,11 @@ func LoadStatsForQuery(conn *pgx.Conn, query string) PgStatStatement {
 
 func LoadStatsForPrimary(query string) PgStatStatement {
 	conn, err := pgx.Connect(context.Background(), "postgres://postgres:postgres@127.0.0.1:45000/postgres")
-	defer conn.Close(context.Background())
 
 	if err != nil {
 		panic(err)
 	}
+	defer conn.Close(context.Background())
 
 	return LoadStatsForQuery(conn, query)
 }
@@ -44,10 +44,10 @@ func LoadStatsForReplicas(query string) []PgStatStatement {
 	for i := range 2 {
 		port := 45001 + i
 		conn, err := pgx.Connect(context.Background(), fmt.Sprintf("postgres://postgres:postgres@127.0.0.1:%d/postgres", port))
-		defer conn.Close(context.Background())
 		if err != nil {
 			panic(err)
 		}
+		defer conn.Close(context.Background())
 
 		stats = append(stats, LoadStatsForQuery(conn, query))
 	}
@@ -59,12 +59,13 @@ func ResetStats() {
 	for i := range 3 {
 		port := 45000 + i
 		conn, err := pgx.Connect(context.Background(), fmt.Sprintf("postgres://postgres:postgres@127.0.0.1:%d/postgres?sslmode=disable", port))
-		defer conn.Close(context.Background())
 
 		if err != nil {
 			panic(err)
 		}
 
+		defer conn.Close(context.Background())
+
 		_, err = conn.Exec(context.Background(), "SELECT pg_stat_statements_reset()")
 		if err != nil {
 			panic(err)

integration/load_balancer/pgx/pg_stat_statements.go

@@ -2,8 +2,8 @@
 
 require_relative 'rspec_helper'
 
-def connect(dbname = 'pgdog')
-  PG.connect(dbname: dbname, user: 'pgdog', password: 'pgdog', port: 6432, host: '127.0.0.1')
+def connect(dbname = 'pgdog', user = 'pgdog')
+  PG.connect(dbname: dbname, user: user, password: 'pgdog', port: 6432, host: '127.0.0.1')
 end
 
 describe 'pg' do
@@ -13,82 +13,90 @@ def connect(dbname = 'pgdog')
 
   it 'simple query' do
     %w[pgdog pgdog_sharded].each do |db|
-      conn = connect db
-      res = conn.exec 'SELECT 1::bigint AS one'
-      expect(res[0]['one']).to eq('1')
-      res = conn.exec 'SELECT $1 AS one, $2 AS two', [1, 2]
-      expect(res[0]['one']).to eq('1')
-      expect(res[0]['two']).to eq('2')
+      %w[pgdog pgdog_2pc].each do |user|
+        conn = connect db, user
+        res = conn.exec 'SELECT 1::bigint AS one'
+        expect(res[0]['one']).to eq('1')
+        res = conn.exec 'SELECT $1 AS one, $2 AS two', [1, 2]
+        expect(res[0]['one']).to eq('1')
+        expect(res[0]['two']).to eq('2')
+      end
     end
   end
 
   it 'prepared statements' do
     %w[pgdog pgdog_sharded].each do |db|
-      conn = connect db
-      15.times do |i|
-        name = "_pg_#{i}"
-        conn.prepare name, 'SELECT $1 AS one'
-        res = conn.exec_prepared name, [i]
-        expect(res[0]['one']).to eq(i.to_s)
-      end
-      30.times do |_i|
+      %w[pgdog pgdog_2pc].each do |user|
+        conn = connect db, user
         15.times do |i|
           name = "_pg_#{i}"
+          conn.prepare name, 'SELECT $1 AS one'
           res = conn.exec_prepared name, [i]
           expect(res[0]['one']).to eq(i.to_s)
         end
+        30.times do |_i|
+          15.times do |i|
+            name = "_pg_#{i}"
+            res = conn.exec_prepared name, [i]
+            expect(res[0]['one']).to eq(i.to_s)
+          end
+        end
       end
     end
   end
 
   it 'sharded' do
-    conn = connect 'pgdog_sharded'
-    conn.exec 'DROP TABLE IF EXISTS sharded'
-    conn.exec 'CREATE TABLE sharded (id BIGINT, value TEXT)'
-    conn.prepare 'insert', 'INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *'
-    conn.prepare 'select', 'SELECT * FROM sharded WHERE id = $1'
-    15.times do |i|
-      [10, 10_000_000_000].each do |num|
-        id = num + i
-        results = []
-        results << conn.exec('INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *', [id, 'value_one'])
-        results << conn.exec('SELECT * FROM sharded WHERE id = $1', [id])
-        results.each do |result|
-          expect(result.num_tuples).to eq(1)
-          expect(result[0]['id']).to eq(id.to_s)
-          expect(result[0]['value']).to eq('value_one')
-        end
-        conn.exec 'TRUNCATE TABLE sharded'
-        results << conn.exec_prepared('insert', [id, 'value_one'])
-        results << conn.exec_prepared('select', [id])
-        results.each do |result|
-          expect(result.num_tuples).to eq(1)
-          expect(result[0]['id']).to eq(id.to_s)
-          expect(result[0]['value']).to eq('value_one')
+    %w[pgdog pgdog_2pc].each do |user|
+      conn = connect 'pgdog_sharded', user
+      conn.exec 'DROP TABLE IF EXISTS sharded'
+      conn.exec 'CREATE TABLE sharded (id BIGINT, value TEXT)'
+      conn.prepare 'insert', 'INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *'
+      conn.prepare 'select', 'SELECT * FROM sharded WHERE id = $1'
+      15.times do |i|
+        [10, 10_000_000_000].each do |num|
+          id = num + i
+          results = []
+          results << conn.exec('INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *', [id, 'value_one'])
+          results << conn.exec('SELECT * FROM sharded WHERE id = $1', [id])
+          results.each do |result|
+            expect(result.num_tuples).to eq(1)
+            expect(result[0]['id']).to eq(id.to_s)
+            expect(result[0]['value']).to eq('value_one')
+          end
+          conn.exec 'TRUNCATE TABLE sharded'
+          results << conn.exec_prepared('insert', [id, 'value_one'])
+          results << conn.exec_prepared('select', [id])
+          results.each do |result|
+            expect(result.num_tuples).to eq(1)
+            expect(result[0]['id']).to eq(id.to_s)
+            expect(result[0]['value']).to eq('value_one')
+          end
         end
       end
     end
   end
 
   it 'transactions' do
     %w[pgdog pgdog_sharded].each do |db|
-      conn = connect db
-      conn.exec 'DROP TABLE IF EXISTS sharded'
-      conn.exec 'CREATE TABLE sharded (id BIGINT, value TEXT)'
-      conn.prepare 'insert', 'INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *'
-      conn.prepare 'select', 'SELECT * FROM sharded WHERE id = $1'
+      %w[pgdog pgdog_2pc].each do |user|
+        conn = connect db, user
+        conn.exec 'DROP TABLE IF EXISTS sharded'
+        conn.exec 'CREATE TABLE sharded (id BIGINT, value TEXT)'
+        conn.prepare 'insert', 'INSERT INTO sharded (id, value) VALUES ($1, $2) RETURNING *'
+        conn.prepare 'select', 'SELECT * FROM sharded WHERE id = $1'
 
-      conn.exec 'BEGIN'
-      res = conn.exec_prepared 'insert', [1, 'test']
-      conn.exec 'COMMIT'
-      expect(res.num_tuples).to eq(1)
-      expect(res[0]['id']).to eq(1.to_s)
-      conn.exec 'BEGIN'
-      res = conn.exec_prepared 'select', [1]
-      expect(res.num_tuples).to eq(1)
-      expect(res[0]['id']).to eq(1.to_s)
-      conn.exec 'ROLLBACK'
-      conn.exec 'SELECT 1'
+        conn.exec 'BEGIN'
+        res = conn.exec_prepared 'insert', [1, 'test']
+        conn.exec 'COMMIT'
+        expect(res.num_tuples).to eq(1)
+        expect(res[0]['id']).to eq(1.to_s)
+        conn.exec 'BEGIN'
+        res = conn.exec_prepared 'select', [1]
+        expect(res.num_tuples).to eq(1)
+        expect(res[0]['id']).to eq(1.to_s)
+        conn.exec 'ROLLBACK'
+        conn.exec 'SELECT 1'
+      end
     end
   end
 

integration/ruby/pg_spec.rb

@@ -0,0 +1,17 @@
+[general]
+two_phase_commit = false
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+shard = 0
+database_name = "shard_0"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+shard = 1
+database_name = "shard_1"
+
+[admin]
+password = "pgdog"

integration/two_pc/pgdog.toml

@@ -0,0 +1,8 @@
+#!/bin/bash
+export PGPASSWORD=pgdog
+export PGHOST=127.0.0.1
+export PGPORT=6432
+export PGUSER=pgdog
+export PGDATABASE=pgdog
+
+pgbench -f script.sql -c 10 -t 1000 -P 1

integration/two_pc/run.sh

@@ -0,0 +1,15 @@
+\set id (1021 * random(1, 10000000))
+
+-- In a transaction.
+BEGIN;
+INSERT INTO sharded_2pc (id, value) VALUES (:id, 'some value') RETURNING *;
+SELECT * FROM sharded_2pc WHERE id = :id AND value = 'some value';
+UPDATE sharded_2pc SET value = 'another value' WHERE id = :id;
+DELETE FROM sharded_2pc WHERE id = :id AND value = 'another value';
+ROLLBACK;
+
+-- Outside a transaction.
+INSERT INTO sharded_2pc (id, value) VALUES (:id, 'some value') RETURNING *;
+SELECT * FROM sharded_2pc WHERE id = :id AND value = 'some value';
+UPDATE sharded_2pc SET value = 'another value' WHERE id = :id;
+DELETE FROM sharded_2pc WHERE id = :id;

integration/two_pc/script.sql

@@ -0,0 +1,8 @@
+#!/bin/bash
+export PGPASSWORD=pgdog
+export PGHOST=127.0.0.1
+export PGPORT=6432
+export PGUSER=pgdog
+export PGDATABASE=pgdog
+
+psql -c 'CREATE TABLE IF NOT EXISTS sharded_2pc (id BIGINT PRIMARY KEY, value VARCHAR)'

integration/two_pc/setup.sh

@@ -0,0 +1,4 @@
+[[users]]
+name = "pgdog"
+database = "pgdog"
+password = "pgdog"

integration/two_pc/users.toml

@@ -30,6 +30,14 @@ server_user = "pgdog"
 two_phase_commit = true
 min_pool_size = 0
 
+[[users]]
+name = "pgdog_2pc"
+database = "pgdog"
+password = "pgdog"
+server_user = "pgdog"
+two_phase_commit = true
+min_pool_size = 0
+
 [[users]]
 name = "pgdog_migrator"
 database = "pgdog_sharded"