From 9068c637c95825b4bb876ebedb796037c149373d Mon Sep 17 00:00:00 2001
From: Shayon Mukherjee <shayonj@users.noreply.github.com>
Date: Wed, 13 Dec 2023 09:17:32 -0500
Subject: [PATCH] Drop existing user with privileges when bootstrapping (#75)

Its possible for there to be hangover from previous failed runs where the internal user
may exist in the DB. When that happens, the bootstrap shouldn't fail but attempt to revoke
priveleges and drop the user, to start fresh again.
---
 lib/pg_easy_replicate.rb       | 36 ++++++++++++++++++++++++++++------
 spec/pg_easy_replicate_spec.rb | 16 +++++++++++++++
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/lib/pg_easy_replicate.rb b/lib/pg_easy_replicate.rb
index 6fc0033..a3f9c09 100644
--- a/lib/pg_easy_replicate.rb
+++ b/lib/pg_easy_replicate.rb
@@ -140,10 +140,10 @@ def cleanup(options)
       if options[:everything]
         # Drop users at last
         logger.info("Dropping replication user on source database")
-        drop_user(conn_string: source_db_url, group_name: options[:group_name])
+        drop_user(conn_string: source_db_url)
 
         logger.info("Dropping replication user on target database")
-        drop_user(conn_string: target_db_url, group_name: options[:group_name])
+        drop_user(conn_string: target_db_url)
       end
     rescue => e
       abort_with("Unable to cleanup: #{e.message}")
@@ -266,8 +266,9 @@ def create_user(
     )
       password = connection_info(conn_string)[:password].gsub("'") { "''" }
 
+      drop_user(conn_string: conn_string)
+
       sql = <<~SQL
-        drop role if exists #{quote_ident(internal_user_name)};
         create role #{quote_ident(internal_user_name)} with password '#{password}' login createdb createrole;
         grant all privileges on database #{quote_ident(db_name(conn_string))} TO #{quote_ident(internal_user_name)};
       SQL
@@ -305,10 +306,13 @@ def create_user(
       raise "Unable to create user: #{e.message}"
     end
 
-    def drop_user(conn_string:, group_name:)
+    def drop_user(conn_string:, user: internal_user_name)
+      return unless user_exists?(conn_string: conn_string, user: user)
+
       sql = <<~SQL
-       revoke all privileges on database #{db_name(conn_string)} from #{quote_ident(internal_user_name)};
+       revoke all privileges on database #{db_name(conn_string)} from #{quote_ident(user)};
       SQL
+
       Query.run(
         query: sql,
         connection_url: conn_string,
@@ -316,7 +320,7 @@ def drop_user(conn_string:, group_name:)
       )
 
       sql = <<~SQL
-        drop role if exists #{quote_ident(internal_user_name)};
+        drop role if exists #{quote_ident(user)};
       SQL
 
       Query.run(
@@ -327,5 +331,25 @@ def drop_user(conn_string:, group_name:)
     rescue => e
       raise "Unable to drop user: #{e.message}"
     end
+
+    def user_exists?(conn_string:, user: internal_user_name)
+      sql = <<~SQL
+        SELECT r.rolname AS username,
+          r1.rolname AS "role"
+        FROM pg_catalog.pg_roles r
+        LEFT JOIN pg_catalog.pg_auth_members m ON (m.member = r.oid)
+        LEFT JOIN pg_roles r1 ON (m.roleid=r1.oid)
+        WHERE r.rolname = '#{user}'
+        ORDER BY 1;
+      SQL
+
+      Query
+        .run(
+          query: sql,
+          connection_url: conn_string,
+          user: db_user(conn_string),
+        )
+        .any? { |q| q[:username] == user }
+    end
   end
 end
diff --git a/spec/pg_easy_replicate_spec.rb b/spec/pg_easy_replicate_spec.rb
index 5b48ecb..1cea2a4 100644
--- a/spec/pg_easy_replicate_spec.rb
+++ b/spec/pg_easy_replicate_spec.rb
@@ -167,6 +167,22 @@
       end
     end
 
+    describe ".drop_user" do
+      it "drops the user" do
+        described_class.create_user(conn_string: connection_url)
+
+        expect(described_class.user_exists?(conn_string: connection_url)).to be(
+          true,
+        )
+
+        described_class.drop_user(conn_string: connection_url)
+
+        expect(described_class.user_exists?(conn_string: connection_url)).to be(
+          false,
+        )
+      end
+    end
+
     describe ".bootstrap" do
       before { setup_tables("james-bond", setup_target_db: false) }